src/FuncTorrent/PeerMsgs.hs

   1 {-
   2  - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
   3  -
   4  - This file is part of FuncTorrent.
   5  -
   6  - FuncTorrent is free software; you can redistribute it and/or modify
   7  - it under the terms of the GNU General Public License as published by
   8  - the Free Software Foundation; either version 3 of the License, or
   9  - (at your option) any later version.
  10  -
  11  - FuncTorrent is distributed in the hope that it will be useful,
  12  - but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  - GNU General Public License for more details.
  15  -
  16  - You should have received a copy of the GNU General Public License
  17  - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
  18  -}
  19
  20 {-# LANGUAGE OverloadedStrings #-}
  21
  22 module FuncTorrent.PeerMsgs
  23        (genHandshakeMsg,
  24         sendMsg,
  25         getMsg,
  26         Peer(..),
  27         makePeer,
  28         PeerMsg(..)
  29        ) where
  30
  31 import Prelude hiding (lookup, concat, replicate, splitAt, take)
  32
  33 import System.IO (Handle)
  34 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
  35 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
  36 import Data.ByteString.Char8 as BC (splitAt)
  37 import qualified Data.ByteString.Char8 as BC (replicate, pack)
  38 import Control.Monad (replicateM)
  39 import Control.Applicative (liftA3)
  40
  41 import Data.Binary (Binary(..), decode, encode)
  42 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  43 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  44
  45 import FuncTorrent.Utils (toIP, toPort)
  46
  47 -- | Peer is a IP address, port tuple
  48 data Peer = Peer IP Port
  49           deriving (Show, Eq)
  50
  51 type ID = String
  52 type IP = String
  53 type Port = Integer
  54
  55 data PeerMsg = KeepAliveMsg
  56              | ChokeMsg
  57              | UnChokeMsg
  58              | InterestedMsg
  59              | NotInterestedMsg
  60              | HaveMsg Integer
  61              | BitFieldMsg ByteString
  62              | RequestMsg Integer Integer Integer
  63              | PieceMsg Integer Integer ByteString
  64              | CancelMsg Integer Integer Integer
  65              | PortMsg Port
  66              | ExtendedMsg Integer ByteString
  67   deriving (Show)
  68
  69 data ExtMetadataMsg = Request Integer
  70                     | Data Integer Integer
  71                     | Reject Integer
  72   deriving (Eq, Show)
  73
  74 instance Binary PeerMsg where
  75   put msg = case msg of
  76              KeepAliveMsg -> putWord32be 0
  77              ChokeMsg -> do putWord32be 1
  78                             putWord8 0
  79              UnChokeMsg -> do putWord32be 1
  80                               putWord8 1
  81              InterestedMsg -> do putWord32be 1
  82                                  putWord8 2
  83              NotInterestedMsg -> do putWord32be 1
  84                                     putWord8 3
  85              HaveMsg i -> do putWord32be 5
  86                              putWord8 4
  87                              putWord32be (fromIntegral i)
  88              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
  89                                   putWord8 5
  90                                   mapM_ putWord8 bfList
  91                                     where bfList = unpack bf
  92                                           bfListLen = length bfList
  93              RequestMsg i o l -> do putWord32be 13
  94                                     putWord8 6
  95                                     putIndexOffsetLength i o l
  96              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
  97                                   putWord8 7
  98                                   putWord32be (fromIntegral i)
  99                                   putWord32be (fromIntegral o)
 100                                   mapM_ putWord8 blockList
 101                                     where blockList = unpack b
 102                                           blocklen = length blockList
 103              CancelMsg i o l -> do putWord32be 13
 104                                    putWord8 8
 105                                    putIndexOffsetLength i o l
 106              PortMsg p -> do putWord32be 3
 107                              putWord8 9
 108                              putWord16be (fromIntegral p)
 109              ExtendedMsg t b-> do putWord32be (fromIntegral blockLen)
 110                                   putWord8 20
 111                                   putWord8 (fromIntegral t) -- 0 => handshake msg
 112                                   -- actual extension msg follows
 113                                   mapM_ putWord8 blockList
 114                                     where blockList = unpack b
 115                                           blockLen  = length blockList
 116
 117     where putIndexOffsetLength i o l = do
 118             putWord32be (fromIntegral i)
 119             putWord32be (fromIntegral o)
 120             putWord32be (fromIntegral l)
 121
 122   get = do
 123     l <- getWord32be
 124     msgid <- getWord8
 125     case msgid of
 126      0 -> return ChokeMsg
 127      1 -> return UnChokeMsg
 128      2 -> return InterestedMsg
 129      3 -> return NotInterestedMsg
 130      4 -> fmap (HaveMsg . fromIntegral) getWord32be
 131      5 -> fmap (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 132      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 133        where getInteger = fromIntegral <$> getWord32be
 134      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 135        where getInteger = fromIntegral <$> getWord32be
 136      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 137        where getInteger = fromIntegral <$> getWord32be
 138      9 -> fmap (PortMsg . fromIntegral) getWord16be
 139      _ -> error ("unknown message ID: " ++ show msgid)
 140
 141 getMsg :: Handle -> IO PeerMsg
 142 getMsg h = do
 143   lBS <- hGet h 4
 144   let l = bsToInt lBS
 145   if l == 0
 146     then return KeepAliveMsg
 147     else do
 148     msg <- hGet h l
 149     return $ decode $ fromStrict $ concat [lBS, msg]
 150
 151 sendMsg :: Handle -> PeerMsg -> IO ()
 152 sendMsg h msg = hPut h bsMsg
 153   where bsMsg = toStrict $ encode msg
 154
 155 genHandshakeMsg :: ByteString -> String -> ByteString
 156 genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved1, reserved2, reserved3, infoHash, peerID]
 157   where pstrlen = singleton 19
 158         pstr = BC.pack "BitTorrent protocol"
 159         reserved1 = BC.replicate 5 '\0'
 160         reserved2 = singleton 0x10 -- support extension protocol
 161         reserved3 = BC.replicate 2 '\0'
 162         peerID = BC.pack peer_id
 163
 164 bsToInt :: ByteString -> Int
 165 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 166
 167 makePeer :: ByteString -> Peer
 168 makePeer peer = Peer (toIP ip') (toPort port')
 169   where (ip', port') = splitAt 4 peer