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