src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      handlePeerMsgs
   5     ) where
   6
   7 import Prelude hiding (lookup, concat, replicate, splitAt, empty, writeFile)
   8
   9 import System.IO (Handle, BufferMode(..), hSetBuffering)
  10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
  11 import Data.ByteString.Lazy (fromStrict, fromChunks)
  12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
  13 import Network (connectTo, PortID(..))
  14 import Data.Binary (Binary(..), decode)
  15 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  16 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  17 import Control.Monad (replicateM, liftM, forever)
  18 import Control.Applicative ((<$>), liftA3)
  19 import Data.Bits
  20 import Data.Word (Word8)
  21
  22 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  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
 149 bsToInt :: ByteString -> Int
 150 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 151
 152 bitfieldToList :: [Word8] -> [Integer]
 153 bitfieldToList bs = go bs 0
 154   where go [] _ = []
 155         go (b:bs') pos =
 156           let setBits = [pos*8 + (toInteger i) | i <- [0..8], testBit b i]
 157           in
 158            setBits ++ (go bs' (pos + 1))
 159
 160 -- downloadPiece :: Integer -> Handle -> IO ()
 161
 162 createDummyFile :: FilePath -> Int -> IO ()
 163 createDummyFile path size = do
 164   writeFile path (BC.replicate size '\0')
 165
 166 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 167 -- recvMsg :: Peer -> Handle -> Msg
 168 msgLoop :: Handle -> ByteString -> IO ()
 169 msgLoop h pieceHash =
 170   let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 171   in
 172    forever $ do
 173      msg <- getMsg h
 174      putStrLn $ "got a " ++ show msg
 175      case msg of
 176       BitFieldMsg bss -> do
 177         let pieceList = bitfieldToList (unpack bss)
 178         putStrLn (show pieceList)
 179         -- download each of the piece in order
 180       _ -> putStrLn (show msg)
 181
 182 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 183 handlePeerMsgs p m peerId logFn = do
 184   h <- handShake p (infoHash m) peerId
 185   logFn $ "handShake"
 186   msgLoop h (pieces (info m))
 187