From: Ramakrishnan Muthukrishnan Date: Sun, 2 Aug 2015 14:47:19 +0000 (+0530) Subject: refactor Peer module X-Git-Url: https://git.rkrishnan.org/pf/content/en/seg/something?a=commitdiff_plain;h=83b8d4b201389fd36b851bad722ac84a0056abf6;p=functorrent.git refactor Peer module --- diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs index cdc0ea6..e809f3f 100644 --- a/src/FuncTorrent/Peer.hs +++ b/src/FuncTorrent/Peer.hs @@ -7,15 +7,10 @@ module FuncTorrent.Peer import Prelude hiding (lookup, concat, replicate, splitAt, take) import System.IO (Handle, BufferMode(..), hSetBuffering) -import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, take, empty) -import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict) -import qualified Data.ByteString.Char8 as BC (replicate, pack, length) +import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty) +import qualified Data.ByteString.Char8 as BC (length) import Network (connectTo, PortID(..)) -import Data.Binary (Binary(..), decode, encode) -import Data.Binary.Put (putWord32be, putWord16be, putWord8) -import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet) -import Control.Monad (replicateM, liftM, forM) -import Control.Applicative ((<$>), liftA3) +import Control.Monad (liftM, forM) import Data.Bits import Data.Word (Word8) import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust) @@ -24,10 +19,7 @@ import qualified Crypto.Hash.SHA1 as SHA1 (hash) import FuncTorrent.Metainfo (Info(..), Metainfo(..)) import FuncTorrent.Utils (splitN, splitNum) import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset) - -type ID = String -type IP = String -type Port = Integer +import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg) -- PeerState is a misnomer data PeerState = PeerState { handle :: Handle @@ -51,22 +43,6 @@ data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have -- which piece is with which peers type PieceMap = Map Integer PieceData --- | Peer is a PeerID, IP address, port tuple -data Peer = Peer ID IP Port - deriving (Show, Eq) - -data PeerMsg = KeepAliveMsg - | ChokeMsg - | UnChokeMsg - | InterestedMsg - | NotInterestedMsg - | HaveMsg Integer - | BitFieldMsg ByteString - | RequestMsg Integer Integer Integer - | PieceMsg Integer Integer ByteString - | CancelMsg Integer Integer Integer - | PortMsg Port - deriving (Show) -- Make the initial Piece map, with the assumption that no peer has the -- piece and that every piece is pending download. @@ -83,103 +59,21 @@ havePiece :: PieceMap -> Integer -> Bool havePiece pm index = state (pm ! index) == Have -genHandShakeMsg :: ByteString -> String -> ByteString -genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID] - where pstrlen = singleton 19 - pstr = BC.pack "BitTorrent protocol" - reserved = BC.replicate 8 '\0' - peerID = BC.pack peer_id - connectToPeer :: Peer -> IO Handle -connectToPeer peer@(Peer _ ip port) = do +connectToPeer (Peer _ ip port) = do h <- connectTo ip (PortNumber (fromIntegral port)) hSetBuffering h LineBuffering return h -doHandShake :: Handle -> Peer -> ByteString -> String -> IO () -doHandShake h peer infoHash peerid = do - let hs = genHandShakeMsg infoHash peerid +doHandshake :: Handle -> Peer -> ByteString -> String -> IO () +doHandshake h peer infoHash peerid = do + let hs = genHandshakeMsg infoHash peerid hPut h hs putStrLn $ "--> handhake to peer: " ++ show peer _ <- hGet h (length (unpack hs)) putStrLn $ "<-- handshake from peer: " ++ show peer return () -instance Binary PeerMsg where - put msg = case msg of - KeepAliveMsg -> putWord32be 0 - ChokeMsg -> do putWord32be 1 - putWord8 0 - UnChokeMsg -> do putWord32be 1 - putWord8 1 - InterestedMsg -> do putWord32be 1 - putWord8 2 - NotInterestedMsg -> do putWord32be 1 - putWord8 3 - HaveMsg i -> do putWord32be 5 - putWord8 4 - putWord32be (fromIntegral i) - BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen) - putWord8 5 - mapM_ putWord8 bfList - where bfList = unpack bf - bfListLen = length bfList - RequestMsg i o l -> do putWord32be 13 - putWord8 6 - putWord32be (fromIntegral i) - putWord32be (fromIntegral o) - putWord32be (fromIntegral l) - PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen) - putWord8 7 - putWord32be (fromIntegral i) - putWord32be (fromIntegral o) - mapM_ putWord8 blockList - where blockList = unpack b - blocklen = length blockList - CancelMsg i o l -> do putWord32be 13 - putWord8 8 - putWord32be (fromIntegral i) - putWord32be (fromIntegral o) - putWord32be (fromIntegral l) - PortMsg p -> do putWord32be 3 - putWord8 9 - putWord16be (fromIntegral p) - get = do - l <- getWord32be - msgid <- getWord8 - case msgid of - 0 -> return ChokeMsg - 1 -> return UnChokeMsg - 2 -> return InterestedMsg - 3 -> return NotInterestedMsg - 4 -> liftM (HaveMsg . fromIntegral) getWord32be - 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8) - 6 -> liftA3 RequestMsg getInteger getInteger getInteger - where getInteger = fromIntegral <$> getWord32be - 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8) - where getInteger = fromIntegral <$> getWord32be - 8 -> liftA3 CancelMsg getInteger getInteger getInteger - where getInteger = fromIntegral <$> getWord32be - 9 -> liftM (PortMsg . fromIntegral) getWord16be - _ -> error ("unknown message ID: " ++ show msgid) - -getMsg :: Handle -> IO PeerMsg -getMsg h = do - lBS <- hGet h 4 - let l = bsToInt lBS - if l == 0 - then return KeepAliveMsg - else do - msg <- hGet h l - return $ decode $ fromStrict $ concat [lBS, msg] - -sendMsg :: Handle -> PeerMsg -> IO () -sendMsg h msg = hPut h bsMsg - where bsMsg = toStrict $ encode msg - -bsToInt :: ByteString -> Int -bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x))) - bitfieldToList :: [Word8] -> [Integer] bitfieldToList bs = go bs 0 where go [] _ = [] @@ -190,7 +84,7 @@ bitfieldToList bs = go bs 0 -- recvMsg :: Peer -> Handle -> Msg msgLoop :: PeerState -> PieceMap -> IO () -msgLoop pState@(PeerState { meInterested = False , heChoking = True }) pieceStatus = +msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus = do -- if me NOT Interested and she is Choking, tell her that -- I am interested. @@ -256,7 +150,7 @@ updatePieceAvailability pieceStatus p pieceList = handlePeerMsgs :: Peer -> Metainfo -> String -> IO () handlePeerMsgs p m peerId = do h <- connectToPeer p - doHandShake h p (infoHash m) peerId + doHandshake h p (infoHash m) peerId let state = PeerState { handle = h , peer = p , heInterested = False diff --git a/src/FuncTorrent/PeerMsgs.hs b/src/FuncTorrent/PeerMsgs.hs new file mode 100644 index 0000000..f4fe97c --- /dev/null +++ b/src/FuncTorrent/PeerMsgs.hs @@ -0,0 +1,124 @@ +{-# LANGUAGE OverloadedStrings #-} +module FuncTorrent.PeerMsgs + (genHandshakeMsg, + sendMsg, + getMsg, + Peer(..), + PeerMsg(..) + ) where + +import Prelude hiding (lookup, concat, replicate, splitAt, take) + +import System.IO (Handle) +import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton) +import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict) +import qualified Data.ByteString.Char8 as BC (replicate, pack) +import Control.Monad (replicateM, liftM) +import Control.Applicative ((<$>), liftA3) + +import Data.Binary (Binary(..), decode, encode) +import Data.Binary.Put (putWord32be, putWord16be, putWord8) +import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet) + +-- | Peer is a PeerID, IP address, port tuple +data Peer = Peer ID IP Port + deriving (Show, Eq) + +type ID = String +type IP = String +type Port = Integer + +data PeerMsg = KeepAliveMsg + | ChokeMsg + | UnChokeMsg + | InterestedMsg + | NotInterestedMsg + | HaveMsg Integer + | BitFieldMsg ByteString + | RequestMsg Integer Integer Integer + | PieceMsg Integer Integer ByteString + | CancelMsg Integer Integer Integer + | PortMsg Port + deriving (Show) + +instance Binary PeerMsg where + put msg = case msg of + KeepAliveMsg -> putWord32be 0 + ChokeMsg -> do putWord32be 1 + putWord8 0 + UnChokeMsg -> do putWord32be 1 + putWord8 1 + InterestedMsg -> do putWord32be 1 + putWord8 2 + NotInterestedMsg -> do putWord32be 1 + putWord8 3 + HaveMsg i -> do putWord32be 5 + putWord8 4 + putWord32be (fromIntegral i) + BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen) + putWord8 5 + mapM_ putWord8 bfList + where bfList = unpack bf + bfListLen = length bfList + RequestMsg i o l -> do putWord32be 13 + putWord8 6 + putWord32be (fromIntegral i) + putWord32be (fromIntegral o) + putWord32be (fromIntegral l) + PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen) + putWord8 7 + putWord32be (fromIntegral i) + putWord32be (fromIntegral o) + mapM_ putWord8 blockList + where blockList = unpack b + blocklen = length blockList + CancelMsg i o l -> do putWord32be 13 + putWord8 8 + putWord32be (fromIntegral i) + putWord32be (fromIntegral o) + putWord32be (fromIntegral l) + PortMsg p -> do putWord32be 3 + putWord8 9 + putWord16be (fromIntegral p) + get = do + l <- getWord32be + msgid <- getWord8 + case msgid of + 0 -> return ChokeMsg + 1 -> return UnChokeMsg + 2 -> return InterestedMsg + 3 -> return NotInterestedMsg + 4 -> liftM (HaveMsg . fromIntegral) getWord32be + 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8) + 6 -> liftA3 RequestMsg getInteger getInteger getInteger + where getInteger = fromIntegral <$> getWord32be + 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8) + where getInteger = fromIntegral <$> getWord32be + 8 -> liftA3 CancelMsg getInteger getInteger getInteger + where getInteger = fromIntegral <$> getWord32be + 9 -> liftM (PortMsg . fromIntegral) getWord16be + _ -> error ("unknown message ID: " ++ show msgid) + +getMsg :: Handle -> IO PeerMsg +getMsg h = do + lBS <- hGet h 4 + let l = bsToInt lBS + if l == 0 + then return KeepAliveMsg + else do + msg <- hGet h l + return $ decode $ fromStrict $ concat [lBS, msg] + +sendMsg :: Handle -> PeerMsg -> IO () +sendMsg h msg = hPut h bsMsg + where bsMsg = toStrict $ encode msg + +genHandshakeMsg :: ByteString -> String -> ByteString +genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID] + where pstrlen = singleton 19 + pstr = BC.pack "BitTorrent protocol" + reserved = BC.replicate 8 '\0' + peerID = BC.pack peer_id + +bsToInt :: ByteString -> Int +bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))