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