src/FuncTorrent/Peer.hs

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