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 isClient h peer infoHash peerid =
  88   let hs = genHandshakeMsg infoHash peerid
  89   in
  90    if isClient
  91    then do
  92      hPut h hs
  93      putStrLn $ "--> handhake to peer: " ++ show peer
  94      _ <- hGet h (length (unpack hs))
  95      putStrLn $ "<-- handshake from peer: " ++ show peer
  96      return ()
  97    else do
  98      putStrLn $ "waiting for a handshake"
  99      hsMsg <- hGet h (length (unpack hs))
 100      putStrLn $ "<-- handshake from peer: " ++ show peer
 101      let rxInfoHash = take 20 $ drop 28 hsMsg
 102      if rxInfoHash /= infoHash
 103        then do
 104        putStrLn $ "infoHashes does not match"
 105        hClose h
 106        return ()
 107        else do
 108        _ <- hPut h hs
 109        putStrLn $ "--> handhake to peer: " ++ show peer
 110        return ()
 111
 112 bitfieldToList :: [Word8] -> [Integer]
 113 bitfieldToList bs = go bs 0
 114   where go [] _ = []
 115         go (b:bs') pos =
 116           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 117           in
 118            setBits ++ go bs' (pos + 1)
 119
 120 -- helper functions to manipulate PeerState
 121 toPeerState :: Handle
 122             -> Peer
 123             -> Bool  -- ^ meChoking
 124             -> Bool  -- ^ meInterested
 125             -> Bool  -- ^ heChoking
 126             -> Bool  -- ^ heInterested
 127             -> PState
 128 toPeerState h p meCh meIn heCh heIn =
 129   PState { handle = h
 130          , peer = p
 131          , heChoking = heCh
 132          , heInterested = heIn
 133          , meChoking = meCh
 134          , meInterested = meIn }
 135
 136 -- simple algorithm to pick piece.
 137 -- pick the first piece from 0 that is not downloaded yet.
 138 pickPiece :: PieceMap -> Maybe Integer
 139 pickPiece =
 140   (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
 141
 142 bytesDownloaded :: PieceMap -> Integer
 143 bytesDownloaded =
 144   sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
 145
 146 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 147 updatePieceAvailability pieceStatus p pieceList =
 148   mapWithKey (\k pd -> if k `elem` pieceList
 149                        then (pd { peers = p : peers pd })
 150                        else pd) pieceStatus
 151
 152 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> IO ()
 153 handlePeerMsgs p peerId m pieceMap isClient = do
 154   h <- connectToPeer p
 155   doHandshake isClient h p (infoHash m) peerId
 156   let pstate = toPeerState h p False False True True
 157       filePath = name (info m)
 158   _ <- runStateT (msgLoop pieceMap filePath) pstate
 159   return ()
 160
 161 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
 162 msgLoop pieceStatus file = do
 163   h <- gets handle
 164   st <- get
 165   case st of
 166     PState { meInterested = False, heChoking = True } -> do
 167       liftIO $ sendMsg h InterestedMsg
 168       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
 169       modify (\st -> st { meInterested = True })
 170       msgLoop pieceStatus file
 171     PState { meInterested = True, heChoking = False } ->
 172       case pickPiece pieceStatus of
 173         Nothing -> liftIO $ putStrLn "Nothing to download"
 174         Just workPiece -> do
 175           let pLen = len (pieceStatus ! workPiece)
 176           liftIO $ putStrLn $ "piece length = " ++ show pLen
 177           pBS <- liftIO $ downloadPiece h workPiece pLen
 178           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 179             then
 180             liftIO $ putStrLn $ "Hash mismatch"
 181             else do
 182             let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
 183             liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
 184             liftIO $ writeFileAtOffset file fileOffset pBS
 185             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
 186     _ -> do
 187       msg <- liftIO $ getMsg h
 188       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
 189       case msg of
 190         KeepAliveMsg -> do
 191           liftIO $ sendMsg h KeepAliveMsg
 192           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 193           msgLoop pieceStatus file
 194         BitFieldMsg bss -> do
 195           p <- gets peer
 196           let pieceList = bitfieldToList (unpack bss)
 197               pieceStatus' = updatePieceAvailability pieceStatus p pieceList
 198           liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
 199           -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 200           -- map with pieceIndex as the key and modify the value to add the peer.
 201           -- download each of the piece in order
 202           msgLoop pieceStatus' file
 203         UnChokeMsg -> do
 204           modify (\st -> st {heChoking = False })
 205           msgLoop pieceStatus file
 206
 207
 208 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 209 downloadPiece h index pieceLength = do
 210   let chunks = splitNum pieceLength 16384
 211   concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
 212                                               sendMsg h (RequestMsg index (i*pLen) pLen)
 213                                               putStrLn $ "--> " ++ "RequestMsg for Piece "
 214                                                 ++ show index ++ ", part: " ++ show i ++ " of length: "
 215                                                 ++ show pLen
 216                                               msg <- getMsg h
 217                                               case msg of
 218                                                 PieceMsg index begin block -> do
 219                                                   putStrLn $ " <-- PieceMsg for Piece: "
 220                                                     ++ show index
 221                                                     ++ ", offset: "
 222                                                     ++ show begin
 223                                                   return block
 224                                                 _ -> do
 225                                                   putStrLn "ignoring irrelevant msg"
 226                                                   return empty)
 227