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, empty)
   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, length)
  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 Data.Bits
  21 import Data.Word (Word8)
  22
  23 type ID = String
  24 type IP = String
  25 type Port = Integer
  26
  27 data PeerState = PeerState { handle :: Handle
  28                            , am_choking :: Bool
  29                            , am_interested :: Bool
  30                            , peer_choking :: Bool
  31                            , peer_interested :: Bool}
  32
  33 -- Maintain info on every piece and the current state of it.
  34 -- should probably be a TVar.
  35 type Pieces = [PieceData]
  36
  37 data PieceState = Pending
  38                 | InProgress
  39                 | Have
  40                 deriving (Show)
  41
  42 data PieceData = PieceData { index :: Int           -- ^ Piece number
  43                            , peers :: [Peer]        -- ^ list of peers who have this piece
  44                            , state :: PieceState }  -- ^ state of the piece from download perspective.
  45
  46 -- | Peer is a PeerID, IP address, port tuple
  47 data Peer = Peer ID IP Port
  48           deriving (Show, Eq)
  49
  50 data PeerMsg = KeepAliveMsg
  51              | ChokeMsg
  52              | UnChokeMsg
  53              | InterestedMsg
  54              | NotInterestedMsg
  55              | HaveMsg Integer
  56              | BitFieldMsg ByteString
  57              | RequestMsg Integer Integer Integer
  58              | PieceMsg Integer Integer ByteString
  59              | CancelMsg Integer Integer Integer
  60              | PortMsg Port
  61              deriving (Show)
  62
  63 genHandShakeMsg :: ByteString -> String -> ByteString
  64 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  65   where pstrlen = singleton 19
  66         pstr = BC.pack "BitTorrent protocol"
  67         reserved = BC.replicate 8 '\0'
  68         peerID = BC.pack peer_id
  69
  70 handShake :: Peer -> ByteString -> String -> IO Handle
  71 handShake (Peer _ ip port) infoHash peerid = do
  72   let hs = genHandShakeMsg infoHash peerid
  73   h <- connectTo ip (PortNumber (fromIntegral port))
  74   hSetBuffering h LineBuffering
  75   hPut h hs
  76   rlenBS <- hGet h (length (unpack hs))
  77   putStrLn $ "got handshake from peer: " ++ show rlenBS
  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     msg <- hGet h l
 146     return $ decode $ fromStrict $ concat [lBS, msg]
 147
 148 bsToInt :: ByteString -> Int
 149 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 150
 151 bitfieldToList :: [Word8] -> Integer -> [Integer]
 152 bitfieldToList [] pos = []
 153 bitfieldToList (b:bs) pos =
 154   let setBits = [pos*8 + (toInteger i) | i <- [0..8], testBit b i]
 155   in
 156    setBits ++ (bitfieldToList bs (pos + 1))
 157
 158 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 159 -- recvMsg :: Peer -> Handle -> Msg
 160
 161 msgLoop :: Handle -> ByteString -> IO ()
 162 msgLoop h pieceHash =
 163   let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 164   in
 165    forever $ do
 166      msg <- getMsg h
 167      putStrLn $ "got a " ++ show msg
 168      case msg of
 169       BitFieldMsg bss -> do
 170                          let pieceList = bitfieldToList (unpack bss) 0
 171                          putStrLn (show pieceList)
 172                          -- doenload each of the piece
 173       _ -> putStrLn (show msg)