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)
  13 import qualified Data.ByteString.Char8 as BC (replicate, pack, readInt)
  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)
  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 1
  75   let rlen = fromIntegral $ (unpack rlenBS) !! 0
  76   hGet h rlen
  77   return h
  78
  79 instance Binary PeerMsg where
  80   put msg = case msg of
  81              KeepAliveMsg -> putWord32be 0
  82              ChokeMsg -> do putWord32be 1
  83                             putWord8 0
  84              UnChokeMsg -> do putWord32be 1
  85                               putWord8 1
  86              InterestedMsg -> do putWord32be 1
  87                                  putWord8 2
  88              NotInterestedMsg -> do putWord32be 1
  89                                     putWord8 3
  90              HaveMsg i -> do putWord32be 5
  91                              putWord8 4
  92                              putWord32be (fromIntegral i)
  93              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
  94                                   putWord8 5
  95                                   mapM_ putWord8 bfList
  96                                     where bfList = unpack bf
  97                                           bfListLen = length bfList
  98              RequestMsg i o l -> do putWord32be 13
  99                                     putWord8 6
 100                                     putWord32be (fromIntegral i)
 101                                     putWord32be (fromIntegral o)
 102                                     putWord32be (fromIntegral l)
 103              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 104                                   putWord8 7
 105                                   putWord32be (fromIntegral i)
 106                                   putWord32be (fromIntegral o)
 107                                   mapM_ putWord8 blockList
 108                                     where blockList = unpack b
 109                                           blocklen = length blockList
 110              CancelMsg i o l -> do putWord32be 13
 111                                    putWord8 8
 112                                    putWord32be (fromIntegral i)
 113                                    putWord32be (fromIntegral o)
 114                                    putWord32be (fromIntegral l)
 115              PortMsg p -> do putWord32be 3
 116                              putWord8 9
 117                              putWord16be (fromIntegral p)
 118   get = do
 119     l <- getWord32be
 120     msgid <- getWord8
 121     case msgid of
 122      0 -> return ChokeMsg
 123      1 -> return UnChokeMsg
 124      2 -> return InterestedMsg
 125      3 -> return NotInterestedMsg
 126      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 127      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 128      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 129        where getInteger = fromIntegral <$> getWord32be
 130      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 131        where getInteger = fromIntegral <$> getWord32be
 132      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 133        where getInteger = fromIntegral <$> getWord32be
 134      9 -> liftM (PortMsg . fromIntegral) getWord16be
 135      _ -> error "unknown message ID"
 136
 137 getMsg :: Handle -> IO PeerMsg
 138 getMsg h = do
 139   lBS <- hGet h 4
 140   let (Just (l, _)) = BC.readInt lBS
 141   msg <- hGet h l
 142   return $ decode $ fromStrict $ append lBS msg
 143
 144
 145 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 146 -- recvMsg :: Peer -> Handle -> Msg
 147
 148 msgLoop :: Handle -> IO ()
 149 msgLoop h = forever $ do
 150   msg <- getMsg h
 151   putStrLn $ "got a " ++ (show msg)