src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      PieceMap,
   5      handlePeerMsgs
   6     ) where
   7
   8 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
   9
  10 import Control.Concurrent.Chan (writeChan)
  11 import Control.Monad.State
  12 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
  13 import Data.Bits
  14 import Data.Word (Word8)
  15 import Data.Map ((!), adjust)
  16 import Network (connectTo, PortID(..))
  17 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
  18
  19 import FuncTorrent.Metainfo (Metainfo(..))
  20 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
  21 import FuncTorrent.Utils (splitNum, verifyHash)
  22 import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
  23 import qualified FuncTorrent.FileSystem as FS (MsgChannel, Msg(..), Piece(..))
  24
  25 data PState = PState { handle :: Handle
  26                      , peer :: Peer
  27                      , meChoking :: Bool
  28                      , meInterested :: Bool
  29                      , heChoking :: Bool
  30                      , heInterested :: Bool}
  31
  32 havePiece :: PieceMap -> Integer -> Bool
  33 havePiece pm index =
  34   dlstate (pm ! index) == Have
  35
  36 connectToPeer :: Peer -> IO Handle
  37 connectToPeer (Peer _ ip port) = do
  38   h <- connectTo ip (PortNumber (fromIntegral port))
  39   hSetBuffering h LineBuffering
  40   return h
  41
  42 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
  43 doHandshake True h p infohash peerid = do
  44   let hs = genHandshakeMsg infohash peerid
  45   hPut h hs
  46   putStrLn $ "--> handhake to peer: " ++ show p
  47   _ <- hGet h (length (unpack hs))
  48   putStrLn $ "<-- handshake from peer: " ++ show p
  49   return ()
  50 doHandshake False h p infohash peerid = do
  51   let hs = genHandshakeMsg infohash peerid
  52   putStrLn "waiting for a handshake"
  53   hsMsg <- hGet h (length (unpack hs))
  54   putStrLn $ "<-- handshake from peer: " ++ show p
  55   let rxInfoHash = take 20 $ drop 28 hsMsg
  56   if rxInfoHash /= infohash
  57     then do
  58     putStrLn "infoHashes does not match"
  59     hClose h
  60     return ()
  61     else do
  62     _ <- hPut h hs
  63     putStrLn $ "--> handhake to peer: " ++ show p
  64     return ()
  65
  66 bitfieldToList :: [Word8] -> [Integer]
  67 bitfieldToList bs = go bs 0
  68   where go [] _ = []
  69         go (b:bs') pos =
  70           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
  71           in
  72            setBits ++ go bs' (pos + 1)
  73
  74 -- helper functions to manipulate PeerState
  75 toPeerState :: Handle
  76             -> Peer
  77             -> Bool  -- ^ meChoking
  78             -> Bool  -- ^ meInterested
  79             -> Bool  -- ^ heChoking
  80             -> Bool  -- ^ heInterested
  81             -> PState
  82 toPeerState h p meCh meIn heCh heIn =
  83   PState { handle = h
  84          , peer = p
  85          , heChoking = heCh
  86          , heInterested = heIn
  87          , meChoking = meCh
  88          , meInterested = meIn }
  89
  90 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
  91 handlePeerMsgs p peerId m pieceMap isClient c = do
  92   h <- connectToPeer p
  93   doHandshake isClient h p (infoHash m) peerId
  94   let pstate = toPeerState h p False False True True
  95   _ <- runStateT (msgLoop pieceMap c) pstate
  96   return ()
  97
  98 msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
  99 msgLoop pieceStatus msgchannel = do
 100   h <- gets handle
 101   st <- get
 102   case st of
 103     PState { meInterested = False, heChoking = True } -> do
 104       liftIO $ sendMsg h InterestedMsg
 105       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
 106       modify (\st' -> st' { meInterested = True })
 107       msgLoop pieceStatus msgchannel
 108     PState { meInterested = True, heChoking = False } ->
 109       case pickPiece pieceStatus of
 110         Nothing -> liftIO $ putStrLn "Nothing to download"
 111         Just workPiece -> do
 112           let pLen = len (pieceStatus ! workPiece)
 113           liftIO $ putStrLn $ "piece length = " ++ show pLen
 114           pBS <- liftIO $ downloadPiece h workPiece pLen
 115           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 116             then
 117             liftIO $ putStrLn "Hash mismatch"
 118             else do
 119             liftIO $ putStrLn $ "Write piece: " ++ show workPiece
 120             liftIO $ writeChan msgchannel $ FS.WritePiece (FS.Piece workPiece pBS)
 121             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
 122     _ -> do
 123       msg <- liftIO $ getMsg h
 124       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
 125       case msg of
 126         KeepAliveMsg -> do
 127           liftIO $ sendMsg h KeepAliveMsg
 128           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 129           msgLoop pieceStatus msgchannel
 130         BitFieldMsg bss -> do
 131           p <- gets peer
 132           let pieceList = bitfieldToList (unpack bss)
 133               pieceStatus' = updatePieceAvailability pieceStatus p pieceList
 134           liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
 135           -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 136           -- map with pieceIndex as the key and modify the value to add the peer.
 137           -- download each of the piece in order
 138           msgLoop pieceStatus' msgchannel
 139         UnChokeMsg -> do
 140           modify (\st' -> st' {heChoking = False })
 141           msgLoop pieceStatus msgchannel
 142         ChokeMsg -> do
 143           modify (\st' -> st' {heChoking = True })
 144           msgLoop pieceStatus msgchannel
 145         InterestedMsg -> do
 146           modify (\st' -> st' {heInterested = True})
 147           msgLoop pieceStatus msgchannel
 148         NotInterestedMsg -> do
 149           modify (\st' -> st' {heInterested = False})
 150           msgLoop pieceStatus msgchannel
 151         CancelMsg _ _ _ -> -- check if valid index, begin, length
 152           msgLoop pieceStatus msgchannel
 153         PortMsg _ ->
 154           msgLoop pieceStatus msgchannel
 155         -- handle RequestMsg, HaveMsg. No need to handle PieceMsg here.
 156         -- also BitFieldMsg
 157
 158
 159 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 160 downloadPiece h index pieceLength = do
 161   let chunks = splitNum pieceLength 16384
 162   concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
 163                                               sendMsg h (RequestMsg index (i*pLen) pLen)
 164                                               putStrLn $ "--> " ++ "RequestMsg for Piece "
 165                                                 ++ show index ++ ", part: " ++ show i ++ " of length: "
 166                                                 ++ show pLen
 167                                               msg <- getMsg h
 168                                               case msg of
 169                                                 PieceMsg index begin block -> do
 170                                                   putStrLn $ " <-- PieceMsg for Piece: "
 171                                                     ++ show index
 172                                                     ++ ", offset: "
 173                                                     ++ show begin
 174                                                   return block
 175                                                 _ -> do
 176                                                   putStrLn $ "ignoring irrelevant msg: " ++ show msg
 177                                                   return empty)
 178