From: Ramakrishnan Muthukrishnan Date: Thu, 3 Sep 2015 11:02:49 +0000 (+0530) Subject: Peer: Use Monad Transformers to simplify code X-Git-Url: https://git.rkrishnan.org/%5B/%5D%20/uri/flags/cyclelanguage?a=commitdiff_plain;h=755a5f299f8ff9768d98fd780dd8a51420ae3192;p=functorrent.git Peer: Use Monad Transformers to simplify code Wrap PState in a State Monad. Reorganization of the code in progress. --- diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs index afc2028..65372d9 100644 --- a/src/FuncTorrent/Peer.hs +++ b/src/FuncTorrent/Peer.hs @@ -10,7 +10,7 @@ import System.IO (Handle, BufferMode(..), hSetBuffering) import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty) import qualified Data.ByteString.Char8 as BC (length) import Network (connectTo, PortID(..)) -import Control.Monad (liftM, forM) +import Control.Monad.State import Data.Bits import Data.Word (Word8) import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust) @@ -21,13 +21,14 @@ import FuncTorrent.Utils (splitN, splitNum) import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset) import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg) --- PeerState is a misnomer -data PeerState = PeerState { handle :: Handle - , peer :: Peer - , meChoking :: Bool - , meInterested :: Bool - , heChoking :: Bool - , heInterested :: Bool} +data PState = PState { handle :: Handle + , peer :: Peer + , meChoking :: Bool + , meInterested :: Bool + , heChoking :: Bool + , heInterested :: Bool} + +type PeerState = State PState data PieceDlState = Pending | InProgress @@ -36,7 +37,7 @@ data PieceDlState = Pending -- todo - map with index to a new data structure (peers who have that piece amd state) data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece - , state :: PieceDlState -- ^ state of the piece from download perspective. + , dlstate :: PieceDlState -- ^ state of the piece from download perspective. , hash :: ByteString -- ^ piece hash , len :: Integer } -- ^ piece length @@ -49,7 +50,7 @@ type PieceMap = Map Integer PieceData mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap mkPieceMap numPieces pieceHash pLengths = fromList kvs where kvs = [(i, PieceData { peers = [] - , state = Pending + , dlstate = Pending , hash = h , len = pLen }) | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths] @@ -57,7 +58,7 @@ mkPieceMap numPieces pieceHash pLengths = fromList kvs havePiece :: PieceMap -> Integer -> Bool havePiece pm index = - state (pm ! index) == Have + dlstate (pm ! index) == Have connectToPeer :: Peer -> IO Handle connectToPeer (Peer _ ip port) = do @@ -82,60 +83,69 @@ bitfieldToList bs = go bs 0 in setBits ++ go bs' (pos + 1) --- recvMsg :: Peer -> Handle -> Msg -msgLoop :: PeerState -> PieceMap -> FilePath -> IO () -msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file = - do - -- if me NOT Interested and she is Choking, tell her that - -- I am interested. - let h = handle pState - sendMsg h InterestedMsg - putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState) - msgLoop (pState { meInterested = True }) pieceStatus file -msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file = - -- if me Interested and she not Choking, send her a request - -- for a piece. - case pickPiece pieceStatus of - Nothing -> putStrLn "Nothing to download" - Just workPiece -> do - let pLen = len (pieceStatus ! workPiece) - putStrLn $ "piece length = " ++ show pLen - pBS <- downloadPiece (handle pState) workPiece pLen - if not $ verifyHash pBS (hash (pieceStatus ! workPiece)) - then - putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS)) - else do - let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1)) - putStrLn $ "Write into file at offset: " ++ show fileOffset - writeFileAtOffset file fileOffset pBS - msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file -msgLoop pState pieceStatus file = do - msg <- getMsg (handle pState) - putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState) - case msg of - KeepAliveMsg -> do - sendMsg (handle pState) KeepAliveMsg - putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState) - msgLoop pState pieceStatus file - BitFieldMsg bss -> do - let pieceList = bitfieldToList (unpack bss) - pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList - putStrLn $ show (length pieceList) ++ " Pieces" - -- for each pieceIndex in pieceList, make an entry in the pieceStatus - -- map with pieceIndex as the key and modify the value to add the peer. - -- download each of the piece in order - msgLoop pState pieceStatus' file - UnChokeMsg -> - msgLoop (pState { heChoking = False }) pieceStatus file - _ -> - msgLoop pState pieceStatus file +-- helper functions to manipulate PeerState +toPeerState :: Handle -> Peer -> Bool -> Bool -> Bool -> Bool -> PState +toPeerState h p meCh meIn heCh heIn = + PState { handle = h + , peer = p + , heChoking = heCh + , heInterested = heIn + , meChoking = meCh + , meInterested = meIn } + +-- -- recvMsg :: Peer -> Handle -> Msg +-- msgLoop :: PeerState -> PieceMap -> FilePath -> IO () +-- msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file = do +-- -- if me NOT Interested and she is Choking, tell her that +-- -- I am interested. +-- let h = handle pState +-- sendMsg h InterestedMsg +-- putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState) +-- msgLoop (pState { meInterested = True }) pieceStatus file +-- msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file = +-- -- if me Interested and she not Choking, send her a request +-- -- for a piece. +-- case pickPiece pieceStatus of +-- Nothing -> putStrLn "Nothing to download" +-- Just workPiece -> do +-- let pLen = len (pieceStatus ! workPiece) +-- putStrLn $ "piece length = " ++ show pLen +-- pBS <- downloadPiece (handle pState) workPiece pLen +-- if not $ verifyHash pBS (hash (pieceStatus ! workPiece)) +-- then +-- putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS)) +-- else do +-- let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1)) +-- putStrLn $ "Write into file at offset: " ++ show fileOffset +-- writeFileAtOffset file fileOffset pBS +-- msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file +-- msgLoop pState pieceStatus file = do +-- msg <- getMsg (handle pState) +-- putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState) +-- case msg of +-- KeepAliveMsg -> do +-- sendMsg (handle pState) KeepAliveMsg +-- putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState) +-- msgLoop pState pieceStatus file +-- BitFieldMsg bss -> do +-- let pieceList = bitfieldToList (unpack bss) +-- pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList +-- putStrLn $ show (length pieceList) ++ " Pieces" +-- -- for each pieceIndex in pieceList, make an entry in the pieceStatus +-- -- map with pieceIndex as the key and modify the value to add the peer. +-- -- download each of the piece in order +-- msgLoop pState pieceStatus' file +-- UnChokeMsg -> +-- msgLoop (pState { heChoking = False }) pieceStatus file +-- _ -> +-- msgLoop pState pieceStatus file -- simple algorithm to pick piece. -- pick the first piece from 0 that is not downloaded yet. pickPiece :: PieceMap -> Maybe Integer pickPiece m = let pieceList = toList m - allPending = filter (\(_, v) -> state v == Pending) pieceList + allPending = filter (\(_, v) -> dlstate v == Pending) pieceList in case allPending of [] -> Nothing @@ -151,12 +161,7 @@ handlePeerMsgs :: Peer -> Metainfo -> String -> IO () handlePeerMsgs p m peerId = do h <- connectToPeer p doHandshake h p (infoHash m) peerId - let state = PeerState { handle = h - , peer = p - , heInterested = False - , heChoking = True - , meInterested = False - , meChoking = True } + let pstate = toPeerState h p False True False True pieceHash = pieces (info m) numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash pLen = pieceLength (info m) @@ -164,8 +169,21 @@ handlePeerMsgs p m peerId = do fileName = name (info m) pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen) createDummyFile fileName (fromIntegral fileLen) - msgLoop state pieceStatus fileName - + (r, _) <- runStateT (msgLoop pieceStatus fileName) pstate + return () + +msgLoop :: PieceMap -> FilePath -> StateT PState IO () +msgLoop pieceStatus file = + StateT(\pState -> do + let h = handle pState + msg <- getMsg h + liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState) + case msg of + KeepAliveMsg -> do + sendMsg h KeepAliveMsg + liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState) + runStateT (msgLoop pieceStatus file) pState) + downloadPiece :: Handle -> Integer -> Integer -> IO ByteString downloadPiece h index pieceLength = do let chunks = splitNum pieceLength 16384