src/FuncTorrent/Peer.hs

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