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)
   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)
  17 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
  18
  19 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  20 import FuncTorrent.Utils (splitN, splitNum)
  21 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
  22 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
  23
  24 data PState = PState { handle :: Handle
  25                      , peer :: Peer
  26                      , meChoking :: Bool
  27                      , meInterested :: Bool
  28                      , heChoking :: Bool
  29                      , heInterested :: Bool}
  30
  31 type PeerState = State PState
  32
  33 data PieceDlState = Pending
  34                   | InProgress
  35                   | Have
  36                   deriving (Show, Eq)
  37
  38 -- todo - map with index to a new data structure (peers who have that piece amd state)
  39 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  40                            , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
  41                            , hash  :: ByteString    -- ^ piece hash
  42                            , len :: Integer }       -- ^ piece length
  43
  44 -- which piece is with which peers
  45 type PieceMap = Map Integer PieceData
  46
  47
  48 -- Make the initial Piece map, with the assumption that no peer has the
  49 -- piece and that every piece is pending download.
  50 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
  51 mkPieceMap numPieces pieceHash pLengths = fromList kvs
  52   where kvs = [(i, PieceData { peers = []
  53                              , dlstate = Pending
  54                              , hash = h
  55                              , len = pLen })
  56               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  57         hashes = splitN 20 pieceHash
  58
  59 havePiece :: PieceMap -> Integer -> Bool
  60 havePiece pm index =
  61   dlstate (pm ! index) == Have
  62
  63 connectToPeer :: Peer -> IO Handle
  64 connectToPeer (Peer _ ip port) = do
  65   h <- connectTo ip (PortNumber (fromIntegral port))
  66   hSetBuffering h LineBuffering
  67   return h
  68
  69 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
  70 doHandshake h peer infoHash peerid = do
  71   let hs = genHandshakeMsg infoHash peerid
  72   hPut h hs
  73   putStrLn $ "--> handhake to peer: " ++ show peer
  74   _ <- hGet h (length (unpack hs))
  75   putStrLn $ "<-- handshake from peer: " ++ show peer
  76   return ()
  77
  78 bitfieldToList :: [Word8] -> [Integer]
  79 bitfieldToList bs = go bs 0
  80   where go [] _ = []
  81         go (b:bs') pos =
  82           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
  83           in
  84            setBits ++ go bs' (pos + 1)
  85
  86 -- helper functions to manipulate PeerState
  87 toPeerState :: Handle
  88             -> Peer
  89             -> Bool  -- ^ meChoking
  90             -> Bool  -- ^ meInterested
  91             -> Bool  -- ^ heChoking
  92             -> Bool  -- ^ heInterested
  93             -> PState
  94 toPeerState h p meCh meIn heCh heIn =
  95   PState { handle = h
  96          , peer = p
  97          , heChoking = heCh
  98          , heInterested = heIn
  99          , meChoking = meCh
 100          , meInterested = meIn }
 101
 102 -- simple algorithm to pick piece.
 103 -- pick the first piece from 0 that is not downloaded yet.
 104 pickPiece :: PieceMap -> Maybe Integer
 105 pickPiece m =
 106   let pieceList = toList m
 107       allPending = filter (\(_, v) -> dlstate v == Pending) pieceList
 108   in
 109    case allPending of
 110     [] -> Nothing
 111     ((i, _):_) -> Just i
 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 True False 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