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, 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, toStrict)
  12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
  13 import Network (connectTo, PortID(..))
  14 import Data.Binary (Binary(..), decode, encode)
  15 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  16 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  17 import Control.Monad (replicateM, liftM, forM)
  18 import Control.Applicative ((<$>), liftA3)
  19 import Data.Bits
  20 import Data.Word (Word8)
  21 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
  22
  23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  24 import FuncTorrent.Utils (splitN, splitNum)
  25
  26 type ID = String
  27 type IP = String
  28 type Port = Integer
  29
  30 -- PeerState is a misnomer
  31 data PeerState = PeerState { handle :: Handle
  32                            , peer :: Peer
  33                            , meChoking :: Bool
  34                            , meInterested :: Bool
  35                            , heChoking :: Bool
  36                            , heInterested :: Bool}
  37
  38 data PieceDlState = Pending
  39                   | InProgress
  40                   | Have
  41                   deriving (Show, Eq)
  42
  43 -- todo - map with index to a new data structure (peers who have that piece amd state)
  44 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  45                            , state :: PieceDlState  -- ^ state of the piece from download perspective.
  46                            , hash  :: ByteString    -- ^ piece hash
  47                            , len :: Integer }       -- ^ piece length
  48
  49 -- which piece is with which peers
  50 type PieceMap = Map Integer PieceData
  51
  52 -- | Peer is a PeerID, IP address, port tuple
  53 data Peer = Peer ID IP Port
  54           deriving (Show, Eq)
  55
  56 data PeerMsg = KeepAliveMsg
  57              | ChokeMsg
  58              | UnChokeMsg
  59              | InterestedMsg
  60              | NotInterestedMsg
  61              | HaveMsg Integer
  62              | BitFieldMsg ByteString
  63              | RequestMsg Integer Integer Integer
  64              | PieceMsg Integer Integer ByteString
  65              | CancelMsg Integer Integer Integer
  66              | PortMsg Port
  67              deriving (Show)
  68
  69 -- Make the initial Piece map, with the assumption that no peer has the
  70 -- piece and that every piece is pending download.
  71 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
  72 mkPieceMap numPieces pieceHash pLengths = fromList kvs
  73   where kvs = [(i, PieceData { peers = []
  74                              , state = Pending
  75                              , hash = h
  76                              , len = pLen })
  77               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  78         hashes = splitN (fromIntegral numPieces) pieceHash
  79
  80 havePiece :: PieceMap -> Integer -> Bool
  81 havePiece pm index =
  82   state (pm ! index) == Have
  83
  84 genHandShakeMsg :: ByteString -> String -> ByteString
  85 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  86   where pstrlen = singleton 19
  87         pstr = BC.pack "BitTorrent protocol"
  88         reserved = BC.replicate 8 '\0'
  89         peerID = BC.pack peer_id
  90
  91 handShake :: Peer -> ByteString -> String -> IO Handle
  92 handShake peer@(Peer _ ip port) infoHash peerid = do
  93   let hs = genHandShakeMsg infoHash peerid
  94   h <- connectTo ip (PortNumber (fromIntegral port))
  95   hSetBuffering h LineBuffering
  96   hPut h hs
  97   putStrLn $ "--> handhake to peer: " ++ show peer
  98   _ <- hGet h (length (unpack hs))
  99   putStrLn $ "<-- handshake from peer: " ++ show peer
 100   return h
 101
 102 instance Binary PeerMsg where
 103   put msg = case msg of
 104              KeepAliveMsg -> putWord32be 0
 105              ChokeMsg -> do putWord32be 1
 106                             putWord8 0
 107              UnChokeMsg -> do putWord32be 1
 108                               putWord8 1
 109              InterestedMsg -> do putWord32be 1
 110                                  putWord8 2
 111              NotInterestedMsg -> do putWord32be 1
 112                                     putWord8 3
 113              HaveMsg i -> do putWord32be 5
 114                              putWord8 4
 115                              putWord32be (fromIntegral i)
 116              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
 117                                   putWord8 5
 118                                   mapM_ putWord8 bfList
 119                                     where bfList = unpack bf
 120                                           bfListLen = length bfList
 121              RequestMsg i o l -> do putWord32be 13
 122                                     putWord8 6
 123                                     putWord32be (fromIntegral i)
 124                                     putWord32be (fromIntegral o)
 125                                     putWord32be (fromIntegral l)
 126              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 127                                   putWord8 7
 128                                   putWord32be (fromIntegral i)
 129                                   putWord32be (fromIntegral o)
 130                                   mapM_ putWord8 blockList
 131                                     where blockList = unpack b
 132                                           blocklen = length blockList
 133              CancelMsg i o l -> do putWord32be 13
 134                                    putWord8 8
 135                                    putWord32be (fromIntegral i)
 136                                    putWord32be (fromIntegral o)
 137                                    putWord32be (fromIntegral l)
 138              PortMsg p -> do putWord32be 3
 139                              putWord8 9
 140                              putWord16be (fromIntegral p)
 141   get = do
 142     l <- getWord32be
 143     msgid <- getWord8
 144     case msgid of
 145      0 -> return ChokeMsg
 146      1 -> return UnChokeMsg
 147      2 -> return InterestedMsg
 148      3 -> return NotInterestedMsg
 149      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 150      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 151      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 152        where getInteger = fromIntegral <$> getWord32be
 153      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 154        where getInteger = fromIntegral <$> getWord32be
 155      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 156        where getInteger = fromIntegral <$> getWord32be
 157      9 -> liftM (PortMsg . fromIntegral) getWord16be
 158      _ -> error ("unknown message ID: " ++ show msgid)
 159
 160 getMsg :: Handle -> IO PeerMsg
 161 getMsg h = do
 162   lBS <- hGet h 4
 163   let l = bsToInt lBS
 164   if l == 0
 165     then return KeepAliveMsg
 166     else do
 167     msg <- hGet h l
 168     return $ decode $ fromStrict $ concat [lBS, msg]
 169
 170 sendMsg :: Handle -> PeerMsg -> IO ()
 171 sendMsg h msg =
 172   let bsMsg = toStrict $ encode msg
 173   in
 174    hPut h bsMsg
 175
 176 bsToInt :: ByteString -> Int
 177 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 178
 179 bitfieldToList :: [Word8] -> [Integer]
 180 bitfieldToList bs = go bs 0
 181   where go [] _ = []
 182         go (b:bs') pos =
 183           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 184           in
 185            setBits ++ go bs' (pos + 1)
 186
 187 createDummyFile :: FilePath -> Int -> IO ()
 188 createDummyFile path size =
 189   writeFile path (BC.replicate size '\0')
 190
 191 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 192 -- recvMsg :: Peer -> Handle -> Msg
 193 msgLoop :: PeerState -> PieceMap -> IO ()
 194 msgLoop pState pieceStatus | meInterested pState == False &&
 195                             heChoking pState == True = do
 196                               -- if me NOT Interested and she is Choking, tell her that
 197                               -- I am interested.
 198                               let h = handle pState
 199                               sendMsg h InterestedMsg
 200                               putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
 201                               msgLoop (pState { meInterested = True }) pieceStatus
 202                           | meInterested pState == True &&
 203                             heChoking pState == False =
 204                               -- if me Interested and she not Choking, send her a request
 205                               -- for a piece.
 206                               case pickPiece pieceStatus of
 207                                Nothing -> putStrLn "Nothing to download"
 208                                Just workPiece -> do
 209                                  let pLen = len (pieceStatus ! workPiece)
 210                                  _ <- downloadPiece (handle pState) workPiece pLen
 211                                  -- sendMsg (handle state) (RequestMsg workPiece 0 pLen)
 212                                  -- putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state) ++ " of length: " ++ show pLen
 213                                  -- msg <- getMsg (handle state)
 214                                  -- putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 215                                  msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
 216                           | otherwise = do
 217                               msg <- getMsg (handle pState)
 218                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
 219                               case msg of
 220                                KeepAliveMsg -> do
 221                                  sendMsg (handle pState) KeepAliveMsg
 222                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
 223                                  msgLoop pState pieceStatus
 224                                BitFieldMsg bss -> do
 225                                  let pieceList = bitfieldToList (unpack bss)
 226                                      pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
 227                                  print pieceList
 228                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 229                                  -- map with pieceIndex as the key and modify the value to add the peer.
 230                                  -- download each of the piece in order
 231                                  msgLoop pState pieceStatus'
 232                                UnChokeMsg -> do
 233                                  msgLoop (pState { heChoking = False }) pieceStatus
 234                                _ -> do
 235                                  msgLoop pState pieceStatus
 236
 237 -- simple algorithm to pick piece.
 238 -- pick the first piece from 0 that is not downloaded yet.
 239 pickPiece :: PieceMap -> Maybe Integer
 240 pickPiece m =
 241   let pieceList = toList m
 242       allPending = filter (\(_, v) -> state v == Pending) pieceList
 243   in
 244    case allPending of
 245     [] -> Nothing
 246     ((i, _):_) -> Just i
 247
 248 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 249 updatePieceAvailability pieceStatus p pieceList =
 250   mapWithKey (\k pd -> if k `elem` pieceList
 251                        then (pd { peers = p : (peers pd) })
 252                        else pd) pieceStatus
 253
 254 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 255 handlePeerMsgs p m peerId = do
 256   h <- handShake p (infoHash m) peerId
 257   let state = PeerState { handle = h
 258                         , peer = p
 259                         , heInterested = False
 260                         , heChoking = True
 261                         , meInterested = False
 262                         , meChoking = True }
 263       pieceHash = pieces (info m)
 264       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 265       pLen = pieceLength (info m)
 266       fileLen = lengthInBytes (info m)
 267       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 268   msgLoop state pieceStatus
 269
 270 downloadPiece :: Handle -> Integer -> Integer -> IO [PeerMsg]
 271 downloadPiece h index pieceLength = do
 272   let chunks = splitNum pieceLength 16384
 273   forM (zip [0..] chunks) (\(i, pLen) -> do
 274                               sendMsg h (RequestMsg index (i*pLen) pLen)
 275                               putStrLn $ "--> " ++ "RequestMsg for Piece "
 276                                 ++ (show index) ++ ", part: " ++ show i ++ " of length: "
 277                                 ++ show pLen
 278                               rMsg <- getMsg h
 279                               case rMsg of
 280                                PieceMsg index begin block ->
 281                                  putStrLn $ " <-- PieceMsg for Piece: "
 282                                  ++ show index
 283                                  ++ ", offset: "
 284                                  ++ show begin
 285                                _ -> putStrLn " <-- UnKnown msg from Peer"
 286                               return rMsg)