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