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