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