src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      handShake,
   5      msgLoop
   6     ) where
   7
   8 import Prelude hiding (lookup, concat, replicate, splitAt)
   9
  10 import System.IO
  11 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
  12 import Data.ByteString.Lazy (fromStrict, fromChunks)
  13 import qualified Data.ByteString.Char8 as BC (replicate, pack)
  14 import Network (connectTo, PortID(..))
  15 import Data.Binary (Binary(..), decode)
  16 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  17 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  18 import Control.Monad (replicateM, liftM, forever)
  19 import Control.Applicative ((<$>), liftA3)
  20
  21 type ID = String
  22 type IP = String
  23 type Port = Integer
  24
  25 data PeerState = PeerState { handle :: Handle
  26                            , am_choking :: Bool
  27                            , am_interested :: Bool
  28                            , peer_choking :: Bool
  29                            , peer_interested :: Bool}
  30
  31 -- Maintain info on every piece and the current state of it.
  32 -- should probably be a TVar.
  33 type Pieces = [PieceData]
  34
  35 data PieceState = Pending
  36                 | InProgress
  37                 | Have
  38                 deriving (Show)
  39
  40 data PieceData = PieceData { index :: Int           -- ^ Piece number
  41                            , peers :: [Peer]        -- ^ list of peers who have this piece
  42                            , state :: PieceState }  -- ^ state of the piece from download perspective.
  43
  44 -- | Peer is a PeerID, IP address, port tuple
  45 data Peer = Peer ID IP Port
  46           deriving (Show, Eq)
  47
  48 data PeerMsg = KeepAliveMsg
  49              | ChokeMsg
  50              | UnChokeMsg
  51              | InterestedMsg
  52              | NotInterestedMsg
  53              | HaveMsg Integer
  54              | BitFieldMsg ByteString
  55              | RequestMsg Integer Integer Integer
  56              | PieceMsg Integer Integer ByteString
  57              | CancelMsg Integer Integer Integer
  58              | PortMsg Port
  59              deriving (Show)
  60
  61 genHandShakeMsg :: ByteString -> String -> ByteString
  62 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  63   where pstrlen = singleton 19
  64         pstr = BC.pack "BitTorrent protocol"
  65         reserved = BC.replicate 8 '\0'
  66         peerID = BC.pack peer_id
  67
  68 handShake :: Peer -> ByteString -> String -> IO Handle
  69 handShake (Peer _ ip port) infoHash peerid = do
  70   let hs = genHandShakeMsg infoHash peerid
  71   h <- connectTo ip (PortNumber (fromIntegral port))
  72   hSetBuffering h LineBuffering
  73   hPut h hs
  74   rlenBS <- hGet h (length (unpack hs))
  75   putStrLn $ "got handshake from peer: " ++ show rlenBS
  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: " ++ show msgid)
 135
 136 getMsg :: Handle -> IO PeerMsg
 137 getMsg h = do
 138   lBS <- hGet h 4
 139   let l = bsToInt lBS
 140   if l == 0
 141     then return KeepAliveMsg
 142     else do
 143     putStrLn $ "len: " ++ show l
 144     msgID <- hGet h 1
 145     putStrLn $ "msg Type: " ++ show msgID
 146     msg <- hGet h (l - 1)
 147     return $ decode $ fromStrict $ concat [lBS, msgID, msg]
 148
 149 bsToInt :: ByteString -> Int
 150 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 151
 152 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 153 -- recvMsg :: Peer -> Handle -> Msg
 154
 155 msgLoop :: Handle -> IO ()
 156 msgLoop h = forever $ do
 157   msg <- getMsg h
 158   putStrLn $ "got a " ++ show msg