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 connectToPeer :: Peer -> IO Handle
  94 connectToPeer peer@(Peer _ ip port) = do
  95   h <- connectTo ip (PortNumber (fromIntegral port))
  96   hSetBuffering h LineBuffering
  97   return h
  98
  99 doHandShake :: Handle -> Peer -> ByteString -> String -> IO ()
 100 doHandShake h peer infoHash peerid = do
 101   let hs = genHandShakeMsg infoHash peerid
 102   hPut h hs
 103   putStrLn $ "--> handhake to peer: " ++ show peer
 104   _ <- hGet h (length (unpack hs))
 105   putStrLn $ "<-- handshake from peer: " ++ show peer
 106   return ()
 107
 108 instance Binary PeerMsg where
 109   put msg = case msg of
 110              KeepAliveMsg -> putWord32be 0
 111              ChokeMsg -> do putWord32be 1
 112                             putWord8 0
 113              UnChokeMsg -> do putWord32be 1
 114                               putWord8 1
 115              InterestedMsg -> do putWord32be 1
 116                                  putWord8 2
 117              NotInterestedMsg -> do putWord32be 1
 118                                     putWord8 3
 119              HaveMsg i -> do putWord32be 5
 120                              putWord8 4
 121                              putWord32be (fromIntegral i)
 122              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
 123                                   putWord8 5
 124                                   mapM_ putWord8 bfList
 125                                     where bfList = unpack bf
 126                                           bfListLen = length bfList
 127              RequestMsg i o l -> do putWord32be 13
 128                                     putWord8 6
 129                                     putWord32be (fromIntegral i)
 130                                     putWord32be (fromIntegral o)
 131                                     putWord32be (fromIntegral l)
 132              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 133                                   putWord8 7
 134                                   putWord32be (fromIntegral i)
 135                                   putWord32be (fromIntegral o)
 136                                   mapM_ putWord8 blockList
 137                                     where blockList = unpack b
 138                                           blocklen = length blockList
 139              CancelMsg i o l -> do putWord32be 13
 140                                    putWord8 8
 141                                    putWord32be (fromIntegral i)
 142                                    putWord32be (fromIntegral o)
 143                                    putWord32be (fromIntegral l)
 144              PortMsg p -> do putWord32be 3
 145                              putWord8 9
 146                              putWord16be (fromIntegral p)
 147   get = do
 148     l <- getWord32be
 149     msgid <- getWord8
 150     case msgid of
 151      0 -> return ChokeMsg
 152      1 -> return UnChokeMsg
 153      2 -> return InterestedMsg
 154      3 -> return NotInterestedMsg
 155      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 156      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 157      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 158        where getInteger = fromIntegral <$> getWord32be
 159      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 160        where getInteger = fromIntegral <$> getWord32be
 161      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 162        where getInteger = fromIntegral <$> getWord32be
 163      9 -> liftM (PortMsg . fromIntegral) getWord16be
 164      _ -> error ("unknown message ID: " ++ show msgid)
 165
 166 getMsg :: Handle -> IO PeerMsg
 167 getMsg h = do
 168   lBS <- hGet h 4
 169   let l = bsToInt lBS
 170   if l == 0
 171     then return KeepAliveMsg
 172     else do
 173     msg <- hGet h l
 174     return $ decode $ fromStrict $ concat [lBS, msg]
 175
 176 sendMsg :: Handle -> PeerMsg -> IO ()
 177 sendMsg h msg = hPut h bsMsg
 178   where bsMsg = toStrict $ encode msg
 179
 180 bsToInt :: ByteString -> Int
 181 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 182
 183 bitfieldToList :: [Word8] -> [Integer]
 184 bitfieldToList bs = go bs 0
 185   where go [] _ = []
 186         go (b:bs') pos =
 187           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 188           in
 189            setBits ++ go bs' (pos + 1)
 190
 191 -- recvMsg :: Peer -> Handle -> Msg
 192 msgLoop :: PeerState -> PieceMap -> IO ()
 193 msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = do
 194                                -- if me NOT Interested and she is Choking, tell her that
 195                                -- I am interested.
 196                                let h = handle pState
 197                                sendMsg h InterestedMsg
 198                                putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
 199                                msgLoop (pState { meInterested = True }) pieceStatus
 200                           | meInterested pState && not (heChoking pState) =
 201                               -- if me Interested and she not Choking, send her a request
 202                               -- for a piece.
 203                               case pickPiece pieceStatus of
 204                                Nothing -> putStrLn "Nothing to download"
 205                                Just workPiece -> do
 206                                  let pLen = len (pieceStatus ! workPiece)
 207                                  putStrLn $ "piece length = " ++ show pLen
 208                                  pBS <- downloadPiece (handle pState) workPiece pLen
 209                                  if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 210                                    then
 211                                    putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
 212                                    else do
 213                                    let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
 214                                    putStrLn $ "Write into file at offset: " ++ show fileOffset
 215                                    writeFileAtOffset "/tmp/download.file" fileOffset pBS
 216                                    msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
 217                           | otherwise = do
 218                               msg <- getMsg (handle pState)
 219                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
 220                               case msg of
 221                                KeepAliveMsg -> do
 222                                  sendMsg (handle pState) KeepAliveMsg
 223                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
 224                                  msgLoop pState pieceStatus
 225                                BitFieldMsg bss -> do
 226                                  let pieceList = bitfieldToList (unpack bss)
 227                                      pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
 228                                  putStrLn $ show (length pieceList) ++ " Pieces"
 229                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 230                                  -- map with pieceIndex as the key and modify the value to add the peer.
 231                                  -- download each of the piece in order
 232                                  msgLoop pState pieceStatus'
 233                                UnChokeMsg ->
 234                                  msgLoop (pState { heChoking = False }) pieceStatus
 235                                _ ->
 236                                  msgLoop pState pieceStatus
 237
 238 -- simple algorithm to pick piece.
 239 -- pick the first piece from 0 that is not downloaded yet.
 240 pickPiece :: PieceMap -> Maybe Integer
 241 pickPiece m =
 242   let pieceList = toList m
 243       allPending = filter (\(_, v) -> state v == Pending) pieceList
 244   in
 245    case allPending of
 246     [] -> Nothing
 247     ((i, _):_) -> Just i
 248
 249 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 250 updatePieceAvailability pieceStatus p pieceList =
 251   mapWithKey (\k pd -> if k `elem` pieceList
 252                        then (pd { peers = p : peers pd })
 253                        else pd) pieceStatus
 254
 255 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 256 handlePeerMsgs p m peerId = do
 257   h <- connectToPeer p
 258   doHandShake h p (infoHash m) peerId
 259   let state = PeerState { handle = h
 260                         , peer = p
 261                         , heInterested = False
 262                         , heChoking = True
 263                         , meInterested = False
 264                         , meChoking = True }
 265       pieceHash = pieces (info m)
 266       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 267       pLen = pieceLength (info m)
 268       fileLen = lengthInBytes (info m)
 269       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 270   createDummyFile "/tmp/download.file" (fromIntegral fileLen)
 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   liftM concat $ 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                                              msg <- getMsg h
 282                                              case msg of
 283                                               PieceMsg index begin block -> do
 284                                                 putStrLn $ " <-- PieceMsg for Piece: "
 285                                                   ++ show index
 286                                                   ++ ", offset: "
 287                                                   ++ show begin
 288                                                 return block
 289                                               _ -> do
 290                                                 putStrLn "ignoring irrelevant msg"
 291                                                 return empty)
 292
 293 verifyHash :: ByteString -> ByteString -> Bool
 294 verifyHash bs pieceHash =
 295   take 20 (SHA1.hash bs) == pieceHash