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