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