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