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