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, unpack, concat, hGet, hPut, singleton)
  11 import Data.ByteString.Char8 (replicate, pack)
  12 import Network (connectTo, PortID(..))
  13 import Data.Binary (Binary(..))
  14 import Data.Binary.Put (putWord32be, putWord8)
  15
  16 type ID = String
  17 type IP = String
  18 type Port = Integer
  19
  20 data PeerState = PeerState { handle :: Handle
  21                            , am_choking :: Bool
  22                            , am_interested :: Bool
  23                            , peer_choking :: Bool
  24                            , peer_interested :: Bool }
  25
  26 -- Maintain info on every piece and the current state of it.
  27 -- should probably be a TVar.
  28 type Pieces = [PieceData]
  29
  30 data PieceState = Pending
  31                 | InProgress
  32                 | Have
  33                 deriving (Show)
  34
  35 data PieceData = PieceData { index :: Int           -- ^ Piece number
  36                            , peers :: [Peer]        -- ^ list of peers who have this piece
  37                            , state :: PieceState }  -- ^ state of the piece from download perspective.
  38
  39 -- | Peer is a PeerID, IP address, port tuple
  40 data Peer = Peer ID IP Port
  41           deriving (Show, Eq)
  42
  43 data PeerMsg = KeepAliveMsg
  44              | ChokeMsg
  45              | UnChokeMsg
  46              | InterestedMsg
  47              | NotInterestedMsg
  48              | HaveMsg Integer
  49              | BitFieldMsg ByteString
  50              | RequestMsg Integer Integer Integer
  51              | PieceMsg Integer Integer ByteString
  52              | CancelMsg Integer Integer Integer
  53              | PortMsg Port
  54              deriving (Show)
  55
  56 genHandShakeMsg :: ByteString -> String -> ByteString
  57 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  58   where pstrlen = singleton 19
  59         pstr = pack "BitTorrent protocol"
  60         reserved = replicate 8 '\0'
  61         peerID = pack peer_id
  62
  63 handShake :: Peer -> ByteString -> String -> IO Handle
  64 handShake (Peer _ ip port) infoHash peerid = do
  65   let hs = genHandShakeMsg infoHash peerid
  66   handle <- connectTo ip (PortNumber (fromIntegral port))
  67   hSetBuffering handle LineBuffering
  68   hPut handle hs
  69   rlenBS <- hGet handle 1
  70   let rlen = fromIntegral $ (unpack rlenBS) !! 0
  71   hGet handle rlen
  72   return handle
  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 index -> do putWord32be 5
  86                                  putWord8 4
  87                                  putWord32be (fromIntegral index)
  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                                     putWord32be (fromIntegral i)
  96                                     putWord32be (fromIntegral o)
  97                                     putWord32be (fromIntegral l)
  98              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
  99                                   putWord8 7
 100                                   putWord32be (fromIntegral i)
 101                                   putWord32be (fromIntegral o)
 102                                   mapM_ putWord8 blockList
 103                                     where blockList = unpack b
 104                                           blocklen = length blockList
 105              CancelMsg i o l -> undefined
 106              PortMsg p -> undefined
 107   get = undefined
 108
 109 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 110 -- recvMsg :: Peer -> Handle -> Msg