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(..), IOMode(..), SeekMode(..), withFile, hSeek, 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 -- write into a file at a specific offet
 192 writeFileAtOffset :: FilePath -> Integer -> ByteString -> IO ()
 193 writeFileAtOffset path offset block =
 194   withFile path WriteMode $ (\h -> do
 195                                 _ <- hSeek h AbsoluteSeek offset
 196                                 hPut h block)
 197
 198 -- recvMsg :: Peer -> Handle -> Msg
 199 msgLoop :: PeerState -> PieceMap -> IO ()
 200 msgLoop pState pieceStatus | meInterested pState == False &&
 201                             heChoking pState == True = do
 202                               -- if me NOT Interested and she is Choking, tell her that
 203                               -- I am interested.
 204                               let h = handle pState
 205                               sendMsg h InterestedMsg
 206                               putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
 207                               msgLoop (pState { meInterested = True }) pieceStatus
 208                           | meInterested pState == True &&
 209                             heChoking pState == False =
 210                               -- if me Interested and she not Choking, send her a request
 211                               -- for a piece.
 212                               case pickPiece pieceStatus of
 213                                Nothing -> putStrLn "Nothing to download"
 214                                Just workPiece -> do
 215                                  let pLen = len (pieceStatus ! workPiece)
 216                                  _ <- downloadPiece (handle pState) workPiece pLen
 217                                  -- TODO: verify the hash
 218                                  msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
 219                           | otherwise = do
 220                               msg <- getMsg (handle pState)
 221                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
 222                               case msg of
 223                                KeepAliveMsg -> do
 224                                  sendMsg (handle pState) KeepAliveMsg
 225                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
 226                                  msgLoop pState pieceStatus
 227                                BitFieldMsg bss -> do
 228                                  let pieceList = bitfieldToList (unpack bss)
 229                                      pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
 230                                  print pieceList
 231                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 232                                  -- map with pieceIndex as the key and modify the value to add the peer.
 233                                  -- download each of the piece in order
 234                                  msgLoop pState pieceStatus'
 235                                UnChokeMsg -> do
 236                                  msgLoop (pState { heChoking = False }) pieceStatus
 237                                _ -> do
 238                                  msgLoop pState pieceStatus
 239
 240 -- simple algorithm to pick piece.
 241 -- pick the first piece from 0 that is not downloaded yet.
 242 pickPiece :: PieceMap -> Maybe Integer
 243 pickPiece m =
 244   let pieceList = toList m
 245       allPending = filter (\(_, v) -> state v == Pending) pieceList
 246   in
 247    case allPending of
 248     [] -> Nothing
 249     ((i, _):_) -> Just i
 250
 251 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 252 updatePieceAvailability pieceStatus p pieceList =
 253   mapWithKey (\k pd -> if k `elem` pieceList
 254                        then (pd { peers = p : (peers pd) })
 255                        else pd) pieceStatus
 256
 257 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 258 handlePeerMsgs p m peerId = do
 259   h <- handShake p (infoHash m) peerId
 260   let state = PeerState { handle = h
 261                         , peer = p
 262                         , heInterested = False
 263                         , heChoking = True
 264                         , meInterested = False
 265                         , meChoking = True }
 266       pieceHash = pieces (info m)
 267       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 268       pLen = pieceLength (info m)
 269       fileLen = lengthInBytes (info m)
 270       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 271   msgLoop state pieceStatus
 272
 273 downloadPiece :: Handle -> Integer -> Integer -> IO [ByteString]
 274 downloadPiece h index pieceLength = do
 275   let chunks = splitNum pieceLength 16384
 276   forM (zip [0..] chunks) (\(i, pLen) -> do
 277                               sendMsg h (RequestMsg index (i*pLen) pLen)
 278                               putStrLn $ "--> " ++ "RequestMsg for Piece "
 279                                 ++ (show index) ++ ", part: " ++ show i ++ " of length: "
 280                                 ++ show pLen
 281                               PieceMsg index begin block <- getMsg h
 282                               putStrLn $ " <-- PieceMsg for Piece: "
 283                                 ++ show index
 284                                 ++ ", offset: "
 285                                 ++ show begin
 286                               return block)