src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      PieceMap,
   5      handlePeerMsgs,
   6      bytesDownloaded,
   7      initPieceMap,
   8      pieceMapFromFile
   9     ) where
  10
  11 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop, filter)
  12
  13 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
  14 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
  15 import qualified Data.ByteString.Char8 as BC (length)
  16 import Network (connectTo, PortID(..))
  17 import Control.Monad.State
  18 import Data.Bits
  19 import Data.Word (Word8)
  20 import Data.Map (Map, fromList, toList, (!), mapWithKey, traverseWithKey, adjust, filter)
  21 import Safe (headMay)
  22
  23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  24 import FuncTorrent.Utils (splitN, splitNum, writeFileAtOffset, readFileAtOffset, verifyHash)
  25 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
  26
  27 data PState = PState { handle :: Handle
  28                      , peer :: Peer
  29                      , meChoking :: Bool
  30                      , meInterested :: Bool
  31                      , heChoking :: Bool
  32                      , heInterested :: Bool}
  33
  34 type PeerState = State PState
  35
  36 data PieceDlState = Pending
  37                   | Downloading
  38                   | Have
  39                   deriving (Show, Eq)
  40
  41 -- todo - map with index to a new data structure (peers who have that piece and state)
  42 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  43                            , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
  44                            , hash  :: ByteString    -- ^ piece hash
  45                            , len :: Integer }       -- ^ piece length
  46
  47 -- which piece is with which peers
  48 type PieceMap = Map Integer PieceData
  49
  50
  51 -- Make the initial Piece map, with the assumption that no peer has the
  52 -- piece and that every piece is pending download.
  53 initPieceMap :: ByteString  -> Integer -> Integer -> PieceMap
  54 initPieceMap pieceHash fileLen pieceLen = fromList kvs
  55   where
  56     numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
  57     kvs = [(i, PieceData { peers = []
  58                          , dlstate = Pending
  59                          , hash = h
  60                          , len = pLen })
  61           | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  62     hashes = splitN 20 pieceHash
  63     pLengths = (splitNum fileLen pieceLen)
  64
  65 pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
  66 pieceMapFromFile filePath pieceMap = do
  67   traverseWithKey f pieceMap
  68     where
  69       f k v = do
  70         let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
  71         isHashValid <- (flip verifyHash) (hash v) <$> (readFileAtOffset filePath offset (len v))
  72         if isHashValid
  73           then return $ v { dlstate = Have }
  74           else return $ v
  75
  76 havePiece :: PieceMap -> Integer -> Bool
  77 havePiece pm index =
  78   dlstate (pm ! index) == Have
  79
  80 connectToPeer :: Peer -> IO Handle
  81 connectToPeer (Peer _ ip port) = do
  82   h <- connectTo ip (PortNumber (fromIntegral port))
  83   hSetBuffering h LineBuffering
  84   return h
  85
  86 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
  87 doHandshake True h peer infoHash peerid = do
  88   let hs = genHandshakeMsg infoHash peerid
  89   hPut h hs
  90   putStrLn $ "--> handhake to peer: " ++ show peer
  91   _ <- hGet h (length (unpack hs))
  92   putStrLn $ "<-- handshake from peer: " ++ show peer
  93   return ()
  94 doHandshake False h peer infoHash peerid = do
  95   let hs = genHandshakeMsg infoHash peerid
  96   putStrLn $ "waiting for a handshake"
  97   hsMsg <- hGet h (length (unpack hs))
  98   putStrLn $ "<-- handshake from peer: " ++ show peer
  99   let rxInfoHash = take 20 $ drop 28 hsMsg
 100   if rxInfoHash /= infoHash
 101     then do
 102     putStrLn $ "infoHashes does not match"
 103     hClose h
 104     return ()
 105     else do
 106     _ <- hPut h hs
 107     putStrLn $ "--> handhake to peer: " ++ show peer
 108     return ()
 109
 110 bitfieldToList :: [Word8] -> [Integer]
 111 bitfieldToList bs = go bs 0
 112   where go [] _ = []
 113         go (b:bs') pos =
 114           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 115           in
 116            setBits ++ go bs' (pos + 1)
 117
 118 -- helper functions to manipulate PeerState
 119 toPeerState :: Handle
 120             -> Peer
 121             -> Bool  -- ^ meChoking
 122             -> Bool  -- ^ meInterested
 123             -> Bool  -- ^ heChoking
 124             -> Bool  -- ^ heInterested
 125             -> PState
 126 toPeerState h p meCh meIn heCh heIn =
 127   PState { handle = h
 128          , peer = p
 129          , heChoking = heCh
 130          , heInterested = heIn
 131          , meChoking = meCh
 132          , meInterested = meIn }
 133
 134 -- simple algorithm to pick piece.
 135 -- pick the first piece from 0 that is not downloaded yet.
 136 pickPiece :: PieceMap -> Maybe Integer
 137 pickPiece =
 138   (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
 139
 140 bytesDownloaded :: PieceMap -> Integer
 141 bytesDownloaded =
 142   sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
 143
 144 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 145 updatePieceAvailability pieceStatus p pieceList =
 146   mapWithKey (\k pd -> if k `elem` pieceList
 147                        then (pd { peers = p : peers pd })
 148                        else pd) pieceStatus
 149
 150 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> IO ()
 151 handlePeerMsgs p peerId m pieceMap isClient = do
 152   h <- connectToPeer p
 153   doHandshake isClient h p (infoHash m) peerId
 154   let pstate = toPeerState h p False False True True
 155       filePath = name (info m)
 156   _ <- runStateT (msgLoop pieceMap filePath) pstate
 157   return ()
 158
 159 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
 160 msgLoop pieceStatus file = do
 161   h <- gets handle
 162   st <- get
 163   case st of
 164     PState { meInterested = False, heChoking = True } -> do
 165       liftIO $ sendMsg h InterestedMsg
 166       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
 167       modify (\st -> st { meInterested = True })
 168       msgLoop pieceStatus file
 169     PState { meInterested = True, heChoking = False } ->
 170       case pickPiece pieceStatus of
 171         Nothing -> liftIO $ putStrLn "Nothing to download"
 172         Just workPiece -> do
 173           let pLen = len (pieceStatus ! workPiece)
 174           liftIO $ putStrLn $ "piece length = " ++ show pLen
 175           pBS <- liftIO $ downloadPiece h workPiece pLen
 176           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 177             then
 178             liftIO $ putStrLn $ "Hash mismatch"
 179             else do
 180             let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
 181             liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
 182             liftIO $ writeFileAtOffset file fileOffset pBS
 183             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
 184     _ -> do
 185       msg <- liftIO $ getMsg h
 186       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
 187       case msg of
 188         KeepAliveMsg -> do
 189           liftIO $ sendMsg h KeepAliveMsg
 190           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 191           msgLoop pieceStatus file
 192         BitFieldMsg bss -> do
 193           p <- gets peer
 194           let pieceList = bitfieldToList (unpack bss)
 195               pieceStatus' = updatePieceAvailability pieceStatus p pieceList
 196           liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
 197           -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 198           -- map with pieceIndex as the key and modify the value to add the peer.
 199           -- download each of the piece in order
 200           msgLoop pieceStatus' file
 201         UnChokeMsg -> do
 202           modify (\st -> st {heChoking = False })
 203           msgLoop pieceStatus file
 204
 205
 206 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 207 downloadPiece h index pieceLength = do
 208   let chunks = splitNum pieceLength 16384
 209   concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
 210                                               sendMsg h (RequestMsg index (i*pLen) pLen)
 211                                               putStrLn $ "--> " ++ "RequestMsg for Piece "
 212                                                 ++ show index ++ ", part: " ++ show i ++ " of length: "
 213                                                 ++ show pLen
 214                                               msg <- getMsg h
 215                                               case msg of
 216                                                 PieceMsg index begin block -> do
 217                                                   putStrLn $ " <-- PieceMsg for Piece: "
 218                                                     ++ show index
 219                                                     ++ ", offset: "
 220                                                     ++ show begin
 221                                                   return block
 222                                                 _ -> do
 223                                                   putStrLn "ignoring irrelevant msg"
 224                                                   return empty)
 225