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