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 amd 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 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 110 updatePieceAvailability pieceStatus p pieceList =
 111   mapWithKey (\k pd -> if k `elem` pieceList
 112                        then (pd { peers = p : peers pd })
 113                        else pd) pieceStatus
 114
 115 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 116 handlePeerMsgs p m peerId = do
 117   h <- connectToPeer p
 118   doHandshake h p (infoHash m) peerId
 119   let pstate = toPeerState h p False False True True
 120       pieceHash = pieces (info m)
 121       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 122       pLen = pieceLength (info m)
 123       fileLen = lengthInBytes (info m)
 124       fileName = name (info m)
 125       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 126   createDummyFile fileName (fromIntegral fileLen)
 127   (r, _) <- runStateT (msgLoop pieceStatus fileName) pstate
 128   return ()
 129
 130 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
 131 msgLoop pieceStatus file = do
 132   h <- gets handle
 133   st <- get
 134   case st of
 135     PState { meInterested = False, heChoking = True } -> do
 136       liftIO $ sendMsg h InterestedMsg
 137       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
 138       modify (\st -> st { meInterested = True })
 139       msgLoop pieceStatus file
 140     PState { meInterested = True, heChoking = False } -> do
 141       case pickPiece pieceStatus of
 142         Nothing -> liftIO $ putStrLn "Nothing to download"
 143         Just workPiece -> do
 144           let pLen = len (pieceStatus ! workPiece)
 145           liftIO $ putStrLn $ "piece length = " ++ show pLen
 146           pBS <- liftIO $ downloadPiece h workPiece pLen
 147           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 148             then
 149             liftIO $ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
 150             else do
 151             let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
 152             liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
 153             liftIO $ writeFileAtOffset file fileOffset pBS
 154             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
 155     _ -> do
 156       msg <- liftIO $ getMsg h
 157       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
 158       case msg of
 159         KeepAliveMsg -> do
 160           liftIO $ sendMsg h KeepAliveMsg
 161           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 162           msgLoop pieceStatus file
 163         BitFieldMsg bss -> do
 164           p <- gets peer
 165           let pieceList = bitfieldToList (unpack bss)
 166               pieceStatus' = updatePieceAvailability pieceStatus p pieceList
 167           liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
 168           -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 169           -- map with pieceIndex as the key and modify the value to add the peer.
 170           -- download each of the piece in order
 171           msgLoop pieceStatus' file
 172         UnChokeMsg -> do
 173           modify (\st -> st {heChoking = False })
 174           msgLoop pieceStatus file
 175
 176
 177 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 178 downloadPiece h index pieceLength = do
 179   let chunks = splitNum pieceLength 16384
 180   liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
 181                                              sendMsg h (RequestMsg index (i*pLen) pLen)
 182                                              putStrLn $ "--> " ++ "RequestMsg for Piece "
 183                                                ++ show index ++ ", part: " ++ show i ++ " of length: "
 184                                                ++ show pLen
 185                                              msg <- getMsg h
 186                                              case msg of
 187                                               PieceMsg index begin block -> do
 188                                                 putStrLn $ " <-- PieceMsg for Piece: "
 189                                                   ++ show index
 190                                                   ++ ", offset: "
 191                                                   ++ show begin
 192                                                 return block
 193                                               _ -> do
 194                                                 putStrLn "ignoring irrelevant msg"
 195                                                 return empty)
 196
 197 verifyHash :: ByteString -> ByteString -> Bool
 198 verifyHash bs pieceHash =
 199   take 20 (SHA1.hash bs) == pieceHash