diff --git a/Combinatorrent.cabal b/Combinatorrent.cabal --- a/Combinatorrent.cabal +++ b/Combinatorrent.cabal @@ -1,6 +1,6 @@ name: Combinatorrent category: Network -version: 0.2.2 +version: 0.3.0 category: Network description: Combinatorrent provides a BitTorrent client, based on STM for concurrency. This is an early preview release which is capable of @@ -18,6 +18,7 @@ maintainer: jesper.louis.andersen@gmail.com stability: experimental synopsis: A concurrent bittorrent client +tested-with: GHC ==6.12.1, GHC ==6.12.2, GHC ==6.13.20100426 build-type: Configure extra-tmp-files: src/Version.hs @@ -51,8 +52,9 @@ build-depends: array >= 0.3, + attoparsec, base >= 3.0, - base <= 5.0, + base < 5.0, bytestring, cereal, containers, @@ -65,6 +67,7 @@ HUnit, mtl, network, + network-bytestring, parsec, pretty, PSQueue, @@ -77,8 +80,10 @@ test-framework-quickcheck2, time - ghc-prof-options: -auto-all -caf-all ghc-options: -Wall -fno-warn-orphans -threaded + if impl(ghc >= 6.13.0) + ghc-options: -rtsopts + if !flag(debug) cpp-options: "-DNDEBUG" diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -34,12 +34,30 @@ [BEP Process](http://www.bittorrent.org/beps/bep_0000.html) document for an explanation of these) - - 003, 004, 020, + - 0003, 0004, 0020, +Combinatorrent implicitly supports these extensions + + - 0027: Support by the virtue of only supporting a single tracker and no + DHT. + +Partially supported extensions: + + - 0007: Combinatorrent understands and uses the "peers6" response from + the tracker to connect clients. On the other hand, it does nothing to + provide the "ipv4=" and "ipv6=" keys on tracker requests. As such, it + can be claimed that 0007 support is available, as everything we left + out is only qualified as MAY. + + - 0023: Combinatorrent supports the "compact" response only, although it + is explicitly stated that the client must support both. In practice it + has little impact as all modern trackers will only return compact + responses anyway. + Combinatorrent is not supporting these BEPs, but strives to do so one day: - - 005, 006, 007, 009, 010, 012, 015, 016, 017, 018, 019, 023, 021, 022, - 024, 026, 027, 028, 029, 030, 031, 032 + - 0005, 0006, 0009, 0010, 0012, 0015, 0016, 0017, 0018, 0019, 0021, 0022, + 0024, 0026, 0028, 0029, 0030, 0031, 0032 Debugging --------- diff --git a/configure b/configure --- a/configure +++ b/configure @@ -3,7 +3,7 @@ # Simple configuration script. ## Git version -GITHEAD="0.2.2 Hackage" +GITHEAD=$(awk '/^version:/ {print $2}' Combinatorrent.cabal) if test -d "${GIT_DIR:-.git}" ; then GITHEAD=`git describe 2>/dev/null` diff --git a/src/Channels.hs b/src/Channels.hs --- a/src/Channels.hs +++ b/src/Channels.hs @@ -15,11 +15,11 @@ import Control.Concurrent.STM import Control.DeepSeq -import Network +import Network.Socket + import Torrent -data Peer = Peer { peerHost :: HostName, - peerPort :: PortID } +data Peer = Peer SockAddr data PeerMessage = ChokePeer | UnchokePeer diff --git a/src/Data/PendingSet.hs b/src/Data/PendingSet.hs --- a/src/Data/PendingSet.hs +++ b/src/Data/PendingSet.hs @@ -51,17 +51,26 @@ -- Each piece is discriminated by a selector function until the first hit is -- found. Then all Pieces of the same priority accepted by the selector is -- chosen for return. -pick :: (PieceNum -> Bool) -> PendingSet -> Maybe [PieceNum] +pick :: (PieceNum -> IO Bool) -> PendingSet -> IO (Maybe [PieceNum]) pick selector ps = findPri (minView . unPS $ ps) - where findPri Nothing = Nothing - findPri (Just (pn :-> p, rest)) = - if selector pn - then pickAtPri p [pn] (minView rest) + where findPri Nothing = return Nothing + findPri (Just (pn :-> p, rest)) = do + r <- selector pn + if r + then pickAtPri numToPick p [pn] (minView rest) else findPri $ minView rest - pickAtPri _p acc Nothing = Just acc - pickAtPri p acc (Just (pn :-> p', rest)) - | p == p' = if selector pn - then pickAtPri p (pn : acc) $ minView rest - else pickAtPri p acc $ minView rest - | otherwise = Just acc + pickAtPri 0 _p acc _ = return $ Just acc + pickAtPri _ _p acc Nothing = return $ Just acc + pickAtPri k p acc (Just (pn :-> p', rest)) + | p == p' = do + r <- selector pn + if r + then pickAtPri (k-1) p (pn : acc) $ minView rest + else pickAtPri k p acc $ minView rest + | otherwise = return $ Just acc +-- | Number of pieces to pick with the picker. Setting an upper limit here because if a lot +-- of peers have all pieces, these numbers grow insanely big, leading to allocation we +-- don't really need. +numToPick :: Int +numToPick = 7 diff --git a/src/Data/Queue.hs b/src/Data/Queue.hs --- a/src/Data/Queue.hs +++ b/src/Data/Queue.hs @@ -9,6 +9,7 @@ , remove , push , pop + , Data.Queue.partition , Data.Queue.filter -- * Test Suite , testSuite @@ -68,7 +69,14 @@ filter :: (a -> Bool) -> Queue a -> Queue a filter p (Queue front back) = Queue (Lst.filter p front) (Lst.filter p back) +-- | Find elements matching an predicate and lift them out of the queue +partition :: (a -> Bool) -> Queue a -> ([a], Queue a) +partition p (Queue front back) = + let (ft, fl) = Lst.partition p front + (bt, bl) = Lst.partition p back + in (ft ++ bt, Queue fl bl) +--------------------------------------------------------------------- -- Tests testSuite :: Test diff --git a/src/Digest.hs b/src/Digest.hs --- a/src/Digest.hs +++ b/src/Digest.hs @@ -1,11 +1,14 @@ -- | Simple abstraction for message digests +{-# LANGUAGE TypeSynonymInstances #-} module Digest ( Digest , digest + , digestBS ) where import Control.Applicative +import Control.DeepSeq import Control.Monad.State import Data.Word @@ -19,20 +22,24 @@ -- Consider newtyping this type Digest = B.ByteString +instance NFData Digest + digest :: L.ByteString -> IO B.ByteString digest bs = {-# SCC "sha1_digest" #-} B.pack <$> digestLBS SSL.SHA1 bs +digestBS :: B.ByteString -> IO B.ByteString +digestBS bs = digest . L.fromChunks $ [bs] + digestLBS :: SSL.MessageDigest -> L.ByteString -> IO [Word8] -digestLBS mdType xs = {-# SCC "sha1_digestLBS" #-} - SSL.mkDigest mdType $ evalStateT (updateLBS xs >> SSL.final) +digestLBS mdType xs = SSL.mkDigest mdType $ evalStateT (updateLBS xs >> SSL.final) updateBS :: B.ByteString -> SSL.Digest () -updateBS bs = {-# SCC "sha1_updateBS" #-} do +updateBS bs = do SSL.DST ctx <- get _ <- liftIO $ unsafeUseAsCStringLen bs $ \(ptr, len) -> SSL.digestUpdate ctx (castPtr ptr) (fromIntegral len) return () updateLBS :: L.ByteString -> SSL.Digest () -updateLBS lbs = {-# SCC "sha1_updateLBS" #-} mapM_ updateBS chunked +updateLBS lbs = mapM_ updateBS chunked where chunked = {-# SCC "sha1_updateLBS_chunked" #-} L.toChunks lbs diff --git a/src/FS.hs b/src/FS.hs --- a/src/FS.hs +++ b/src/FS.hs @@ -56,19 +56,20 @@ ) $ r -} -projectHandles (Handles handles@((h1, length1):handles')) offset size +projectHandles (Handles handles@((h1, length1):handles')) offs size | size <= 0 = fail "FS: Should have already stopped projection" | null handles = fail "FS: Attempt to read beyond torrent length" - | offset >= length1 = - projectHandles (Handles handles') (offset - length1) size + | offs >= length1 = + projectHandles (Handles handles') (offs - length1) size | otherwise = - let size1 = length1 - offset -- ^How much of h1 to take? + let size1 = length1 - offs -- ^How much of h1 to take? in if size1 >= size - then [(h1, offset, size)] - else (h1, offset, size1) : + then [(h1, offs, size)] + else (h1, offs, size1) : projectHandles (Handles handles') 0 (size - size1) +projectHandles (Handles []) _ _ = fail "FS: Empty Handles list, can't happen" pInfoLookup :: PieceNum -> PieceMap -> IO PieceInfo pInfoLookup pn mp = return $ mp ! pn @@ -80,8 +81,8 @@ do pInfo <- pInfoLookup pn mp bs <- L.concat `fmap` forM (projectHandles handles (offset pInfo) (len pInfo)) - (\(h, offset, size) -> - do hSeek h AbsoluteSeek offset + (\(h, offs, size) -> + do hSeek h AbsoluteSeek offs L.hGet h (fromInteger size) ) if L.length bs == (fromInteger . len $ pInfo) @@ -97,8 +98,8 @@ B.concat `fmap` forM (projectHandles handles (offset pInfo + (fromIntegral $ blockOffset blk)) (fromIntegral $ blockSize blk)) - (\(h, offset, size) -> - do hSeek h AbsoluteSeek offset + (\(h, offs, size) -> + do hSeek h AbsoluteSeek offs B.hGet h $ fromInteger size ) @@ -110,12 +111,12 @@ {-# SCC "writeBlock" #-} do when lenFail $ fail "Writing block of wrong length" pInfo <- pInfoLookup n pm - foldM_ (\blkData (h, offset, size) -> + foldM_ (\blkData' (h, offs, size) -> do let size' = fromInteger size - hSeek h AbsoluteSeek offset - B.hPut h $ B.take size' blkData + hSeek h AbsoluteSeek offs + B.hPut h $ B.take size' blkData' hFlush h - return $ B.drop size' blkData + return $ B.drop size' blkData' ) blkData (projectHandles handles (position pInfo) (fromIntegral $ B.length blkData)) return () where @@ -129,8 +130,8 @@ checkPiece inf handles = {-# SCC "checkPiece" #-} do bs <- L.concat `fmap` forM (projectHandles handles (offset inf) (fromInteger $ len inf)) - (\(h, offset, size) -> - do hSeek h AbsoluteSeek offset + (\(h, offs, size) -> + do hSeek h AbsoluteSeek offs L.hGet h (fromInteger size) ) dgs <- liftIO $ D.digest bs @@ -185,13 +186,13 @@ do handles <- Handles `fmap` forM files - (\(path, length) -> + (\(path, l) -> do let dir = joinPath $ init path when (dir /= "") $ createDirectoryIfMissing True dir let fpath = joinPath path h <- openBinaryFile fpath ReadWriteMode - return (h, length) + return (h, l) ) have <- checkFile handles pieceMap return (handles, have, pieceMap) diff --git a/src/Process/ChokeMgr.hs b/src/Process/ChokeMgr.hs --- a/src/Process/ChokeMgr.hs +++ b/src/Process/ChokeMgr.hs @@ -3,6 +3,7 @@ -- * Types, Channels ChokeMgrChannel , RateTVar + , PeerRateInfo(..) , ChokeMgrMsg(..) -- * Interface , start @@ -47,8 +48,16 @@ type ChokeMgrChannel = TChan ChokeMgrMsg -type RateTVar = TVar [(ThreadId, (Double, Double, Bool, Bool, Bool))] +data PeerRateInfo = PRI { + peerUpRate :: Double, + peerDownRate :: Double, + peerInterested :: Bool, + peerSeeding :: Bool, + peerChokingUs :: Bool } + deriving Show +type RateTVar = TVar [(ThreadId, PeerRateInfo)] + data CF = CF { mgrCh :: ChokeMgrChannel , rateTV :: RateTVar } @@ -69,7 +78,7 @@ start ch rtv ur supC = do _ <- registerSTM roundTickSecs ch Tick spawnP (CF ch rtv) (initPeerDB $ calcUploadSlots ur Nothing) - (catchP pgm + ({-# SCC "ChokeMgr" #-} catchP pgm (defaultStopHandler supC)) where initPeerDB slots = PeerDB 2 slots S.empty M.empty [] @@ -111,14 +120,19 @@ , pChannel :: PeerChannel } +instance Show Peer where + show (Peer tid _ _) = "(Peer " ++ show tid ++ "...)" + -- | Peer upload and download ratio data PRate = PRate { pUpRate :: Double, pDownRate :: Double } + deriving Show -- | Current State of the peer data PState = PState { pChokingUs :: Bool -- ^ True if the peer is choking us , pInterestedInUs :: Bool -- ^ Reflection from Peer DB , pIsASeeder :: Bool -- ^ True if the peer is a seeder } + deriving Show type RateMap = M.Map ThreadId (PRate, PState) data PeerDB = PeerDB @@ -140,11 +154,12 @@ return q case rateUpdate of [] -> return () - updates -> let f old (tid, (uprt, downrt, interested, seeder, choking)) = - M.insert tid (PRate { pUpRate = uprt, pDownRate = downrt }, - PState { pInterestedInUs = interested, - pIsASeeder = seeder, - pChokingUs = choking }) old + updates -> let f old (tid, pri) = + M.insert tid (PRate { pUpRate = peerUpRate pri, + pDownRate = peerDownRate pri }, + PState { pInterestedInUs = peerInterested pri, + pIsASeeder = peerSeeding pri, + pChokingUs = peerChokingUs pri }) old nm m = foldl f m $ reverse updates in do debugP $ "Rate updates since last round: " ++ show updates @@ -191,7 +206,10 @@ sd <- gets seeding rm <- gets rateMap debugP $ "Chain is: " ++ show (map pThreadId chn) + debugP $ "RateMap is: " ++ show rm let (seed, down) = splitSeedLeech sd rm chn + debugP $ "Seeders " ++ show seed + debugP $ "Downloaders " ++ show down electedPeers <- selectPeers us down seed performChokingUnchoking electedPeers chn @@ -313,8 +331,12 @@ where -- Partition the peers in elected and non-elected (electedPeers, nonElectedPeers) = partition (\rd -> S.member (pThreadId rd) elected) peers - unchoke p = msgPeer (pChannel p) UnchokePeer - choke p = msgPeer (pChannel p) ChokePeer + unchoke p = do + debugP $ "Unchoking: " ++ show p + msgPeer (pChannel p) UnchokePeer + choke p = do + debugP $ "Choking: " ++ show p + msgPeer (pChannel p) ChokePeer -- If we have k optimistic slots, @optChoke k peers@ will unchoke the first -- @k@ peers interested in us. The rest will either be unchoked if they are diff --git a/src/Process/Console.hs b/src/Process/Console.hs --- a/src/Process/Console.hs +++ b/src/Process/Console.hs @@ -39,7 +39,8 @@ start waitC statusC supC = do cmdC <- readerP -- We shouldn't be doing this in the long run wrtC <- writerP - spawnP (CF cmdC wrtC) () (catchP eventLoop (defaultStopHandler supC)) + spawnP (CF cmdC wrtC) () ({-# SCC "Console" #-} + catchP eventLoop (defaultStopHandler supC)) where eventLoop = do c <- asks cmdCh @@ -69,14 +70,15 @@ writerP = do wrt <- newTChanIO _ <- forkIO $ lp wrt return wrt - where lp wCh = forever (do m <- atomically $ readTChan wCh - putStrLn m) + where lp wCh = {-# SCC "writerP" #-} + forever (do m <- atomically $ readTChan wCh + putStrLn m) readerP :: IO CmdChannel readerP = do cmd <- newTChanIO _ <- forkIO $ lp cmd return cmd - where lp cmd = forever $ + where lp cmd = {-# SCC "readerP" #-} forever $ do l <- getLine atomically $ writeTChan cmd (case l of diff --git a/src/Process/DirWatcher.hs b/src/Process/DirWatcher.hs --- a/src/Process/DirWatcher.hs +++ b/src/Process/DirWatcher.hs @@ -37,7 +37,8 @@ -> IO ThreadId start fp chan supC = do spawnP (CF chan fp) S.empty - (catchP pgm (defaultStopHandler supC)) + ({-# SCC "DirWatcher" #-} + catchP pgm (defaultStopHandler supC)) where pgm = do q <- liftIO $ registerDelay (5 * 1000000) liftIO . atomically $ do diff --git a/src/Process/FS.hs b/src/Process/FS.hs --- a/src/Process/FS.hs +++ b/src/Process/FS.hs @@ -47,7 +47,8 @@ start :: FS.Handles -> FS.PieceMap -> FSPChannel -> SupervisorChannel -> IO ThreadId start handles pm fspC supC = - spawnP (CF fspC) (ST handles pm) (catchP lp (defaultStopHandler supC)) + spawnP (CF fspC) (ST handles pm) ({-# SCC "FS" #-} + catchP lp (defaultStopHandler supC)) where lp = do c <- asks fspCh diff --git a/src/Process/Listen.hs b/src/Process/Listen.hs --- a/src/Process/Listen.hs +++ b/src/Process/Listen.hs @@ -5,10 +5,10 @@ import Control.Concurrent import Control.Concurrent.STM -import Control.Monad import Control.Monad.Reader import Network +import Network.Socket as S import Process import Process.PeerMgr hiding (start) @@ -19,14 +19,23 @@ instance Logging CF where logName _ = "Process.Listen" -start :: PortID -> PeerMgrChannel -> SupervisorChannel -> IO ThreadId +start :: PortNumber -> PeerMgrChannel -> SupervisorChannel -> IO ThreadId start port peerMgrC supC = do - spawnP (CF peerMgrC) () (catchP (openListen >>= pgm) + spawnP (CF peerMgrC) () ({-# SCC "Listen" #-} catchP (openListen port >>= eventLoop) (defaultStopHandler supC)) -- TODO: Close socket resource! - where openListen = liftIO $ listenOn port - pgm sock = forever lp - where lp = do c <- asks peerMgrCh - liftIO $ do - conn <- accept sock - atomically $ writeTChan c (NewIncoming conn) + +openListen :: PortNumber -> Process CF () Socket +openListen port = liftIO $ do + s <- S.socket S.AF_INET S.Stream S.defaultProtocol + S.bindSocket s (S.SockAddrInet port S.iNADDR_ANY) + S.listen s 10 + return s + +eventLoop :: Socket -> Process CF () () +eventLoop sockFd = do + c <- asks peerMgrCh + liftIO $ do + conn <- S.accept sockFd + atomically $ writeTChan c (NewIncoming conn) + eventLoop sockFd diff --git a/src/Process/Peer.hs b/src/Process/Peer.hs --- a/src/Process/Peer.hs +++ b/src/Process/Peer.hs @@ -5,6 +5,8 @@ PeerMessage(..) -- * Interface , Process.Peer.start + -- * Tests + , Process.Peer.testSuite ) where @@ -23,24 +25,30 @@ import Data.Array import Data.Bits import qualified Data.ByteString as B -import qualified Data.ByteString.Lazy as L +import Data.Function (on) import qualified Data.PieceSet as PS import Data.Maybe +import Data.Monoid(Monoid(..), Last(..)) import Data.Set as S hiding (map) import Data.Time.Clock import Data.Word -import System.IO +import Network.Socket hiding (KeepAlive) +import Test.Framework +import Test.Framework.Providers.HUnit +import Test.HUnit hiding (Path, Test, assert) + import Channels +import Digest import Process import Process.FS import Process.PieceMgr import RateCalc as RC import Process.Status -import Process.ChokeMgr (RateTVar) +import qualified Process.ChokeMgr as ChokeMgr (RateTVar, PeerRateInfo(..)) import Process.Timer import Supervisor import Torrent @@ -53,18 +61,18 @@ -- INTERFACE ---------------------------------------------------------------------- -start :: Handle -> MgrChannel -> RateTVar -> PieceMgrChannel +start :: Socket -> [Capabilities] -> MgrChannel -> ChokeMgr.RateTVar -> PieceMgrChannel -> FSPChannel -> TVar [PStat] -> PieceMap -> Int -> InfoHash -> IO Children -start h pMgrC rtv pieceMgrC fsC stv pm nPieces ih = do +start s caps pMgrC rtv pieceMgrC fsC stv pm nPieces ih = do queueC <- newTChanIO senderMV <- newEmptyTMVarIO receiverC <- newTChanIO sendBWC <- newTChanIO - return [Worker $ Sender.start h senderMV, - Worker $ SenderQ.start queueC senderMV sendBWC fsC, - Worker $ Receiver.start h receiverC, - Worker $ peerP pMgrC rtv pieceMgrC pm nPieces + return [Worker $ Sender.start s senderMV, + Worker $ SenderQ.start caps queueC senderMV sendBWC fsC, + Worker $ Receiver.start s receiverC, + Worker $ peerP caps pMgrC rtv pieceMgrC pm nPieces queueC receiverC sendBWC stv ih] -- INTERNAL FUNCTIONS @@ -78,12 +86,13 @@ , sendBWCh :: BandwidthChannel , timerCh :: TChan () , statTV :: TVar [PStat] - , rateTV :: RateTVar + , rateTV :: ChokeMgr.RateTVar , pcInfoHash :: InfoHash , pieceMap :: !PieceMap , piecesDoneTV :: TMVar [PieceNum] , interestTV :: TMVar Bool , grabBlockTV :: TMVar Blocks + , extConf :: ExtensionConfig } instance Logging CF where @@ -102,12 +111,126 @@ , lastMessage :: !Int } +data ExtensionConfig = ExtensionConfig + { handleHaveAll :: Last (Process CF ST ()) + , handleHaveNone :: Last (Process CF ST ()) + , handleBitfield :: Last (BitField -> Process CF ST ()) + , handleSuggest :: Last (PieceNum -> Process CF ST ()) + , handleAllowedFast :: Last (PieceNum -> Process CF ST ()) + , handleRejectRequest :: Last (PieceNum -> Block -> Process CF ST ()) + , handleExtendedMsg :: Last (Int -> B.ByteString -> Process CF ST ()) + , handleRequestMsg :: Last (PieceNum -> Block -> Process CF ST ()) + , handleChokeMsg :: Last (Process CF ST ()) + , handleCancelMsg :: Last (PieceNum -> Block -> Process CF ST ()) + , handlePieceMsg :: Last (PieceNum -> Int -> B.ByteString -> Process CF ST ()) + } -peerP :: MgrChannel -> RateTVar -> PieceMgrChannel -> PieceMap -> Int +emptyExtensionConfig :: ExtensionConfig +emptyExtensionConfig = ExtensionConfig + mempty mempty mempty mempty mempty mempty mempty mempty mempty mempty mempty + +appendEConfig :: ExtensionConfig -> ExtensionConfig -> ExtensionConfig +appendEConfig a b = + ExtensionConfig { + handleHaveAll = app handleHaveAll a b + , handleHaveNone = app handleHaveNone a b + , handleBitfield = app handleBitfield a b + , handleSuggest = app handleSuggest a b + , handleAllowedFast = app handleAllowedFast a b + , handleRejectRequest = app handleRejectRequest a b + , handleExtendedMsg = app handleExtendedMsg a b + , handleRequestMsg = app handleRequestMsg a b + , handleChokeMsg = app handleChokeMsg a b + , handleCancelMsg = app handleCancelMsg a b + , handlePieceMsg = app handlePieceMsg a b + } + where app f = mappend `on` f + +instance Monoid ExtensionConfig where + mempty = emptyExtensionConfig + mappend = appendEConfig + +-- | Constructor for 'Last' values. +ljust :: a -> Last a +ljust = Last . Just + +-- | Deconstructor for 'Last' values. +fromLJ :: (ExtensionConfig -> Last a) -- ^ Field to access. + -> ExtensionConfig -- ^ Default to use. + -> a +fromLJ f cfg = case f cfg of + Last Nothing -> fromLJ f extensionBase + Last (Just a) -> a + +extensionBase :: ExtensionConfig +extensionBase = ExtensionConfig + (ljust errorHaveAll) + (ljust errorHaveNone) + (ljust bitfieldMsg) + (ljust errorSuggest) + (ljust errorAllowedFast) + (ljust errorRejectRequest) + (ljust errorExtendedMsg) + (ljust requestMsg) + (ljust chokeMsg) + (ljust cancelBlock) + (ljust pieceMsg) + +fastExtension :: ExtensionConfig +fastExtension = ExtensionConfig + (ljust haveAllMsg) + (ljust haveNoneMsg) + (ljust bitfieldMsg) + (ljust ignoreSuggest) + (ljust ignoreAllowedFast) + (ljust rejectMsg) + mempty + (ljust requestFastMsg) + (ljust fastChokeMsg) + (ljust fastCancelBlock) + (ljust fastPieceMsg) + +ignoreSuggest :: PieceNum -> Process CF ST () +ignoreSuggest _ = debugP "Ignoring SUGGEST message" + +ignoreAllowedFast :: PieceNum -> Process CF ST () +ignoreAllowedFast _ = debugP "Ignoring ALLOWEDFAST message" + +errorHaveAll :: Process CF ST () +errorHaveAll = do + errorP "Received a HAVEALL, but the extension is not enabled" + stopP + +errorHaveNone :: Process CF ST () +errorHaveNone = do + errorP "Received a HAVENONE, but the extension is not enabled" + stopP + +errorSuggest :: PieceNum -> Process CF ST () +errorSuggest _ = do + errorP "Received a SUGGEST PIECE, but the extension is not enabled" + stopP + +errorAllowedFast :: PieceNum -> Process CF ST () +errorAllowedFast _ = do + errorP "Received a ALLOWEDFAST, but the extension is not enabled" + stopP + +errorRejectRequest :: PieceNum -> Block -> Process CF ST () +errorRejectRequest _ _ = do + errorP "Received a REJECT REQUEST, but the extension is not enabled" + stopP + +errorExtendedMsg :: Int -> B.ByteString -> Process CF ST () +errorExtendedMsg _ _ = do + errorP "Received an EXTENDED MESSAGE, but the extension is not enabled" + stopP + +peerP :: [Capabilities] -> MgrChannel -> ChokeMgr.RateTVar -> PieceMgrChannel -> PieceMap -> Int -> TChan SenderQ.SenderQMsg -> TChan (Message, Integer) -> BandwidthChannel -> TVar [PStat] -> InfoHash -> SupervisorChannel -> IO ThreadId -peerP pMgrC rtv pieceMgrC pm nPieces outBound inBound sendBWC stv ih supC = do +peerP caps pMgrC rtv pieceMgrC pm nPieces outBound inBound sendBWC stv ih supC = do ch <- newTChanIO tch <- newTChanIO ct <- getCurrentTime @@ -115,11 +238,17 @@ intmv <- newEmptyTMVarIO gbtmv <- newEmptyTMVarIO pieceSet <- PS.new nPieces + let cs = configCapabilities caps spawnP (CF inBound outBound pMgrC pieceMgrC ch sendBWC tch stv rtv ih pm - pdtmv intmv gbtmv) + pdtmv intmv gbtmv cs) (ST True False S.empty True False pieceSet nPieces (RC.new ct) (RC.new ct) False 0) - (cleanupP (startup nPieces) (defaultStopHandler supC) cleanup) + ({-# SCC "PeerControl" #-} + cleanupP (startup nPieces) (defaultStopHandler supC) cleanup) +configCapabilities :: [Capabilities] -> ExtensionConfig +configCapabilities caps = + mconcat [mempty, if Fast `elem` caps then fastExtension else mempty] + startup :: Int -> Process CF ST () startup nPieces = do tid <- liftIO $ myThreadId @@ -138,9 +267,8 @@ cleanup = do t <- liftIO myThreadId pieces <- gets peerPieces >>= PS.toList - ch <- asks pieceMgrCh ch2 <- asks peerMgrCh - liftIO . atomically $ writeTChan ch (PeerUnhave pieces) + msgPieceMgr (PeerUnhave pieces) liftIO . atomically $ writeTChan ch2 (Disconnect t) readOp :: Process CF ST Operation @@ -160,8 +288,10 @@ op <- readOp case op of PeerMsgEvt (m, sz) -> peerMsg m sz - ChokeMgrEvt m -> chokeMsg m - UpRateEvent up -> modify (\s -> s { upRate = RC.update up $ upRate s}) + ChokeMgrEvt m -> chokeMgrMsg m + UpRateEvent up -> do s <- get + u <- return $ RC.update up $ upRate s + put $! s { upRate = u } TimerEvent -> timerTick eventLoop @@ -174,34 +304,48 @@ -- | Return a list of pieces which are currently done by us getPiecesDone :: Process CF ST [PieceNum] getPiecesDone = do - ch <- asks pieceMgrCh c <- asks piecesDoneTV - liftIO $ do - atomically $ writeTChan ch (GetDone c) - atomically $ takeTMVar c + msgPieceMgr (GetDone c) + liftIO $ do atomically $ takeTMVar c -- | Process an event from the Choke Manager -chokeMsg :: PeerMessage -> Process CF ST () -chokeMsg msg = do - debugP "ChokeMgrEvent" +chokeMgrMsg :: PeerMessage -> Process CF ST () +chokeMgrMsg msg = do case msg of - PieceCompleted pn -> outChan $ SenderQ.SenderQM $ Have pn + PieceCompleted pn -> do + debugP "Telling about Piece Completion" + outChan $ SenderQ.SenderQM $ Have pn + considerInterest ChokePeer -> do choking <- gets weChoke when (not choking) - (do outChan $ SenderQ.SenderOChoke - debugP "ChokePeer" + (do t <- liftIO myThreadId + debugP $ "Pid " ++ show t ++ " choking" + outChan $ SenderQ.SenderOChoke modify (\s -> s {weChoke = True})) UnchokePeer -> do choking <- gets weChoke when choking - (do outChan $ SenderQ.SenderQM Unchoke - debugP "UnchokePeer" + (do t <- liftIO myThreadId + debugP $ "Pid " ++ show t ++ " unchoking" + outChan $ SenderQ.SenderQM Unchoke modify (\s -> s {weChoke = False})) CancelBlock pn blk -> do - modify (\s -> s { blockQueue = S.delete (pn, blk) $ blockQueue s }) - outChan $ SenderQ.SenderQRequestPrune pn blk + cf <- asks extConf + fromLJ handleCancelMsg cf pn blk -processLastMessage :: Process CF ST () -processLastMessage = do +cancelBlock :: PieceNum -> Block -> Process CF ST () +cancelBlock pn blk = do + s <- get + put $! s { blockQueue = S.delete (pn, blk) $ blockQueue s } + outChan $ SenderQ.SenderQRequestPrune pn blk + +fastCancelBlock :: PieceNum -> Block -> Process CF ST () +fastCancelBlock pn blk = do + bq <- gets blockQueue + when (S.member (pn, blk) bq) + $ outChan $ SenderQ.SenderQRequestPrune pn blk + +checkKeepAlive :: Process CF ST () +checkKeepAlive = do lm <- gets lastMessage if lm >= 24 then do outChan $ SenderQ.SenderQM KeepAlive @@ -215,8 +359,7 @@ timerTick :: Process CF ST () timerTick = do mTid <- liftIO myThreadId - processLastMessage - debugP "TimerEvent" + checkKeepAlive tch <- asks timerCh _ <- registerSTM 5 tch () -- Tell the ChokeMgr about our progress @@ -232,11 +375,15 @@ rtv <- asks rateTV liftIO . atomically $ do q <- readTVar rtv - writeTVar rtv ((mTid, (up, down, i, seed, pchoke)) : q) + writeTVar rtv ((mTid, ChokeMgr.PRI { + ChokeMgr.peerUpRate = up, + ChokeMgr.peerDownRate = down, + ChokeMgr.peerInterested = i, + ChokeMgr.peerSeeding = seed, + ChokeMgr.peerChokingUs = pchoke }) : q) -- Tell the Status Process about our progress let (upCnt, nuRate) = RC.extractCount $ nur (downCnt, ndRate) = RC.extractCount $ ndr - debugP $ "Sending peerStats: " ++ show upCnt ++ ", " ++ show downCnt stv <- asks statTV ih <- asks pcInfoHash liftIO .atomically $ do @@ -247,33 +394,59 @@ modify (\s -> s { upRate = nuRate, downRate = ndRate }) +chokeMsg :: Process CF ST () +chokeMsg = do + putbackBlocks + s <- get + put $! s { peerChoke = True } + +fastChokeMsg :: Process CF ST () +fastChokeMsg = do + s <- get + put $! s { peerChoke = True} + +unchokeMsg :: Process CF ST () +unchokeMsg = do + s <- get + put $! s { peerChoke = False } + fillBlocks + -- | Process an Message from the peer in the other end of the socket. peerMsg :: Message -> Integer -> Process CF ST () peerMsg msg sz = do modify (\s -> s { downRate = RC.update sz $ downRate s}) case msg of KeepAlive -> return () - Choke -> do putbackBlocks - modify (\s -> s { peerChoke = True }) - Unchoke -> do modify (\s -> s { peerChoke = False }) - fillBlocks + Choke -> asks extConf >>= fromLJ handleChokeMsg + Unchoke -> unchokeMsg Interested -> modify (\s -> s { peerInterested = True }) NotInterested -> modify (\s -> s { peerInterested = False }) Have pn -> haveMsg pn BitField bf -> bitfieldMsg bf - Request pn blk -> requestMsg pn blk - Piece n os bs -> do pieceMsg n os bs + Request pn blk -> do cf <- asks extConf + fromLJ handleRequestMsg cf pn blk + Piece n os bs -> do cf <- asks extConf + fromLJ handlePieceMsg cf n os bs fillBlocks Cancel pn blk -> cancelMsg pn blk Port _ -> return () -- No DHT yet, silently ignore + HaveAll -> fromLJ handleHaveAll =<< asks extConf + HaveNone -> fromLJ handleHaveNone =<< asks extConf + Suggest pn -> do cf <- asks extConf + fromLJ handleSuggest cf pn + AllowedFast pn -> do cf <- asks extConf + fromLJ handleAllowedFast cf pn + RejectRequest pn blk -> do cf <- asks extConf + fromLJ handleRejectRequest cf pn blk + ExtendedMsg idx bs -> do cf <- asks extConf + fromLJ handleExtendedMsg cf idx bs -- | Put back blocks for other peer processes to grab. This is done whenever -- the peer chokes us, or if we die by an unknown cause. putbackBlocks :: Process CF ST () putbackBlocks = do blks <- gets blockQueue - pmch <- asks pieceMgrCh - liftIO . atomically $ writeTChan pmch (PutbackBlocks (S.toList blks)) + msgPieceMgr (PutbackBlocks (S.toList blks)) modify (\s -> s { blockQueue = S.empty }) -- | Process a HAVE message from the peer. Note we also update interest as a side effect @@ -283,17 +456,18 @@ let (lo, hi) = bounds pm if pn >= lo && pn <= hi then do PS.insert pn =<< gets peerPieces - pmch <- asks pieceMgrCh - liftIO . atomically $ writeTChan pmch (PeerHave [pn]) + debugP $ "Peer has pn: " ++ show pn + msgPieceMgr (PeerHave [pn]) decMissingCounter 1 considerInterest + fillBlocks else do warningP "Unknown Piece" stopP -- True if the peer is a seeder isASeeder :: Process CF ST Bool isASeeder = do sdr <- gets missingPieces - sdr `deepseq` (return $! sdr == 0) + return $! sdr == 0 -- Decrease the counter of missing pieces for the peer decMissingCounter :: Int -> Process CF ST () @@ -319,34 +493,90 @@ pp <- createPeerPieces nPieces bf modify (\s -> s { peerPieces = pp }) peerLs <- PS.toList pp - pmch <- asks pieceMgrCh - liftIO . atomically $ writeTChan pmch (PeerHave peerLs) + msgPieceMgr (PeerHave peerLs) decMissingCounter (length peerLs) considerInterest else do infoP "Got out of band Bitfield request, dying" stopP +haveAllNoneMsg :: String -> Bool -> Process CF ST () +haveAllNoneMsg ty a = do + pieces <- gets peerPieces + piecesNull <- PS.null pieces + if piecesNull + then do nPieces <- succ . snd . bounds <$> asks pieceMap + pp <- createAllPieces nPieces a + modify (\s -> s { peerPieces = pp}) + peerLs <- PS.toList pp + msgPieceMgr (PeerHave peerLs) + decMissingCounter (length peerLs) + considerInterest + else do infoP $ "Got out of band " ++ ty ++ " request, dying" + stopP + +haveNoneMsg :: Process CF ST () +haveNoneMsg = haveAllNoneMsg "HaveNone" False + +haveAllMsg :: Process CF ST () +haveAllMsg = haveAllNoneMsg "HaveAll" True + + -- | Process a request message from the Peer requestMsg :: PieceNum -> Block -> Process CF ST () requestMsg pn blk = do choking <- gets weChoke - unless (choking) - (outChan $ SenderQ.SenderQPiece pn blk) + unless choking + (do debugP $ "Peer requested: " ++ show pn ++ "(" ++ show blk ++ ")" + outChan $ SenderQ.SenderQPiece pn blk) +requestFastMsg :: PieceNum -> Block -> Process CF ST () +requestFastMsg pn blk = do + choking <- gets weChoke + debugP $ "Peer fastRequested: " ++ show pn ++ "(" ++ show blk ++ ")" + if choking + then outChan $ SenderQ.SenderQM (RejectRequest pn blk) + else outChan $ SenderQ.SenderQPiece pn blk + -- | Handle a Piece Message incoming from the peer pieceMsg :: PieceNum -> Int -> B.ByteString -> Process CF ST () -pieceMsg n os bs = do +pieceMsg pn offs bs = pieceMsg' pn offs bs >> return () + +fastPieceMsg :: PieceNum -> Int -> B.ByteString -> Process CF ST () +fastPieceMsg pn offs bs = do + r <- pieceMsg' pn offs bs + unless r + (do infoP "Peer sent out-of-band piece we did not request, closing" + stopP) + +pieceMsg' :: PieceNum -> Int -> B.ByteString -> Process CF ST Bool +pieceMsg' n os bs = do let sz = B.length bs blk = Block os sz e = (n, blk) q <- gets blockQueue -- When e is not a member, the piece may be stray, so ignore it. -- Perhaps print something here. - when (S.member e q) - (do storeBlock n blk bs - bq <- gets blockQueue >>= return . S.delete e - bq `deepseq` modify (\s -> s { blockQueue = bq })) + if S.member e q + then do storeBlock n blk bs + bq <- gets blockQueue >>= return . S.delete e + s <- get + bq `deepseq` put $! s { blockQueue = bq } + return True + else return False +rejectMsg :: PieceNum -> Block -> Process CF ST () +rejectMsg pn blk = do + let e = (pn, blk) + q <- gets blockQueue + if S.member e q + then do + msgPieceMgr (PutbackBlocks [e]) + s <- get + put $! s { blockQueue = S.delete e q } + else do + infoP "Peer rejected piece/block we never requested, stopping" + stopP + -- | Handle a cancel message from the peer cancelMsg :: PieceNum -> Block -> Process CF ST () cancelMsg n blk = outChan $ SenderQ.SenderQCancel n blk @@ -357,21 +587,27 @@ considerInterest = do c <- asks interestTV pcs <- gets peerPieces - pmch <- asks pieceMgrCh - interested <- liftIO $ do - atomically $ writeTChan pmch (AskInterested pcs c) - atomically $ takeTMVar c + msgPieceMgr (AskInterested pcs c) + interested <- liftIO $ do atomically $ takeTMVar c if interested - then do modify (\s -> s { weInterested = True }) - outChan $ SenderQ.SenderQM Interested - else modify (\s -> s { weInterested = False}) + then do oldI <- gets weInterested + when (interested /= oldI) + (do modify (\s -> s { weInterested = True }) + debugP "We are interested" + outChan $ SenderQ.SenderQM Interested) + else do oldI <- gets weInterested + when (interested /= oldI) + (do modify (\s -> s { weInterested = False}) + debugP "We are not interested" + outChan $ SenderQ.SenderQM NotInterested) -- | Try to fill up the block queue at the peer. The reason we pipeline a -- number of blocks is to get around the line delay present on the internet. fillBlocks :: Process CF ST () fillBlocks = do choked <- gets peerChoke - unless choked checkWatermark + interested <- gets weInterested + when (not choked && interested) checkWatermark -- | check the current Watermark level. If we are below the lower one, then -- fill till the upper one. This in turn keeps the pipeline of pieces full as @@ -384,17 +620,15 @@ let sz = S.size q mark = if eg then endgameLoMark else loMark when (sz < mark) - (do - toQueue <- grabBlocks (hiMark - sz) - debugP $ "Got " ++ show (length toQueue) ++ " blocks: " ++ show toQueue - queuePieces toQueue) + (do toQueue <- grabBlocks (hiMark - sz) + queuePieces toQueue) -- These three values are chosen rather arbitrarily at the moment. loMark :: Int -loMark = 10 +loMark = 5 hiMark :: Int -hiMark = 15 +hiMark = 25 -- Low mark when running in endgame mode endgameLoMark :: Int @@ -404,18 +638,19 @@ -- | Queue up pieces for retrieval at the Peer queuePieces :: [(PieceNum, Block)] -> Process CF ST () queuePieces toQueue = do - mapM_ (uncurry pushRequest) toQueue - modify (\s -> s { blockQueue = S.union (blockQueue s) (S.fromList toQueue) }) - --- | Push a request to the peer so he can send it to us -pushRequest :: PieceNum -> Block -> Process CF ST () -pushRequest pn blk = outChan $ SenderQ.SenderQM $ Request pn blk + s <- get + let bq = blockQueue s + q <- forM toQueue + (\(p, b) -> do + if S.member (p, b) bq + then return Nothing -- Ignore pieces which are already in queue + else do outChan $ SenderQ.SenderQM $ Request p b + return $ Just (p, b)) + put $! s { blockQueue = S.union bq (S.fromList $ catMaybes q) } -- | Tell the PieceManager to store the given block storeBlock :: PieceNum -> Block -> B.ByteString -> Process CF ST () -storeBlock n blk bs = do - pmch <- asks pieceMgrCh - liftIO . atomically $ writeTChan pmch (StoreBlock n blk bs) +storeBlock n blk bs = msgPieceMgr (StoreBlock n blk bs) -- | The call @grabBlocks n@ will attempt to grab (up to) @n@ blocks from the -- piece Manager for request at the peer. @@ -423,19 +658,21 @@ grabBlocks n = do c <- asks grabBlockTV ps <- gets peerPieces - pmch <- asks pieceMgrCh - blks <- liftIO $ do - atomically $ writeTChan pmch (GrabBlocks n ps c) - atomically $ takeTMVar c + msgPieceMgr (GrabBlocks n ps c) + blks <- liftIO $ do atomically $ takeTMVar c case blks of Leech bs -> return bs Endgame bs -> modify (\s -> s { runningEndgame = True }) >> return bs -createPeerPieces :: MonadIO m => Int -> L.ByteString -> m PS.PieceSet +createAllPieces :: MonadIO m => Int -> Bool -> m PS.PieceSet +createAllPieces n False = PS.fromList n [] +createAllPieces n True = PS.fromList n [0..(n-1)] + +createPeerPieces :: MonadIO m => Int -> B.ByteString -> m PS.PieceSet createPeerPieces nPieces = - PS.fromList nPieces . map fromIntegral . concat . decodeBytes 0 . L.unpack + PS.fromList nPieces . map fromIntegral . concat . decodeBytes 0 . B.unpack where decodeByte :: Int -> Word8 -> [Maybe Int] decodeByte soFar w = let dBit n = if testBit w (7-n) @@ -451,3 +688,56 @@ modify (\st -> st { lastMessage = 0 }) ch <- asks outCh liftIO . atomically $ writeTChan ch qm + +msgPieceMgr :: PieceMgrMsg -> Process CF ST () +msgPieceMgr m = do + pmc <- asks pieceMgrCh + {-# SCC "Channel_Write" #-} liftIO . atomically $ writeTChan pmc m + + +-- IP address is given in host byte-order +allowedFast :: Word32 -> InfoHash -> Int -> Int -> IO [Word32] +allowedFast ip ihash sz n = generate n [] x [] + where + -- Take pieces from the generated ones and refill when it is exhausted. + -- While taking pieces, kill duplicates + generate 0 pcs _ _ = return $ reverse pcs + generate k pcs hsh (p : rest) + | p `elem` pcs = generate k pcs hsh rest + | otherwise = generate (k-1) (p : pcs) hsh rest + generate k pcs hsh [] = do + nhsh <- Digest.digestBS hsh + generate k pcs nhsh (genPieces nhsh) + + genPieces hash | B.null hash = [] + | otherwise = + let (h, rest) = B.splitAt 4 hash + bytes :: [Word32] + bytes = [fromIntegral z `shiftL` s | + (z, s) <- zip (B.unpack h) [24,16,8,0]] + ntohl = fromIntegral . sum + in ((ntohl bytes) `mod` fromIntegral sz) : genPieces rest + -- To prevent a Peer to reconnect, obtain a new IP and thus new FAST-set pieces, we mask out + -- the lower bits + ipBytes = B.pack $ map fromIntegral + [ (ip .&. 0xFF000000) `shiftR` 24 + , (ip .&. 0x00FF0000) `shiftR` 16 + , (ip .&. 0x0000FF00) `shiftR` 8 + , 0 + ] + x = B.concat [ipBytes, ihash] + +testSuite :: Test +testSuite = testGroup "Process/Peer" + [ testCase "AllowedFast" testAllowedFast ] + +testAllowedFast :: Assertion +testAllowedFast = do + pcs <- allowedFast (ip32 [80,4,4,200]) tHash 1313 7 + assertEqual "Test1" [1059,431,808,1217,287,376,1188] pcs + pcs' <- allowedFast (ip32 [80,4,4,200]) tHash 1313 9 + assertEqual "Test2" [1059,431,808,1217,287,376,1188,353,508] pcs' + where ip32 :: [Int] -> Word32 + ip32 bytes = fromIntegral $ sum [b `shiftL` s | (b, s) <- zip bytes [24,16,8,0]] + tHash :: B.ByteString + tHash = B.pack $ take 20 (repeat 0xaa) diff --git a/src/Process/Peer/Receiver.hs b/src/Process/Peer/Receiver.hs --- a/src/Process/Peer/Receiver.hs +++ b/src/Process/Peer/Receiver.hs @@ -8,15 +8,13 @@ import Control.Monad.Reader import Control.Monad.State -import Prelude hiding (catch, log) - import qualified Data.ByteString as B -import qualified Data.Serialize.Get as G - +import Prelude hiding (catch, log) -import Data.Word +import qualified Data.Attoparsec as A -import System.IO +import Network.Socket hiding (send, sendTo, recv, recvFrom) +import Network.Socket.ByteString import Process import Supervisor @@ -28,32 +26,29 @@ instance Logging CF where logName _ = "Process.Peer.Receiver" -start :: Handle -> TChan (Message, Integer) +start :: Socket -> TChan (Message, Integer) -> SupervisorChannel -> IO ThreadId -start h ch supC = do - hSetBuffering h NoBuffering - spawnP (CF ch) h - (catchP readSend +start s ch supC = do + spawnP (CF ch) s + ({-# SCC "Receiver" #-} catchP readSend (defaultStopHandler supC)) -readSend :: Process CF Handle () +readSend :: Process CF Socket () readSend = do - h <- get + s <- get c <- asks rpMsgCh - bs' <- liftIO $ B.hGet h 4 - l <- conv bs' - if (l == 0) - then return () - else do bs <- {-# SCC "hGet_From_BS" #-} liftIO $ B.hGet h (fromIntegral l) - case G.runGet decodeMsg bs of - Left _ -> do warningP "Incorrect parse in receiver, dying!" - stopP - Right msg -> liftIO . atomically $ writeTChan c (msg, fromIntegral l) - readSend + bs <- liftIO $ recv s 2048 + when (B.length bs == 0) stopP + loop s c (A.parse getMsg bs) + where loop s c (A.Done r msg) = do + liftIO . atomically $ writeTChan c (msg, fromIntegral $ msgSize msg) + loop s c (A.parse getMsg r) + loop s c (prt@(A.Partial _)) = do + bs <- liftIO $ recv s 4096 + when (B.length bs == 0) stopP + loop s c (A.feed prt bs) + loop _ _ (A.Fail _ ctx err) = + do warningP $ "Incorrect parse in receiver, context: " + ++ show ctx ++ ", " ++ show err + stopP -conv :: B.ByteString -> Process CF Handle Word32 -conv bs = do - case G.runGet G.getWord32be bs of - Left err -> do warningP $ "Incorrent parse in receiver, dying: " ++ show err - stopP - Right i -> return i diff --git a/src/Process/Peer/Sender.hs b/src/Process/Peer/Sender.hs --- a/src/Process/Peer/Sender.hs +++ b/src/Process/Peer/Sender.hs @@ -9,31 +9,33 @@ import qualified Data.ByteString.Lazy as L -import System.IO +import Network.Socket hiding (send, sendTo, recv, recvFrom) +import Network.Socket.ByteString.Lazy +import Prelude hiding (getContents) import Process import Supervisor data CF = CF { chan :: TMVar L.ByteString - , hndl :: Handle } + , sock :: Socket } instance Logging CF where logName _ = "Process.Peer.Sender" -- | The raw sender process, it does nothing but send out what it syncs on. -start :: Handle -> TMVar L.ByteString -> SupervisorChannel -> IO ThreadId -start h ch supC = spawnP (CF ch h) () (catchP pgm +start :: Socket -> TMVar L.ByteString -> SupervisorChannel -> IO ThreadId +start s ch supC = spawnP (CF ch s) () ({-# SCC "Sender" #-} + catchP pgm (do t <- liftIO $ myThreadId liftIO . atomically $ writeTChan supC $ IAmDying t - liftIO $ hClose h)) + liftIO $ sClose s)) pgm :: Process CF () () pgm = do ch <- asks chan - h <- asks hndl - liftIO $ do + s <- asks sock + _ <- liftIO $ do r <- atomically $ takeTMVar ch - L.hPut h r - hFlush h + send s r pgm diff --git a/src/Process/Peer/SenderQ.hs b/src/Process/Peer/SenderQ.hs --- a/src/Process/Peer/SenderQ.hs +++ b/src/Process/Peer/SenderQ.hs @@ -14,6 +14,7 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Lazy as L +import Data.List (foldl') import Channels import Process @@ -35,6 +36,7 @@ , bandwidthCh :: BandwidthChannel , readBlockTV :: TMVar B.ByteString , fsCh :: FSPChannel + , fastExtension :: Bool } data ST = ST { outQueue :: !(Q.Queue (Either Message (PieceNum, Block))) @@ -46,12 +48,13 @@ -- | sendQueue Process, simple version. -- TODO: Split into fast and slow. -start :: TChan SenderQMsg -> TMVar L.ByteString -> BandwidthChannel +start :: [Capabilities] -> TChan SenderQMsg -> TMVar L.ByteString -> BandwidthChannel -> FSPChannel -> SupervisorChannel -> IO ThreadId -start inC outC bandwC fspC supC = do +start caps inC outC bandwC fspC supC = do rbtv <- liftIO newEmptyTMVarIO - spawnP (CF inC outC bandwC rbtv fspC) (ST Q.empty 0) - (catchP pgm + spawnP (CF inC outC bandwC rbtv fspC + (Fast `elem` caps)) (ST Q.empty 0) + ({-# SCC "SenderQ" #-} catchP pgm (defaultStopHandler supC)) pgm :: Process CF ST () @@ -81,11 +84,40 @@ case m of SenderQM msg -> modifyQ (Q.push $ Left msg) SenderQPiece n blk -> modifyQ (Q.push $ Right (n, blk)) - SenderQCancel n blk -> modifyQ (Q.filter (filterPiece n blk)) - SenderOChoke -> do modifyQ (Q.filter filterAllPiece) - modifyQ (Q.push $ Left Choke) - SenderQRequestPrune n blk -> - modifyQ (Q.filter (filterRequest n blk)) + SenderQCancel n blk -> do + fe <- asks fastExtension + if fe + then do + piece <- partitionQ (pickPiece n blk) + case piece of + [] -> return () -- Piece must have been sent + [_] -> modifyQ (Q.push (Left $ RejectRequest n blk)) + ps -> fail $ "Impossible case, SenderQCancel " ++ show (length ps) + else modifyQ (Q.filter (filterPiece n blk)) + SenderOChoke -> do + fe <- asks fastExtension + if fe + then do + -- In the fast extension, we explicitly reject all pieces + pieces <- partitionQ filterAllPiece + modifyQ (Q.push $ Left Choke) + let rejects = map (\(Right (pn, blk)) -> Left $ RejectRequest pn blk) + pieces + modifyQ (flip (foldl' (flip Q.push)) rejects) + else do modifyQ (Q.filter filterAllPiece) + modifyQ (Q.push $ Left Choke) + SenderQRequestPrune n blk -> do + fe <- asks fastExtension + piece <- partitionQ (pickRequest n blk) + case piece of + [] -> modifyQ (Q.push (Left $ Cancel n blk)) -- Request must have been sent + [_] -> if fe + then modifyQ -- This is a hack for now + (Q.push (Left $ Cancel n blk) . + Q.push (Left $ Request n blk)) + else return () + ps -> fail $ "Impossible case, SenderQRequestPrune " + ++ show ps ++ ", " ++ show fe pgm rateUpdateEvent :: Process CF ST () @@ -95,24 +127,37 @@ liftIO . atomically $ writeTChan bwc l modify (\s -> s { bytesTransferred = 0 }) -filterAllPiece :: Either Message (PieceNum, Block) -> Bool +-- The type of the Outgoing queue +type OutQT = Either Message (PieceNum, Block) + +filterAllPiece :: OutQT -> Bool filterAllPiece (Right _) = True filterAllPiece (Left _) = False -filterPiece :: PieceNum -> Block -> Either Message (PieceNum, Block) -> Bool +pickPiece :: PieceNum -> Block -> OutQT -> Bool +pickPiece n blk (Right (n1, blk1)) | n == n1 && blk == blk1 = True +pickPiece _ _ _ = False + +filterPiece :: PieceNum -> Block -> OutQT -> Bool filterPiece n blk (Right (n1, blk1)) | n == n1 && blk == blk1 = False | otherwise = True filterPiece _ _ _ = True -filterRequest :: PieceNum -> Block -> Either Message (PieceNum, Block) -> Bool -filterRequest n blk (Left (Request n1 blk1)) | n == n1 && blk == blk1 = False - | otherwise = True -filterRequest _ _ _ = True +pickRequest :: PieceNum -> Block -> OutQT -> Bool +pickRequest n blk (Left (Request n1 blk1)) | n == n1 && blk == blk1 = True +pickRequest _ _ _ = False -modifyQ :: (Q.Queue (Either Message (PieceNum, Block)) -> - Q.Queue (Either Message (PieceNum, Block))) +modifyQ :: (Q.Queue (OutQT) -> + Q.Queue (OutQT)) -> Process CF ST () modifyQ f = modify (\s -> s { outQueue = f $! outQueue s }) + +partitionQ :: (OutQT -> Bool) -> Process CF ST [OutQT] +partitionQ p = do + s <- get + let (as, nq) = Q.partition p $ outQueue s + put $! s { outQueue = nq } + return as -- | Read a block from the filesystem for sending readBlock :: PieceNum -> Block -> Process CF ST B.ByteString diff --git a/src/Process/PeerMgr.hs b/src/Process/PeerMgr.hs --- a/src/Process/PeerMgr.hs +++ b/src/Process/PeerMgr.hs @@ -22,8 +22,7 @@ import Data.Array import qualified Data.Map as M -import Network -import System.IO +import qualified Network.Socket as Sock import System.Log.Logger import Channels @@ -39,7 +38,7 @@ import Torrent hiding (infoHash) data PeerMgrMsg = PeersFromTracker InfoHash [Peer] - | NewIncoming (Handle, HostName, PortNumber) + | NewIncoming (Sock.Socket, Sock.SockAddr) | NewTorrent InfoHash TorrentLocal | StopTorrent InfoHash @@ -82,7 +81,7 @@ fakeChan <- newTChanIO pool <- liftM snd $ oneForOne "PeerPool" [] fakeChan spawnP (CF ch mgrC pool chokeMgrC rtv) - (ST [] M.empty pid cmap) (catchP lp + (ST [] M.empty pid cmap) ({-# SCC "PeerMgr" #-} catchP lp (defaultStopHandler supC)) where cmap = M.empty @@ -104,14 +103,14 @@ PeersFromTracker ih ps -> do debugP "Adding peers to queue" modify (\s -> s { peersInQueue = (map (ih,) ps) ++ peersInQueue s }) - NewIncoming conn@(h, _, _) -> do + NewIncoming conn@(s, _) -> do sz <- liftM M.size $ gets peers if sz < numPeers then do debugP "New incoming peer, handling" _ <- addIncoming conn return () else do debugP "Already too many peers, closing!" - liftIO $ hClose h + liftIO $ Sock.sClose s NewTorrent ih tl -> do modify (\s -> s { cmMap = M.insert ih tl (cmMap s)}) StopTorrent _ih -> do @@ -147,15 +146,15 @@ modify (\s -> s { peersInQueue = rest })) addPeer :: (InfoHash, Peer) -> Process CF ST ThreadId -addPeer (ih, (Peer hn prt)) = do +addPeer (ih, (Peer addr)) = do ppid <- gets peerId pool <- asks peerPool mgrC <- asks mgrCh cm <- gets cmMap v <- asks chokeRTV - liftIO $ connect (hn, prt, ppid, ih) pool mgrC v cm + liftIO $ connect (addr, ppid, ih) pool mgrC v cm -addIncoming :: (Handle, HostName, PortNumber) -> Process CF ST ThreadId +addIncoming :: (Sock.Socket, Sock.SockAddr) -> Process CF ST ThreadId addIncoming conn = do ppid <- gets peerId pool <- asks peerPool @@ -164,56 +163,58 @@ cm <- gets cmMap liftIO $ acceptor conn pool ppid mgrC v cm -type ConnectRecord = (HostName, PortID, PeerId, InfoHash) +type ConnectRecord = (Sock.SockAddr, PeerId, InfoHash) connect :: ConnectRecord -> SupervisorChannel -> MgrChannel -> RateTVar -> ChanManageMap -> IO ThreadId -connect (host, port, pid, ih) pool mgrC rtv cmap = +connect (addr, pid, ih) pool mgrC rtv cmap = forkIO (connector >> return ()) where - connector = - do debugM "Process.PeerMgr.connect" $ - "Connecting to " ++ show host ++ " (" ++ showPort port ++ ")" - h <- connectTo host port + connector = {-# SCC "connect" #-} + do sock <- Sock.socket Sock.AF_INET Sock.Stream Sock.defaultProtocol + debugM "Process.PeerMgr.connect" $ "Connecting to: " ++ show addr + Sock.connect sock addr debugM "Process.PeerMgr.connect" "Connected, initiating handShake" - r <- initiateHandshake h pid ih + r <- initiateHandshake sock pid ih debugM "Process.PeerMgr.connect" "Handshake run" case r of Left err -> do debugM "Process.PeerMgr.connect" - ("Peer handshake failure at host " ++ host + ("Peer handshake failure at host " ++ show addr ++ " with error " ++ err) return () - Right (_caps, _rpid, ihsh) -> + Right (caps, _rpid, ihsh) -> do debugM "Process.PeerMgr.connect" "entering peerP loop code" let tc = case M.lookup ihsh cmap of Nothing -> error "Impossible (2), I hope" Just x -> x - children <- Peer.start h mgrC rtv (tcPcMgrCh tc) (tcFSCh tc) (tcStatTV tc) - (tcPM tc) (succ . snd . bounds $ tcPM tc) ihsh + children <- Peer.start sock caps mgrC rtv + (tcPcMgrCh tc) (tcFSCh tc) (tcStatTV tc) + (tcPM tc) (succ . snd . bounds $ tcPM tc) + ihsh atomically $ writeTChan pool $ SpawnNew (Supervisor $ allForOne "PeerSup" children) return () -acceptor :: (Handle, HostName, PortNumber) -> SupervisorChannel +acceptor :: (Sock.Socket, Sock.SockAddr) -> SupervisorChannel -> PeerId -> MgrChannel -> RateTVar -> ChanManageMap -> IO ThreadId -acceptor (h,hn,pn) pool pid mgrC rtv cmmap = +acceptor (s,sa) pool pid mgrC rtv cmmap = forkIO (connector >> return ()) where ihTst k = M.member k cmmap - connector = do + connector = {-# SCC "acceptor" #-} do debugLog "Handling incoming connection" - r <- receiveHandshake h pid ihTst + r <- receiveHandshake s pid ihTst debugLog "RecvHandshake run" case r of - Left err -> do debugLog ("Incoming Peer handshake failure with " ++ show hn ++ "(" - ++ show pn ++ "), error: " ++ err) - return () - Right (_caps, _rpid, ih) -> + Left err -> do debugLog ("Incoming Peer handshake failure with " + ++ show sa ++ ", error: " ++ err) + return() + Right (caps, _rpid, ih) -> do debugLog "entering peerP loop code" let tc = case M.lookup ih cmmap of Nothing -> error "Impossible, I hope" Just x -> x - children <- Peer.start h mgrC rtv (tcPcMgrCh tc) (tcFSCh tc) + children <- Peer.start s caps mgrC rtv (tcPcMgrCh tc) (tcFSCh tc) (tcStatTV tc) (tcPM tc) (succ . snd . bounds $ tcPM tc) ih atomically $ writeTChan pool $ @@ -221,6 +222,3 @@ return () debugLog = debugM "Process.PeerMgr.acceptor" -showPort :: PortID -> String -showPort (PortNumber pn) = show pn -showPort _ = "N/A" diff --git a/src/Process/PieceMgr.hs b/src/Process/PieceMgr.hs --- a/src/Process/PieceMgr.hs +++ b/src/Process/PieceMgr.hs @@ -138,9 +138,8 @@ start :: PieceMgrChannel -> FSPChannel -> ChokeMgrChannel -> StatusChannel -> ST -> InfoHash -> SupervisorChannel -> IO ThreadId start mgrC fspC chokeC statC db ih supC = - {-# SCC "PieceMgr" #-} spawnP (CF mgrC fspC chokeC statC ih) db - (catchP eventLoop + ({-# SCC "PieceMgr" #-} catchP eventLoop (defaultStopHandler supC)) eventLoop :: Process CF ST () @@ -155,24 +154,22 @@ dl <- gets donePush if (null dl) then return () - else sendChokeMgr (head dl) >> drainSend - -sendChokeMgr :: ChokeMgrMsg -> Process CF ST () -sendChokeMgr e = do - c <- asks chokeCh - liftIO . atomically $ writeTChan c e - modify (\db -> db { donePush = tail (donePush db) }) + else do + c <- asks chokeCh + liftIO . atomically $ writeTChan c (head dl) + s <- get + put $! s { donePush = tail (donePush s) } traceMsg :: PieceMgrMsg -> Process CF ST () traceMsg m = do tb <- gets traceBuffer - let ntb = (trace $! show m) tb + let !ntb = (trace $! show m) tb modify (\db -> db { traceBuffer = ntb }) rpcMessage :: Process CF ST () rpcMessage = do ch <- asks pieceMgrCh - m <- liftIO . atomically $ readTChan ch + m <- {-# SCC "Channel_Read" #-} liftIO . atomically $ readTChan ch traceMsg m case m of GrabBlocks n eligible c -> {-# SCC "GrabBlocks" #-} @@ -191,7 +188,6 @@ intr <- askInterested pieces liftIO . atomically $ do putTMVar retC intr -- And this neither too! - storeBlock :: PieceNum -> Block -> B.ByteString -> Process CF ST () storeBlock pn blk d = do debugP $ "Storing block: " ++ show (pn, blk) @@ -270,7 +266,7 @@ createPieceDb mmap pmap = do pending <- filt (==False) done <- filt (==True) - return $ ST pending done [] M.empty [] pmap False PendS.empty 0 (Tracer.new 20) + return $ ST pending done [] M.empty [] pmap False PendS.empty 0 (Tracer.new 25) where filt f = PS.fromList (succ . snd . bounds $ pmap) . M.keys $ M.filter f mmap @@ -388,36 +384,41 @@ -- returns a list of Blocks which where grabbed. grabBlocks :: Int -> PS.PieceSet -> PieceMgrProcess Blocks grabBlocks k eligible = {-# SCC "grabBlocks" #-} do - ps' <- PS.copy eligible - blocks <- tryGrabProgress k ps' [] + blocks <- tryGrab k eligible pend <- gets pendingPieces pendN <- PS.null pend if blocks == [] && pendN - then do ps'' <- PS.copy eligible - blks <- grabEndGame k ps'' - modify (\db -> db { endGaming = True }) + then do blks <- grabEndGame k eligible + db <- get + put $! db { endGaming = True } debugP $ "PieceMgr entered endgame." return $ Endgame blks - else do modify (\s -> s { downloading = blocks ++ (downloading s) }) + else do s <- get + let !dld = downloading s + put $! s { downloading = blocks ++ dld } return $ Leech blocks -- Grabbing blocks is a state machine implemented by tail calls -- Try grabbing pieces from the pieces in progress first -tryGrabProgress :: PieceNum -> PS.PieceSet -> [(PieceNum, Block)] +tryGrab :: PieceNum -> PS.PieceSet -> Process CF ST [(PieceNum, Block)] +tryGrab k ps = {-# SCC "tryGrabProgress" #-} + tryGrabProgress k ps [] =<< (M.keys <$> gets inProgress) + +tryGrabProgress :: PieceNum -> PS.PieceSet -> [(PieceNum, Block)] -> [PieceNum] -> Process CF ST [(PieceNum, Block)] -tryGrabProgress 0 _ captured = return captured -tryGrabProgress n ps captured = grabber =<< (M.keys <$> gets inProgress) - where - grabber [] = tryGrabPending n ps captured - grabber (i : is) = do m <- PS.member i ps - if m - then grabFromProgress n ps i captured - else grabber is +tryGrabProgress 0 _ captured _ = return captured +tryGrabProgress k ps captured [] = {-# SCC "tryGrabProgress_k_e" #-} + tryGrabPending k ps captured +tryGrabProgress k ps captured (i : is) = {-# SCC "tryGrabProgress_k_is" #-} + do m <- PS.member i ps + if m + then grabFromProgress k ps i captured is + else tryGrabProgress k ps captured is -- The Piece @p@ was found, grab it -grabFromProgress :: PieceNum -> PS.PieceSet -> PieceNum -> [(PieceNum, Block)] +grabFromProgress :: PieceNum -> PS.PieceSet -> PieceNum -> [(PieceNum, Block)] -> [PieceNum] -> Process CF ST [(PieceNum, Block)] -grabFromProgress n ps p captured = do +grabFromProgress n ps p captured nxt = {-# SCC "grabFromProgress" #-} do inprog <- gets inProgress ipp <- case M.lookup p inprog of Nothing -> fail "grabFromProgress: could not lookup piece" @@ -426,26 +427,30 @@ nIpp = ipp { ipPendingBlocks = rest } -- This rather ugly piece of code should be substituted with something better if grabbed == [] - -- All pieces are taken, try the next one. - then do PS.delete p ps --TODO: Dangerous since we will NEVER reconsider that piece then! - tryGrabProgress n ps captured - else do modify (\db -> db { inProgress = M.insert p nIpp inprog }) - tryGrabProgress (n - length grabbed) ps ([(p,g) | g <- grabbed] ++ captured) + then tryGrabProgress n ps captured nxt + else do modify (\db -> db { inProgress = M.insert p nIpp inprog }) + tryGrabProgress + (n - length grabbed) ps + ([(p,g) | g <- grabbed] ++ captured) + nxt -- Try grabbing pieces from the pending blocks -tryGrabPending :: PieceNum -> PS.PieceSet -> [(PieceNum, Block)] -> Process CF ST [(PieceNum, Block)] -tryGrabPending n ps captured = do +tryGrabPending :: PieceNum -> PS.PieceSet -> [(PieceNum, Block)] + -> Process CF ST [(PieceNum, Block)] +tryGrabPending n ps captured = {-# SCC "tryGrabPending" #-} do histo <- gets histogram pending <- gets pendingPieces - selector <- PS.freeze ps - pendingS <- PS.freeze pending - let culprits = PendS.pick (\p -> selector p && pendingS p) histo + culprits <- {-# SCC "PendS.pick" #-} + liftIO $ PendS.pick (\p -> do a <- PS.member p ps + b <- PS.member p pending + return $ a && b) + histo case culprits of - Nothing -> do + Nothing -> {-# SCC "No_Culprits" #-} do isn <- PS.intersection ps pending assert (null isn) (return ()) return captured - Just pieces -> do + Just pieces -> {-# SCC "Build_Blocks" #-} do h <- pickRandom pieces inProg <- gets inProgress blockList <- createBlock h @@ -453,10 +458,11 @@ ipp = InProgressPiece sz S.empty blockList PS.delete h =<< gets pendingPieces modify (\db -> db { inProgress = M.insert h ipp inProg }) - tryGrabProgress n ps captured + tryGrabProgress n ps captured [h] grabEndGame :: PieceNum -> PS.PieceSet -> Process CF ST [(PieceNum, Block)] -grabEndGame n ps = do -- In endgame we are allowed to grab from the downloaders +grabEndGame n ps = {-# SCC "grabEndGame" #-} do + -- In endgame we are allowed to grab from the downloaders dls <- filterM (\(p, _) -> PS.member p ps) =<< gets downloading take n . shuffle' dls (length dls) <$> liftIO newStdGen @@ -481,7 +487,7 @@ assertST = {-# SCC "assertST" #-} do c <- gets assertCount if c == 0 - then do modify (\db -> db { assertCount = 10 }) + then do modify (\db -> db { assertCount = 25 }) assertSets >> assertInProgress >> assertDownloading sizes <- sizeReport debugP sizes diff --git a/src/Process/Status.hs b/src/Process/Status.hs --- a/src/Process/Status.hs +++ b/src/Process/Status.hs @@ -95,7 +95,7 @@ start fp statusC tv supC = do r <- newIORef (0,0) spawnP (CF statusC tv) M.empty - (cleanupP (pgm r) (defaultStopHandler supC) (cleanup r)) + ({-# SCC "Status" #-} cleanupP (pgm r) (defaultStopHandler supC) (cleanup r)) where cleanup r = do st <- liftIO $ readIORef r @@ -146,6 +146,7 @@ liftIO . atomically $ writeTChan (trackerMsgCh ns) Complete modify (\s -> M.insert ih (ns { state = Seeding}) s) + fetchUpdates :: IORef (Integer, Integer) -> Process CF ST () fetchUpdates r = do tv <- asks statusTV @@ -159,7 +160,8 @@ ndn = d + down liftIO $ nup `deepseq` ndn `deepseq` writeIORef r (nup, ndn) s <- get - put $! M.adjust (\st -> - st `deepseq` st { uploaded = (uploaded st) + up - , downloaded = (downloaded st) + down }) ih s) updates + let adjusted = M.adjust (\st -> + st { uploaded = (uploaded st) + up + , downloaded = (downloaded st) + down }) ih s + deepseq adjusted (put adjusted)) updates diff --git a/src/Process/Timer.hs b/src/Process/Timer.hs --- a/src/Process/Timer.hs +++ b/src/Process/Timer.hs @@ -15,6 +15,6 @@ import Control.Monad.Trans registerSTM :: MonadIO m => Int -> TChan a -> a -> m ThreadId -registerSTM secs c m = liftIO $ forkIO $ do +registerSTM secs c m = liftIO $ forkIO $ {-# SCC "Timer" #-} do threadDelay (secs * 1000000) atomically $ writeTChan c m diff --git a/src/Process/TorrentManager.hs b/src/Process/TorrentManager.hs --- a/src/Process/TorrentManager.hs +++ b/src/Process/TorrentManager.hs @@ -59,7 +59,7 @@ -> IO ThreadId start chan statusC stv chokeC pid peerC supC = spawnP (CF chan statusC stv pid peerC chokeC) (ST []) - (catchP pgm (defaultStopHandler supC)) + ({-# SCC "TorrentManager" #-} catchP pgm (defaultStopHandler supC)) where pgm = startStop >> dirMsg >> pgm dirMsg = do c <- asks tCh diff --git a/src/Process/Tracker.hs b/src/Process/Tracker.hs --- a/src/Process/Tracker.hs +++ b/src/Process/Tracker.hs @@ -18,11 +18,13 @@ import Control.Monad.Reader import Control.Monad.State +import Data.Bits import Data.Char (ord, chr) import Data.List (intersperse) import qualified Data.ByteString as B +import Data.Word -import Network +import Network.Socket as S import Network.HTTP hiding (port) import Network.URI hiding (unreserved) @@ -88,16 +90,16 @@ torrentInfo :: TorrentInfo , peerId :: PeerId , state :: TrackerEvent - , localPort :: PortID + , localPort :: S.PortNumber , nextTick :: Integer } -start :: InfoHash -> TorrentInfo -> PeerId -> PortID +start :: InfoHash -> TorrentInfo -> PeerId -> S.PortNumber -> Status.StatusChannel -> TrackerChannel -> PeerMgr.PeerMgrChannel -> SupervisorChannel -> IO ThreadId start ih ti pid port statusC msgC pc supC = spawnP (CF statusC msgC pc ih) (ST ti pid Stopped port 0) - (cleanupP loop + ({-# SCC "Tracker" #-} cleanupP loop (defaultStopHandler supC) stopEvent) where @@ -200,17 +202,45 @@ <*> (pure $ BCode.trackerMinInterval d) -decodeIps :: B.ByteString -> [PeerMgr.Peer] -decodeIps = decodeIps' . fromBS +decodeIps :: (B.ByteString, B.ByteString) -> [PeerMgr.Peer] +decodeIps (ipv4, ipv6) = decodeIps4 ipv4 ++ decodeIps6 ipv6 --- Decode a list of IP addresses. We expect these to be a compact response by default. -decodeIps' :: String -> [PeerMgr.Peer] -decodeIps' [] = [] -decodeIps' (b1:b2:b3:b4:p1:p2 : rest) = PeerMgr.Peer ip port : decodeIps' rest - where ip = concat . intersperse "." . map (show . ord) $ [b1, b2, b3, b4] - port = PortNumber . fromIntegral $ ord p1 * 256 + ord p2 -decodeIps' _xs = [] +decodeIps4 :: B.ByteString -> [PeerMgr.Peer] +decodeIps4 bs | B.null bs = [] + | B.length bs >= 6 = + let (ip, r1) = B.splitAt 4 bs + (port, r2) = B.splitAt 2 r1 + i' = cW32 ip + p' = PortNum $ cW16 port + in PeerMgr.Peer (S.SockAddrInet p' i') : decodeIps4 r2 + | otherwise = [] -- Some trackers fail spectacularly +decodeIps6 :: B.ByteString -> [PeerMgr.Peer] +decodeIps6 bs | B.null bs = [] + | B.length bs >= 18 = + let (ip6, r1) = B.splitAt 16 bs + (port, r2) = B.splitAt 2 r1 + i' = cW128 ip6 + p' = PortNum $ cW16 port + in PeerMgr.Peer (S.SockAddrInet6 p' 0 i' 0) : decodeIps6 r2 + | otherwise = [] -- Some trackers fail spectacularly + +cW32 :: B.ByteString -> Word32 +cW32 bs = fromIntegral . sum $ s + where up = B.unpack bs + s = [ fromIntegral b `shiftL` sa | (b, sa) <- zip up [0,8,16,24]] :: [Word32] + +cW16 :: B.ByteString -> Word16 +cW16 bs = fromIntegral . sum $ s + where s = [ fromIntegral b `shiftL` sa | (b, sa) <- zip (B.unpack bs) [0,8]] :: [Word16] + +cW128 :: B.ByteString -> (Word32, Word32, Word32, Word32) +cW128 bs = + let (q1, r1) = B.splitAt 4 bs + (q2, r2) = B.splitAt 4 r1 + (q3, q4) = B.splitAt 4 r2 + in (cW32 q1, cW32 q2, cW32 q3, cW32 q4) + trackerRequest :: URI -> Process CF ST (Either String TrackerResponse) trackerRequest uri = do resp <- liftIO $ simpleHTTP request @@ -257,8 +287,7 @@ prt :: Process CF ST Integer prt = do lp <- gets localPort case lp of - PortNumber pnum -> return $ fromIntegral pnum - _ -> do fail "Unknown port type" + PortNum pnum -> return $ fromIntegral pnum trackerfyEvent ev = case ev of Running -> [] diff --git a/src/Protocol/BCode.hs b/src/Protocol/BCode.hs --- a/src/Protocol/BCode.hs +++ b/src/Protocol/BCode.hs @@ -250,8 +250,10 @@ trackerError = searchStr "failure reason" trackerWarning = searchStr "warning mesage" -trackerPeers :: BCode -> Maybe B.ByteString -trackerPeers = searchStr "peers" +trackerPeers :: BCode -> Maybe (B.ByteString, B.ByteString) +trackerPeers bc = do v4 <- searchStr "peers" bc + v6 <- return $ maybe (B.empty) id $ searchStr "peers6" bc + return (v4, v6) info :: BCode -> Maybe BCode info = search [toPS "info"] diff --git a/src/Protocol/Wire.hs b/src/Protocol/Wire.hs --- a/src/Protocol/Wire.hs +++ b/src/Protocol/Wire.hs @@ -7,8 +7,11 @@ module Protocol.Wire ( Message(..) + , msgSize , encodePacket , decodeMsg + , getMsg -- ^ Run attoparsec-based parser on input + , getAPMsg , BitField , constructBitField -- Handshaking @@ -25,12 +28,19 @@ import qualified Data.ByteString as B import qualified Data.ByteString.Lazy as L +import Data.Attoparsec as A +import Data.Bits (testBit, setBit) + import Data.Serialize import Data.Serialize.Put import Data.Serialize.Get import Data.Char -import System.IO +import Data.Maybe (catMaybes) +import Network.Socket hiding (send, sendTo, recv, recvFrom, KeepAlive) +import Network.Socket.ByteString +import qualified Network.Socket.ByteString.Lazy as Lz + import System.Log.Logger import Test.Framework @@ -41,7 +51,7 @@ ------------------------------------------------------------ -type BitField = L.ByteString +type BitField = B.ByteString data Message = KeepAlive | Choke @@ -54,15 +64,44 @@ | Piece PieceNum Int B.ByteString | Cancel PieceNum Block | Port Integer + | HaveAll + | HaveNone + | Suggest PieceNum + | RejectRequest PieceNum Block + | AllowedFast PieceNum + | ExtendedMsg Int B.ByteString deriving (Eq, Show) +msgSize :: Message -> Int +msgSize KeepAlive = 0 +msgSize Choke = 1 +msgSize Unchoke = 1 +msgSize Interested = 1 +msgSize NotInterested = 1 +msgSize (Have _) = 5 +msgSize (BitField bf) = B.length bf + 1 +msgSize (Request _ _) = 13 +msgSize (Piece _ _ bs) = 9 + B.length bs +msgSize (Cancel _ _) = 13 +msgSize (Port _) = 3 +msgSize HaveAll = 1 +msgSize HaveNone = 1 +msgSize (Suggest _) = 5 +msgSize (RejectRequest _ _) = 13 +msgSize (AllowedFast _) = 5 +msgSize (ExtendedMsg _ bs) = B.length bs + 5 + instance Arbitrary Message where arbitrary = oneof [return KeepAlive, return Choke, return Unchoke, return Interested, - return NotInterested, + return NotInterested, return HaveAll, return HaveNone, + Suggest <$> pos, + RejectRequest <$> pos <*> arbitrary, + AllowedFast <$> pos, Have <$> pos, BitField <$> arbitrary, Request <$> pos <*> arbitrary, Piece <$> pos <*> pos <*> arbitrary, + ExtendedMsg <$> pos <*> arbitrary, Cancel <$> pos <*> arbitrary, Port <$> choose (0,16383)] where @@ -75,7 +114,9 @@ protocolHeader = "BitTorrent protocol" extensionBasis :: Word64 -extensionBasis = 0 +extensionBasis = + (flip setBit 2) -- Fast extension + 0 p8 :: Word8 -> Put p8 = putWord8 @@ -99,33 +140,100 @@ put Interested = p8 2 put NotInterested = p8 3 put (Have pn) = p8 4 *> p32be pn - put (BitField bf) = p8 5 *> putLazyByteString bf + put (BitField bf) = p8 5 *> putByteString bf put (Request pn (Block os sz)) = p8 6 *> mapM_ p32be [pn,os,sz] put (Piece pn os c) = p8 7 *> mapM_ p32be [pn,os] *> putByteString c put (Cancel pn (Block os sz)) = p8 8 *> mapM_ p32be [pn,os,sz] put (Port p) = p8 9 *> (putWord16be . fromIntegral $ p) + put (Suggest pn) = p8 0x0D *> p32be pn + put (ExtendedMsg idx bs) + = p8 20 *> p32be idx *> putByteString bs + put HaveAll = p8 0x0E + put HaveNone = p8 0x0F + put (RejectRequest pn (Block os sz)) + = p8 0x10 *> mapM_ p32be [pn,os,sz] + put (AllowedFast pn) + = p8 0x11 *> p32be pn get = getKA <|> getChoke <|> getUnchoke <|> getIntr <|> getNI <|> getHave <|> getBF <|> getReq <|> getPiece <|> getCancel - <|> getPort + <|> getPort <|> getHaveAll + <|> getHaveNone + <|> getSuggest <|> getRejectRequest + <|> getAllowedFast + <|> getExtendedMsg +getMsg :: Parser Message +getMsg = do + l <- apW32be + if l == 0 + then return KeepAlive + else getAPMsg l + +getAPMsg :: Int -> Parser Message +getAPMsg l = do + c <- A.anyWord8 + case c of + 0 -> return Choke + 1 -> return Unchoke + 2 -> return Interested + 3 -> return NotInterested + 4 -> (Have <$> apW32be) + 5 -> (BitField <$> (A.take (l-1))) + 6 -> (Request <$> apW32be <*> (Block <$> apW32be <*> apW32be)) + 7 -> (Piece <$> apW32be <*> apW32be <*> A.take (l - 9)) + 8 -> (Cancel <$> apW32be <*> (Block <$> apW32be <*> apW32be)) + 9 -> (Port . fromIntegral <$> apW16be) + 0x0D -> (Suggest <$> apW32be) + 0x0E -> return HaveAll + 0x0F -> return HaveNone + 0x10 -> (RejectRequest <$> apW32be <*> (Block <$> apW32be <*> apW32be)) + 0x11 -> (AllowedFast <$> apW32be) + 20 -> (ExtendedMsg <$> apW32be <*> A.take (l - 5)) + k -> fail $ "Illegal parse, code: " ++ show k + +apW32be :: Parser Int +apW32be = do + [b1,b2,b3,b4] <- replicateM 4 A.anyWord8 + let b1' = fromIntegral b1 + b2' = fromIntegral b2 + b3' = fromIntegral b3 + b4' = fromIntegral b4 + return (b4' + (256 * b3') + (256 * 256 * b2') + (256 * 256 * 256 * b1')) + +apW16be :: Parser Int +apW16be = do + [b1, b2] <- replicateM 2 A.anyWord8 + let b1' = fromIntegral b1 + b2' = fromIntegral b2 + return (b2' + 256 * b1') + + getBF, getChoke, getUnchoke, getIntr, getNI, getHave, getReq :: Get Message getPiece, getCancel, getPort, getKA :: Get Message -getChoke = byte 0 *> return Choke -getUnchoke = byte 1 *> return Unchoke -getIntr = byte 2 *> return Interested -getNI = byte 3 *> return NotInterested -getHave = byte 4 *> (Have <$> gw32) -getBF = byte 5 *> (BitField <$> (remaining >>= getLazyByteString . fromIntegral)) -getReq = byte 6 *> (Request <$> gw32 <*> (Block <$> gw32 <*> gw32)) -getPiece = byte 7 *> (Piece <$> gw32 <*> gw32 <*> (remaining >>= getByteString)) -getCancel = byte 8 *> (Cancel <$> gw32 <*> (Block <$> gw32 <*> gw32)) -getPort = byte 9 *> (Port . fromIntegral <$> getWord16be) +getRejectRequest, getAllowedFast, getSuggest, getHaveAll, getHaveNone :: Get Message +getExtendedMsg :: Get Message +getChoke = byte 0 *> return Choke +getUnchoke = byte 1 *> return Unchoke +getIntr = byte 2 *> return Interested +getNI = byte 3 *> return NotInterested +getHave = byte 4 *> (Have <$> gw32) +getBF = byte 5 *> (BitField <$> (remaining >>= getByteString . fromIntegral)) +getReq = byte 6 *> (Request <$> gw32 <*> (Block <$> gw32 <*> gw32)) +getPiece = byte 7 *> (Piece <$> gw32 <*> gw32 <*> (remaining >>= getByteString)) +getCancel = byte 8 *> (Cancel <$> gw32 <*> (Block <$> gw32 <*> gw32)) +getPort = byte 9 *> (Port . fromIntegral <$> getWord16be) +getSuggest = byte 0x0D *> (Suggest <$> gw32) +getHaveAll = byte 0x0E *> return HaveAll +getHaveNone = byte 0x0F *> return HaveNone +getRejectRequest = byte 0x10 *> (RejectRequest <$> gw32 <*> (Block <$> gw32 <*> gw32)) +getAllowedFast = byte 0x11 *> (AllowedFast <$> gw32) +getExtendedMsg = byte 20 *> (ExtendedMsg <$> gw32 <*> (remaining >>= getByteString)) getKA = do empty <- isEmpty if empty @@ -161,11 +269,18 @@ -- | Receive the header parts from the other end -receiveHeader :: Handle -> Int -> (InfoHash -> Bool) +receiveHeader :: Socket -> Int -> (InfoHash -> Bool) -> IO (Either String ([Capabilities], L.ByteString, InfoHash)) -receiveHeader h sz ihTst = parseHeader `fmap` B.hGet h sz +receiveHeader sock sz ihTst = parseHeader `fmap` loop [] sz where parseHeader = runGet (headerParser ihTst) + loop :: [B.ByteString] -> Int -> IO B.ByteString + loop bs z | z == 0 = return . B.concat $ reverse bs + | otherwise = do + nbs <- recv sock z + when (B.length nbs == 0) $ fail "Socket is dead" + loop (nbs : bs) (z - B.length nbs) + headerParser :: (InfoHash -> Bool) -> Get ([Capabilities], L.ByteString, InfoHash) headerParser ihTst = do hdSz <- getWord8 @@ -178,19 +293,18 @@ pid <- getLazyByteString 20 return (decodeCapabilities caps, pid, ihR) -data Capabilities decodeCapabilities :: Word64 -> [Capabilities] -decodeCapabilities _ = [] +decodeCapabilities w = catMaybes + [ if testBit w 2 then Just Fast else Nothing ] -- | Initiate a handshake on a socket -initiateHandshake :: Handle -> PeerId -> InfoHash +initiateHandshake :: Socket -> PeerId -> InfoHash -> IO (Either String ([Capabilities], L.ByteString, InfoHash)) -initiateHandshake handle peerid infohash = do +initiateHandshake sock peerid infohash = do debugM "Protocol.Wire" "Sending off handshake message" - L.hPut handle msg - hFlush handle + _ <- Lz.send sock msg debugM "Protocol.Wire" "Receiving handshake from other end" - receiveHeader handle sz (== infohash) + receiveHeader sock sz (== infohash) where msg = handShakeMessage peerid infohash sz = fromIntegral (L.length msg) @@ -202,17 +316,16 @@ putByteString . toBS $ pid] -- | Receive a handshake on a socket -receiveHandshake :: Handle -> PeerId -> (InfoHash -> Bool) +receiveHandshake :: Socket -> PeerId -> (InfoHash -> Bool) -> IO (Either String ([Capabilities], L.ByteString, InfoHash)) -receiveHandshake h pid ihTst = do +receiveHandshake s pid ihTst = do debugM "Protocol.Wire" "Receiving handshake from other end" - r <- receiveHeader h sz ihTst + r <- receiveHeader s sz ihTst case r of Left err -> return $ Left err Right (caps, rpid, ih) -> do debugM "Protocol.Wire" "Sending back handshake message" - L.hPut h (msg ih) - hFlush h + _ <- Lz.send s (msg ih) return $ Right (caps, rpid, ih) where msg ih = handShakeMessage pid ih sz = fromIntegral (L.length $ msg (B.pack $ replicate 20 32)) -- Dummy value @@ -220,8 +333,8 @@ -- | The call @constructBitField pieces@ will return the a ByteString suitable for inclusion in a -- BITFIELD message to a peer. -constructBitField :: Int -> [PieceNum] -> L.ByteString -constructBitField sz pieces = L.pack . build $ m +constructBitField :: Int -> [PieceNum] -> B.ByteString +constructBitField sz pieces = B.pack . build $ m where m = map (`elem` pieces) [0..sz-1 + pad] pad = 8 - (sz `mod` 8) build [] = [] diff --git a/src/RateCalc.hs b/src/RateCalc.hs --- a/src/RateCalc.hs +++ b/src/RateCalc.hs @@ -11,6 +11,7 @@ where +import Control.DeepSeq import Data.Time.Clock -- | A Rate is a record of information used for calculating the rate @@ -23,6 +24,10 @@ , rateSince :: !UTCTime -- ^ From where is the rate measured } +instance NFData Rate where + rnf (Rate r b c _ _ _) = + rnf r `seq` rnf b `seq` rnf c + fudge :: NominalDiffTime fudge = fromInteger 5 -- Seconds @@ -40,9 +45,9 @@ -- | The call @update n rt@ updates the rate structure @rt@ with @n@ new bytes update :: Integer -> Rate -> Rate -update n rt = nb `seq` nc `seq` rt { bytes = nb + n, count = nc + n} - where nb = bytes rt - nc = count rt +update n rt = rt { bytes = nb, count = nc} + where nb = bytes rt + n + nc = count rt + n -- | The call @extractRate t rt@ extracts the current rate from the rate structure and updates the rate @@ -56,21 +61,22 @@ n = bytes rt r = (rate rt * oldWindow + (fromIntegral n)) / newWindow expectN = min 5 (round $ (fromIntegral n / (max r 0.0001))) + nrt = rt { rate = r + , bytes = 0 + , nextExpected = addUTCTime (fromInteger expectN) t + , lastExt = t + , rateSince = max (rateSince rt) (addUTCTime (-maxRatePeriod) t) + } in -- Update the rate and book-keep the missing pieces. The total is simply a built-in -- counter. The point where we expect the next update is pushed at most 5 seconds ahead -- in time. But it might come earlier if the rate is high. -- Last is updated with the current time. Finally, we move the windows earliest value -- forward if it is more than 20 seconds from now. - (r, rt { rate = r - , bytes = 0 - , nextExpected = addUTCTime (fromInteger expectN) t - , lastExt = t - , rateSince = max (rateSince rt) (addUTCTime (-maxRatePeriod) t) - }) + nrt `deepseq` (r, nrt) -- | The call @extractCount rt@ extract the bytes transferred since last extraction extractCount :: Rate -> (Integer, Rate) -extractCount rt = crt `seq` (crt, rt { count = 0 }) +extractCount rt = (crt, rt { count = 0 }) where crt = count rt diff --git a/src/Test.hs b/src/Test.hs --- a/src/Test.hs +++ b/src/Test.hs @@ -12,7 +12,7 @@ import qualified Data.Queue (testSuite) import qualified Protocol.BCode (testSuite) import qualified Protocol.Wire (testSuite) - +import qualified Process.Peer (testSuite) runTests :: IO () runTests = @@ -23,6 +23,7 @@ , Data.PieceSet.testSuite , Protocol.BCode.testSuite , Protocol.Wire.testSuite + , Process.Peer.testSuite ] testSuite :: Test diff --git a/src/Torrent.hs b/src/Torrent.hs --- a/src/Torrent.hs +++ b/src/Torrent.hs @@ -14,6 +14,7 @@ , PieceInfo(..) , BlockSize , Block(..) + , Capabilities(..) -- * Interface , determineState , bytesLeft @@ -36,6 +37,7 @@ import qualified Data.Map as M import Network +import Network.Socket import Numeric import System.Random @@ -66,6 +68,12 @@ instance NFData TorrentState +---------------------------------------------------------------------- +-- Capabilities + +data Capabilities = Fast + deriving (Show, Eq) + -- PIECES ---------------------------------------------------------------------- type PieceNum = Int @@ -120,8 +128,8 @@ defaultOptimisticSlots = 2 -- | Default port to communicate on -defaultPort :: PortID -defaultPort = PortNumber $ fromInteger 1579 +defaultPort :: PortNumber +defaultPort = PortNum $ fromInteger 1579 -- | The current version of the Combinatorrent protocol string. This is bumped -- whenever we make a radical change to the protocol communication or fix a grave bug.