src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      handShake
   5     ) where
   6
   7 import Prelude hiding (lookup, concat, replicate, splitAt)
   8
   9 import System.IO
  10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
  11 import qualified Data.ByteString.Char8 as BC (replicate, pack)
  12 import Network (connectTo, PortID(..))
  13 import Data.Binary (Binary(..))
  14 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  15 import Data.Binary.Get (getWord32be, getWord16be, getWord8)
  16 import Control.Monad (replicateM, liftM)
  17 import Control.Applicative ((<$>), liftA3)
  18
  19 type ID = String
  20 type IP = String
  21 type Port = Integer
  22
  23 data PeerState = PeerState { handle :: Handle
  24                            , am_choking :: Bool
  25                            , am_interested :: Bool
  26                            , peer_choking :: Bool
  27                            , peer_interested :: Bool }
  28
  29 -- Maintain info on every piece and the current state of it.
  30 -- should probably be a TVar.
  31 type Pieces = [PieceData]
  32
  33 data PieceState = Pending
  34                 | InProgress
  35                 | Have
  36                 deriving (Show)
  37
  38 data PieceData = PieceData { index :: Int           -- ^ Piece number
  39                            , peers :: [Peer]        -- ^ list of peers who have this piece
  40                            , state :: PieceState }  -- ^ state of the piece from download perspective.
  41
  42 -- | Peer is a PeerID, IP address, port tuple
  43 data Peer = Peer ID IP Port
  44           deriving (Show, Eq)
  45
  46 data PeerMsg = KeepAliveMsg
  47              | ChokeMsg
  48              | UnChokeMsg
  49              | InterestedMsg
  50              | NotInterestedMsg
  51              | HaveMsg Integer
  52              | BitFieldMsg ByteString
  53              | RequestMsg Integer Integer Integer
  54              | PieceMsg Integer Integer ByteString
  55              | CancelMsg Integer Integer Integer
  56              | PortMsg Port
  57              deriving (Show)
  58
  59 genHandShakeMsg :: ByteString -> String -> ByteString
  60 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  61   where pstrlen = singleton 19
  62         pstr = BC.pack "BitTorrent protocol"
  63         reserved = BC.replicate 8 '\0'
  64         peerID = BC.pack peer_id
  65
  66 handShake :: Peer -> ByteString -> String -> IO Handle
  67 handShake (Peer _ ip port) infoHash peerid = do
  68   let hs = genHandShakeMsg infoHash peerid
  69   handle <- connectTo ip (PortNumber (fromIntegral port))
  70   hSetBuffering handle LineBuffering
  71   hPut handle hs
  72   rlenBS <- hGet handle 1
  73   let rlen = fromIntegral $ (unpack rlenBS) !! 0
  74   hGet handle rlen
  75   return handle
  76
  77 instance Binary PeerMsg where
  78   put msg = case msg of
  79              KeepAliveMsg -> putWord32be 0
  80              ChokeMsg -> do putWord32be 1
  81                             putWord8 0
  82              UnChokeMsg -> do putWord32be 1
  83                               putWord8 1
  84              InterestedMsg -> do putWord32be 1
  85                                  putWord8 2
  86              NotInterestedMsg -> do putWord32be 1
  87                                     putWord8 3
  88              HaveMsg index -> do putWord32be 5
  89                                  putWord8 4
  90                                  putWord32be (fromIntegral index)
  91              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
  92                                   putWord8 5
  93                                   mapM_ putWord8 bfList
  94                                     where bfList = unpack bf
  95                                           bfListLen = length bfList
  96              RequestMsg i o l -> do putWord32be 13
  97                                     putWord8 6
  98                                     putWord32be (fromIntegral i)
  99                                     putWord32be (fromIntegral o)
 100                                     putWord32be (fromIntegral l)
 101              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 102                                   putWord8 7
 103                                   putWord32be (fromIntegral i)
 104                                   putWord32be (fromIntegral o)
 105                                   mapM_ putWord8 blockList
 106                                     where blockList = unpack b
 107                                           blocklen = length blockList
 108              CancelMsg i o l -> do putWord32be 13
 109                                    putWord8 8
 110                                    putWord32be (fromIntegral i)
 111                                    putWord32be (fromIntegral o)
 112                                    putWord32be (fromIntegral l)
 113              PortMsg p -> do putWord32be 3
 114                              putWord8 9
 115                              putWord16be (fromIntegral p)
 116   get = do
 117     l <- getWord32be
 118     id <- getWord8
 119     case id of
 120      0 -> return ChokeMsg
 121      1 -> return UnChokeMsg
 122      2 -> return InterestedMsg
 123      3 -> return NotInterestedMsg
 124      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 125      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 126      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 127        where getInteger = fromIntegral <$> getWord32be
 128      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 129        where getInteger = fromIntegral <$> getWord32be
 130      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 131        where getInteger = fromIntegral <$> getWord32be
 132      9 -> liftM (PortMsg . fromIntegral) getWord16be
 133      _ -> error "unknown message ID"
 134
 135 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 136 -- recvMsg :: Peer -> Handle -> Msg