Combinatorrent (empty) → 0.2.0
raw patch · 38 files changed
+4959/−0 lines, 38 filesdep +HTTPdep +HUnitdep +PSQueuesetup-changed
Dependencies added: HTTP, HUnit, PSQueue, QuickCheck, array, base, bytestring, cereal, containers, directory, filepath, hopenssl, hslogger, mtl, network, parsec, pretty, random, random-shuffle, stm, test-framework, test-framework-hunit, test-framework-quickcheck2, time
Files
- AUTHORS +5/−0
- Combinatorrent.cabal +93/−0
- LICENSE +25/−0
- README.md +131/−0
- Setup.lhs +23/−0
- configure +21/−0
- src/Channels.hs +44/−0
- src/Combinatorrent.hs +147/−0
- src/Data/PendingSet.hs +67/−0
- src/Data/PieceSet.hs +198/−0
- src/Data/Queue.hs +132/−0
- src/Digest.hs +39/−0
- src/FS.hs +197/−0
- src/Process.hs +117/−0
- src/Process/ChokeMgr.hs +349/−0
- src/Process/Console.hs +86/−0
- src/Process/DirWatcher.hs +62/−0
- src/Process/FS.hs +73/−0
- src/Process/Listen.hs +32/−0
- src/Process/Peer.hs +422/−0
- src/Process/Peer/Receiver.hs +56/−0
- src/Process/Peer/Sender.hs +52/−0
- src/Process/Peer/SenderQ.hs +101/−0
- src/Process/PeerMgr.hs +217/−0
- src/Process/PieceMgr.hs +548/−0
- src/Process/Status.hs +156/−0
- src/Process/Timer.hs +20/−0
- src/Process/TorrentManager.hs +116/−0
- src/Process/Tracker.hs +278/−0
- src/Protocol/BCode.hs +372/−0
- src/Protocol/Wire.hs +263/−0
- src/RateCalc.hs +73/−0
- src/Supervisor.hs +184/−0
- src/Test.hs +36/−0
- src/TestInstance.hs +51/−0
- src/Torrent.hs +145/−0
- src/Tracer.hs +20/−0
- src/Version.hs.in +8/−0
+ AUTHORS view
@@ -0,0 +1,5 @@+Alex Mason <axman6@gmail.com>+Astro <astro@spaceboyz.net>+Jesper Louis Andersen <jesper.louis.andersen@gmail.com>+John Gunderman <johngunderman@gmail.com>+Thomas Christensen <thomas@chrstnsn.dk>
+ Combinatorrent.cabal view
@@ -0,0 +1,93 @@+name: Combinatorrent+category: Network+version: 0.2.0+category: Network+description: Combinatorrent provides a BitTorrent client, based on STM+ for concurrency. This is an early preview release which is capable of+ downloading files from various torrent trackers, but have not yet+ demonstrated to be correct in all aspects.++ It is expected that the package currently contains numerous and even+ grave bugs. Patches to fix any problem are welcome!+cabal-version: >= 1.6++license: BSD3+license-file: LICENSE+copyright: (c) 2009 Jesper Louis Andersen+author: Jesper Louis Andersen+maintainer: jesper.louis.andersen@gmail.com+stability: experimental+synopsis: A concurrent bittorrent client+build-type: Configure++extra-tmp-files: src/Version.hs+extra-source-files: src/Version.hs.in, configure+data-files: AUTHORS, README.md++flag debug+ description: Enable debug support+ default: False++flag threadscope+ description: Enable the eventlog necessary for ThreadScope+ default: False++executable Combinatorrent+ hs-source-dirs: src+ main-is: Combinatorrent.hs+ other-modules: Protocol.BCode, Protocol.Wire,+ Data.Queue, Data.PieceSet, Data.PendingSet+ Process.ChokeMgr, Process.Console, Process.FS, Process.Listen,+ Process.PeerMgr, Process.Peer, Process.PieceMgr, Process.Status,+ Process.Timer, Process.Tracker, Process.TorrentManager+ Digest, FS, Channels, Process, RateCalc,+ Supervisor, Torrent, Test, TestInstance, Process.DirWatcher,+ Tracer,+ Process.Peer.Sender,+ Process.Peer.SenderQ,+ Process.Peer.Receiver++ extensions: CPP++ build-depends:+ array >= 0.3,+ base >= 3.0,+ base <= 5.0,+ bytestring,+ cereal,+ containers,+ directory,+ filepath,+ hopenssl,+ hslogger,+ HTTP,+ HUnit,+ mtl,+ network,+ parsec,+ pretty,+ PSQueue,+ QuickCheck >= 2,+ random,+ random-shuffle,+ stm,+ test-framework,+ test-framework-hunit,+ test-framework-quickcheck2,+ time++ ghc-prof-options: -auto-all -ignore-scc -caf-all+ ghc-options: -Wall -fno-warn-orphans -threaded+ if !flag(debug)+ cpp-options: "-DNDEBUG"++ if flag(threadscope)+ ghc-options: -eventlog++source-repository head+ type: git+ location: git://github.com/jlouis/combinatorrent.git+ branch: master+++
+ LICENSE view
@@ -0,0 +1,25 @@+Copyright (c) 2009, 2010, Jesper Louis Andersen+All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are+met:++* Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++* Redistributions in binary form must reproduce the above copyright+ notice, this list of conditions and the following disclaimer in the+ documentation and/or other materials provided with the distribution.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,131 @@+Combinatorrent - a bittorrent client.+=====================================++Introduction+----------++This is a bittorrent client. I am the introduction document and I need to be+written by some generous soul!++Installation+------------++Here is what I do to install haskell torrrent locally on my machine:++ cabal install --prefix=$HOME --user++Since we are using the magnificient cabal, this is enough to install haskell+torrent in our $HOME/bin directory.++Usage+-----------------++Combinatorrent can currently only do one very simple thing. If you call it with++ Combinatorrent foo.torrent++then it will begin downloading the file in foo.torrent to the current+directory via the Bittorrent protocol.++Protocol support+----------------++Currently haskell-torrent supports the following BEPs (See the+[BEP Process](http://www.bittorrent.org/beps/bep_0000.html) document for an+explanation of these)++ - 003, 004, 020,++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++Debugging+---------++For debugging, jlouis tends to use the following:++ make conf build test++This builds Combinatorrent with the *Debug* flag set and also builds the+software with profiling by default so it is easy to hunt down performance+regressions. It also runs the internal test-suite for various values. There+are a couple of interesting targets in the top-level Makefile++Reading material for hacking Combinatorrent:+--------------------------------------------++ - [Protocol specification - BEP0003](http://www.bittorrent.org/beps/bep_0003.html):+ This is the original protocol specification, tracked into the BEP+ process. It is worth reading because it explains the general overview+ and the precision with which the original protocol was written down.++ - [Bittorrent Enhancement Process - BEP0000](http://www.bittorrent.org/beps/bep_0000.html)+ The BEP process is an official process for adding extensions on top of+ the BitTorrent protocol. It allows implementors to mix and match the+ extensions making sense for their client and it allows people to+ discuss extensions publicly in a forum. It also provisions for the+ deprecation of certain features in the long run as they prove to be of+ less value.++ - [wiki.theory.org](http://wiki.theory.org/Main_Page)+ An alternative description of the protocol. This description is in+ general much more detailed than the BEP structure. It is worth a read+ because it acts somewhat as a historic remark and a side channel. Note+ that there are some commentary on these pages which can be disputed+ quite a lot.++ - ["Supervisor Behaviour"](http://erlang.org/doc/design_principles/sup_princ.html)+ From the Erlang documentation. How the Erlang Supervisor behaviour+ works. The Supervisor and process structure of Combinatorrent is+ somewhat inspired by the Erlang ditto.++Source code Hierarchy+---------------------++ - **Data**: Data structures.+ - **Queue**: Functional queues. Standard variant with two lists.+ - **PendingSet**: A wrapper around Data.PSQueue for tracking how+ common a piece is.+ - **PieceSet**: BitArrays of pieces and their operations.++ - **Process**: Process definitions for the different processes comprising+ Combinatorrent+ - **ChokeMgr**: Manages choking and unchoking of peers, based upon the current speed of the peer+ and its current state. Global for multiple torrents.+ - **Console**: Simple console process. Only responds to 'quit' at the moment.+ - **DirWatcher**: Watches a directory and adds any torrent present in+ it.+ - **FS**: Process managing the file system.+ - **Listen**: Not used at the moment. Step towards listening sockets.+ - **Peer**: Several process definitions for handling peers. Two for sending, one for receiving+ and one for controlling the peer and handle the state.+ - **PeerMgr**: Management of a set of peers for a single torrent.+ - **PieceMgr**: Keeps track of what pieces have been downloaded and what are missing. Also hands+ out blocks for downloading to the peers.+ - **Status**: Keeps track of uploaded/downloaded/left bytes for a single torrent. Could be globalized.+ - **Timer**: Timer events.+ - **TorrentManager**: Manages torrents at the top-level.+ - **Tracker**: Communication with the tracker.++ - **Protocol**: Modules for interacting with the various bittorrent protocols.+ - **BCode**: The bittorrent BCode coding. Used by several protocols.+ - **Wire**: The protocol used for communication between peers.++ - **Top Level**:+ - **Channels**: Various Channel definitions.+ - **Combinatorrent**: Main entry point to the code. Sets up processes.+ - **Digest**: SHA1 digests as used in the bittorrent protocol.+ - **FS**: Low level Filesystem code. Interacts with files.+ - **Process**: Code for Erlang-inspired processes.+ - **RateCalc**: Rate calculations for a network socket. We use this to keep track of the+ current speed of a peer in one direction.+ - **Supervisor**: Erlang-inspired Supervisor processes.+ - **Test.hs**: Code for test-framework+ - **TestInstance.hs**: Various helper instances not present in the test framework by default+ - **Torrent**: Various helpers and types for Torrents.+ - **Tracer**: Code for simple "ring"-like tracing.+ - **Version.hs.in**: Generates **Version.hs** via the configure script.++# vim: filetype=none tw=76 expandtab
+ Setup.lhs view
@@ -0,0 +1,23 @@+#!/usr/bin/env runhaskell++In principle, we could do with a lot less than autoconfUserhooks, but simpleUserHooks+is not running 'configure'.++We will need Distribution.Simple to do the heavyweight lifting, and we will need some filePath magic.++> import System.FilePath+> import System.Process+> import Distribution.Simple+> import Distribution.Simple.LocalBuildInfo+> import Distribution.PackageDescription++The main program is just to make Cabal lift it. But we will override testing.++> main = defaultMainWithHooks hooks+> where hooks = autoconfUserHooks { runTests = runTests' }++Running tests is to call Combinatorrent with its parameters for tests:++> runTests' :: Args -> Bool -> PackageDescription -> LocalBuildInfo -> IO ()+> runTests' _ _ _ lbi = system testprog >> return ()+> where testprog = (buildDir lbi) </> "Combinatorrent" </> "Combinatorrent --tests"
+ configure view
@@ -0,0 +1,21 @@+#!/bin/sh+#+# Simple configuration script.++## Git version+GITHEAD="0.2 Hackage"+if test -d "${GIT_DIR:-.git}" ; then+ GITHEAD=`git describe 2>/dev/null`++ if test -z ${GITHEAD} ; then+ GITHEAD=`git rev-parse HEAD`+ fi++ if test -n "`git diff-index -m --name-only HEAD`" ; then+ GITHEAD="${GITHEAD}-dirty"+ fi+fi++sed -e "s/@GITHEAD@/$GITHEAD/g" src/Version.hs.in > src/Version.hs++
+ src/Channels.hs view
@@ -0,0 +1,44 @@+module Channels+ ( Peer(..)+ , PeerMessage(..)+ , PeerChannel+ , MgrMessage(..)+ , MgrChannel+ , BandwidthChannel+ , TrackerMsg(..)+ , TrackerChannel+ )+where++import Control.Concurrent+import Control.Concurrent.STM++import Network+import Torrent++data Peer = Peer { peerHost :: HostName,+ peerPort :: PortID }++data PeerMessage = ChokePeer+ | UnchokePeer+ | PieceCompleted PieceNum+ | CancelBlock PieceNum Block++type PeerChannel = TChan PeerMessage++---- TRACKER++-- | Messages to the tracker process+data TrackerMsg = Stop -- ^ Ask the Tracker to stop+ | TrackerTick Integer -- ^ Ticking in the tracker, used to contact again+ | Start -- ^ Ask the tracker to Start+ | Complete -- ^ Ask the tracker to Complete the torrent+type TrackerChannel = TChan TrackerMsg++data MgrMessage = Connect InfoHash ThreadId PeerChannel+ | Disconnect ThreadId++type MgrChannel = TChan MgrMessage++-- | A Channel type we use for transferring the amount of data we transferred+type BandwidthChannel = TChan Integer
+ src/Combinatorrent.hs view
@@ -0,0 +1,147 @@+module Main (main)+where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad+import Control.Monad.State++import Data.List++import System.Environment+import System.Random++import System.Console.GetOpt+import System.Directory (doesDirectoryExist)+import System.FilePath ()+import System.Log.Logger+import System.Log.Handler.Simple+import System.IO as SIO++import qualified Process.Console as Console+import qualified Process.PeerMgr as PeerMgr+import qualified Process.ChokeMgr as ChokeMgr (start)+import qualified Process.Listen as Listen+import qualified Process.DirWatcher as DirWatcher (start)+import qualified Process.Status as Status (start, StatusChannel, PStat)+import qualified Process.TorrentManager as TorrentManager (start, TorrentMgrChan, TorrentManagerMsg(..))++import Supervisor+import Torrent+import Version+import qualified Test++main :: IO ()+main = do args <- getArgs+ if "--tests" `elem` args+ then Test.runTests+ else progOpts args >>= run++-- COMMAND LINE PARSING++data Flag = Version | Debug | LogFile FilePath | WatchDir FilePath | StatFile FilePath+ deriving (Eq, Show)++options :: [OptDescr Flag]+options =+ [ Option ['V','?'] ["version"] (NoArg Version) "Show version number"+ , Option ['D'] ["debug"] (NoArg Debug) "Spew extra debug information"+ , Option [] ["logfile"] (ReqArg LogFile "FILE") "Choose a filepath on which to log"+ , Option ['W'] ["watchdir"] (ReqArg WatchDir "DIR") "Choose a directory to watch for torrents"+ , Option ['S'] ["statfile"] (ReqArg StatFile "FILE") "Choose a file to gather stats into"+ ]++progOpts :: [String] -> IO ([Flag], [String])+progOpts args = do+ case getOpt Permute options args of+ (o,n,[] ) -> return (o, n)+ (_,_,errs) -> ioError (userError (concat errs ++ usageInfo header options))+ where header = "Usage: Combinatorrent [OPTION...] file"++run :: ([Flag], [String]) -> IO ()+run (flags, files) = do+ if Version `elem` flags+ then progHeader+ else case files of+ [] -> putStrLn "No torrentfile input"+ names -> progHeader >> download flags names++progHeader :: IO ()+progHeader = putStrLn $ "This is Combinatorrent \x2620 version " ++ version ++ "\n" +++ " For help type 'help'\n"++setupLogging :: [Flag] -> IO ()+setupLogging flags = do+ fLog <- case logFlag flags of+ Nothing -> streamHandler SIO.stdout DEBUG+ Just (LogFile fp) -> fileHandler fp DEBUG+ Just _ -> error "Impossible match"+ when (Debug `elem` flags)+ (updateGlobalLogger rootLoggerName+ (setHandlers [fLog] . (setLevel DEBUG)))+ where logFlag = find (\e -> case e of+ LogFile _ -> True+ _ -> False)++setupDirWatching :: [Flag] -> TorrentManager.TorrentMgrChan -> IO [Child]+setupDirWatching flags watchC = do+ case dirWatchFlag flags of+ Nothing -> return []+ Just (WatchDir dir) -> do+ ex <- doesDirectoryExist dir+ if ex+ then do return [ Worker $ DirWatcher.start dir watchC ]+ else do putStrLn $ "Directory does not exist, not watching"+ return []+ Just _ -> error "Impossible match"+ where dirWatchFlag = find (\e -> case e of+ WatchDir _ -> True+ _ -> False)++setupStatus :: [Flag] -> Status.StatusChannel -> TVar [Status.PStat] -> Child+setupStatus flags statusC stv =+ case statFileFlag flags of+ Nothing -> Worker $ Status.start Nothing statusC stv+ Just (StatFile fn) -> Worker $ Status.start (Just fn) statusC stv+ Just _ -> error "Impossible match"+ where statFileFlag = find (\e -> case e of+ StatFile _ -> True+ _ -> False)++generatePeerId :: IO PeerId+generatePeerId = do+ gen <- getStdGen+ return $ mkPeerId gen++download :: [Flag] -> [String] -> IO ()+download flags names = do+ setupLogging flags+ watchC <- liftIO newTChanIO+ workersWatch <- setupDirWatching flags watchC+ -- setup channels+ statusC <- liftIO $ newTChanIO+ waitC <- liftIO $ newEmptyTMVarIO+ supC <- liftIO newTChanIO+ pmC <- liftIO $ newTChanIO+ chokeC <- liftIO $ newTChanIO+ rtv <- atomically $ newTVar []+ stv <- atomically $ newTVar []+ debugM "Main" "Created channels"+ pid <- generatePeerId+ tid <- allForOne "MainSup"+ (workersWatch +++ [ Worker $ Console.start waitC statusC+ , Worker $ TorrentManager.start watchC statusC stv chokeC pid pmC+ , setupStatus flags statusC stv+ , Worker $ PeerMgr.start pmC pid chokeC rtv+ , Worker $ ChokeMgr.start chokeC rtv 100 -- 100 is upload rate in KB+ , Worker $ Listen.start defaultPort pmC+ ]) supC+ atomically $ writeTChan watchC (map TorrentManager.AddedTorrent names)+ _ <- atomically $ takeTMVar waitC+ infoM "Main" "Closing down, giving processes 10 seconds to cool off"+ atomically $ writeTChan supC (PleaseDie tid)+ threadDelay $ 10*1000000+ infoM "Main" "Done..."+ return ()+
+ src/Data/PendingSet.hs view
@@ -0,0 +1,67 @@+module Data.PendingSet+ ( PendingSet+ , Data.PendingSet.empty+ , Data.PendingSet.size+ , have+ , unhave+ , haves+ , unhaves+ , pick+ )+where++import Data.PSQueue hiding (foldl)++import Torrent++-- | Representation of Pending Sets.+newtype PendingSet = PendingSet { unPS :: PSQ PieceNum Int }++-- | The empty pending set.+empty :: PendingSet+empty = PendingSet Data.PSQueue.empty++size :: PendingSet -> Int+size = Data.PSQueue.size . unPS++-- | A peer has a given piece. Reflect this in the PendingSet.+have :: PieceNum -> PendingSet -> PendingSet+have pn = PendingSet . alter f pn . unPS+ where f Nothing = Just 1+ f (Just x) = Just (x + 1)++-- | A Peer does not have a given piece anymore (TODO: Not used in practice)+unhave :: PieceNum -> PendingSet -> PendingSet+unhave pn = PendingSet . alter f pn . unPS+ where f Nothing = error "Data.PendingSet.unhave"+ f (Just 1) = Nothing+ f (Just x) = Just (x-1)++-- | Add all pieces in a bitfield+haves :: [PieceNum] -> PendingSet -> PendingSet+haves pns = flip (foldl f) pns+ where f e = flip have e++-- | Remove all pieces in a bitfield+unhaves :: [PieceNum] -> PendingSet -> PendingSet+unhaves pns = flip (foldl f) pns+ where f e = flip unhave e++-- | Crawl through the set of pending pieces in decreasing order of rarity.+-- 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 selector ps = findPri (minView . unPS $ ps)+ where findPri Nothing = Nothing+ findPri (Just (pn :-> p, rest)) =+ if selector pn+ then pickAtPri 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+
+ src/Data/PieceSet.hs view
@@ -0,0 +1,198 @@+-- | Module for the representation of PieceSets. Exist so we can abstract on the implementation later+module Data.PieceSet+ ( PieceSet+ , new+ , size+ , full+ , copy+ , delete+ , Data.PieceSet.null+ , insert+ , intersection+ , member+ , fromList+ , toList+ , Data.PieceSet.freeze+ -- * Tests+ , testSuite+ )+where++import Control.Applicative+import Control.Monad+import Control.Monad.Trans+import Data.Array.IO+import Data.Array.Unboxed ((!), UArray)+import qualified Data.Foldable as F+import Data.List ((\\), partition, sort, null)+import Prelude hiding (null)++import Test.Framework+import Test.Framework.Providers.HUnit+import Test.HUnit hiding (Path, Test)+import TestInstance() -- Pull arbitraries++import Torrent++newtype PieceSet = PieceSet { unPieceSet :: IOUArray Int Bool }++new :: MonadIO m => Int -> m PieceSet+new n = {-# SCC "Data.PieceSet/new" #-}+ liftIO $ PieceSet <$> newArray (0, n-1) False++all :: MonadIO m => (Bool -> Bool) -> PieceSet -> m Bool+all f ps = liftIO $ do+ elems <- getElems $ unPieceSet ps+ return $ Prelude.all f elems++null :: MonadIO m => PieceSet -> m Bool+null = Data.PieceSet.all (==False) ++full :: MonadIO m => PieceSet -> m Bool+full = {-# SCC "Data.PieceSet/full" #-} Data.PieceSet.all (==True)++insert :: MonadIO m => Int -> PieceSet -> m ()+insert n (PieceSet ps) = {-# SCC "Data.PieceSet/insert" #-}+ liftIO $ writeArray ps n True++copy :: MonadIO m => PieceSet -> m PieceSet+copy (PieceSet ps) = liftIO $ do+ newArr <- mapArray id ps+ return $ PieceSet $ newArr++size :: MonadIO m => PieceSet -> m Int+size (PieceSet arr) = {-# SCC "Data.PieceSet/size" #-}+ liftIO $ do+ (l, u) <- getBounds arr+ walk [l..u] 0+ where walk [] n = return n+ walk (x : xs) n = readArray arr x >>= \p -> if p then walk xs (n+1) else walk xs n++member :: MonadIO m => Int -> PieceSet -> m Bool+member n (PieceSet arr) = {-# SCC "Data.PieceSet/member" #-}+ liftIO $ readArray arr n++delete :: MonadIO m => Int -> PieceSet -> m ()+delete n (PieceSet arr) = {-# SCC "Data.PieceSet/delete" #-} liftIO $+ writeArray arr n False++intersection :: MonadIO m => PieceSet -> PieceSet -> m [Int]+intersection (PieceSet arr1) (PieceSet arr2) = liftIO $ do+ eqSize <- (==) <$> getBounds arr1 <*> getBounds arr2+ if not eqSize+ then error "Wrong intersection sizes"+ else do+ elems <- getAssocs arr1+ F.foldlM mem [] elems+ where+ mem ls (_, False) = return ls+ mem ls (i, True) = do+ m <- readArray arr2 i+ return $ if m then (i : ls) else ls+++fromList :: MonadIO m => Int -> [Int] -> m PieceSet+fromList n elems = {-# SCC "Data.PieceSet/fromList" #-} liftIO $ do+ nArr <- newArray (0, n-1) False+ mapM_ (flip (writeArray nArr) True) elems+ return $ PieceSet nArr++toList :: MonadIO m => PieceSet -> m [Int]+toList (PieceSet arr) = {-# SCC "Data.PieceSet/toList" #-} liftIO $ do+ elems <- getAssocs arr+ return [i | (i, e) <- elems, e == True]++freeze :: MonadIO m => PieceSet -> m (PieceNum -> Bool)+freeze (PieceSet ps) = do+ frozen <- liftIO $ (Data.Array.IO.freeze ps :: IO (UArray Int Bool))+ return $ (frozen !)++-- Tests++testSuite :: Test+testSuite = testGroup "Data/PieceSet"+ [ testCase "New/Size" testNewSize+ , testCase "Full" testFull+ , testCase "Build" testBuild+ , testCase "Full" testFull+ , testCase "Intersection" testIntersect+ , testCase "Membership" testMember+ , testCase "Insert/Delete" testInsertDelete+ , testCase "Copy" testCopy+ ]++testNewSize :: Assertion+testNewSize = do+ a <- new 1337+ sz <- size a+ assertEqual "For a new PieceSet" sz 0+ insert 3 a+ insert 5 a+ sz2 <- size a+ assertEqual "For inserted" sz2 2++testFull :: Assertion+testFull = do+ let maxElem = 1337+ ps <- new maxElem+ _ <- forM [0..maxElem-1] (flip insert ps)+ tst <- liftM and $ mapM (flip member ps) [0..maxElem-1]+ assertBool "for a full PieceSet" tst++testBuild :: Assertion+testBuild = do+ let nums = [0..1336]+ m = 1336 + 1+ ps <- fromList m nums+ sz <- size ps+ assertEqual "for size" sz (length nums)++testIntersect :: Assertion+testIntersect = do+ let (evens, odds) = partition (\x -> x `mod` 2 == 0) [0..99]+ evPS <- fromList 100 evens+ oddPS <- fromList 100 odds+ is1 <- intersection evPS oddPS+ assertBool "for intersection" (Data.List.null is1)+ ps1 <- fromList 10 [1,2,3,4,9]+ ps2 <- fromList 10 [0,2,5,4,8 ]+ is2 <- intersection ps1 ps2+ assertBool "for simple intersection" (sort is2 == [2,4])++testMember :: Assertion+testMember = do+ let evens = filter (\x -> x `mod` 2 == 0) [0..999]+ m = 1000+ notThere = [0..999] \\ evens+ ps <- fromList m evens+ a <- liftM and $ mapM (flip member ps) evens+ b <- liftM and $ mapM (liftM not . flip member ps) notThere+ assertBool "for members" a+ assertBool "for non-members" b++testInsertDelete :: Assertion+testInsertDelete = do+ ps <- new 10+ insert 3 ps+ insert 4 ps+ assertBool "Ins/del #1" =<< member 3 ps+ assertBool "Ins/del #2" =<< liftM not (member 5 ps)+ delete 3 ps+ assertBool "Ins/del #3" =<< member 4 ps+ assertBool "Ins/del #4" =<< liftM not (member 3 ps)+ insert 5 ps+ assertBool "Ins/del #5" =<< member 5 ps++testCopy :: Assertion+testCopy = do+ ps <- new 10+ insert 3 ps+ pc <- copy ps+ insert 4 pc+ delete 3 pc+ assertBool "#1" =<< member 3 ps+ assertBool "#2" =<< liftM not (member 3 pc)+ assertBool "#3" =<< member 4 pc+ assertBool "#4" =<< liftM not (member 4 ps)++
+ src/Data/Queue.hs view
@@ -0,0 +1,132 @@+-- | Simple Functional queues based on a double list. This usually achieves good amortized bounds+module Data.Queue (+ -- * Types+ Queue+ -- * Functions+ , empty+ , null+ , first+ , remove+ , push+ , pop+ , Data.Queue.filter+ -- * Test Suite+ , testSuite+ )+where++import Data.List as Lst hiding (null)+import Data.Maybe (fromJust)++import Prelude hiding (null)++import Test.QuickCheck+import Test.Framework+import Test.Framework.Providers.QuickCheck2+import Test.Framework.Providers.HUnit+import Test.HUnit hiding (Path, Test)++data Queue a = Queue [a] [a]+ deriving (Eq, Show)++-- | The definition of an empty Queue+empty :: Queue a+empty = Queue [] []++-- | Returns True on an empty Queue, and False otherwise.+null :: Queue a -> Bool+null (Queue [] []) = True+null _ = False++-- | Pushes a new element to the tail of the list.+-- Operates in constant time.+push :: a -> Queue a -> Queue a+push e (Queue front back) = Queue front (e : back)+++-- | Pops the top most element off the queue.+-- Operates in amortized constant time+pop :: Queue a -> Maybe (a, Queue a)+pop (Queue [] []) = Nothing+pop (Queue (e : es) back) = Just (e, Queue es back)+pop (Queue [] back) = pop (Queue (reverse back) [])++-- | Return the head of the queue, if any+first :: Queue a -> Maybe a+first (Queue [] []) = Nothing+first (Queue (e : _) _) = Just e+first (Queue [] back) = Just $ last back -- Yeah slow++-- | Kill the first element in the queue+remove :: Queue a -> Queue a+remove (Queue [] []) = Queue [] []+remove (Queue (_ : es) b) = Queue es b+remove (Queue [] b) = remove (Queue (reverse b) [])++-- | Generates a new Queue only containing elements for which+-- p returns true.+filter :: (a -> Bool) -> Queue a -> Queue a+filter p (Queue front back) = Queue (Lst.filter p front) (Lst.filter p back)+++-- Tests++testSuite :: Test+testSuite = testGroup "Data/Queue"+ [ testCase "Empty Queue is Empty" testEmptyIsEmpty+ , testCase "First/Remove" testFirstRemove+ , testProperty "Simple push/pop" testPushPopSimple+ , testProperty "push/pop more" testPushPopMore+ , testProperty "push/pop interleave" testPushPopInterleave+ ]++-- Rudimentary boring simple tests+testEmptyIsEmpty :: Assertion+testEmptyIsEmpty = do+ assertBool "for Empty Q" (null empty)+ assertBool "for non-Empty Q" (not $ null (push "Foo" empty))+ assertEqual "for popping the Empty Q" (pop $ snd . fromJust . pop $ push "Foo" empty) Nothing++testFirstRemove :: Assertion+testFirstRemove = do+ -- Should really cover this more+ let nq = push 2 (push (1 :: Integer) empty)+ assertEqual "first" (first nq) (Just 1)+ assertEqual "first/removed" (first (remove nq)) (Just 2)+ assertEqual "emptied" (first (remove (remove nq))) Nothing++testPushPopSimple :: String -> Bool+testPushPopSimple s =+ let nq = pop (push s empty)+ in case nq of+ Nothing -> False+ Just (r, q) -> r == s && null q++testPushPopMore :: [String] -> Bool+testPushPopMore ls =+ let nq = foldl (flip push) empty ls+ popAll = unfoldr pop nq+ in popAll == ls++data Operation = Push | Pop+ deriving (Eq, Show)++instance Arbitrary Operation where+ arbitrary = oneof [return Push, return Pop]++testPushPopInterleave :: [Operation] -> [String] -> Bool+testPushPopInterleave ops ls = testQ empty ops ls []+ where+ testQ q op lst res =+ case op of+ [] -> popAll q == reverse res+ Pop : r -> case pop q of+ Nothing -> testQ empty r lst []+ Just (e, nq) ->+ if (last res) == e+ then testQ nq r lst (init res)+ else False+ Push : r -> case lst of+ [] -> testQ q r lst res+ (e : es) -> testQ (push e q) r es (e : res)+ popAll = unfoldr pop
+ src/Digest.hs view
@@ -0,0 +1,39 @@+-- | Simple abstraction for message digests+module Digest+ ( Digest+ , digest+ )+where++import Data.Char+import Data.Word+import Control.Monad.State++import Foreign.Ptr+import qualified Data.ByteString as B+import Data.ByteString.Unsafe+import qualified Data.ByteString.Lazy as L+import qualified OpenSSL.Digest as SSL++-- Consider newtyping this+type Digest = String++digest :: L.ByteString -> IO Digest+digest bs = {-# SCC "sha1_digest" #-} do+ upack <- digestLBS SSL.SHA1 bs+ return $ map (chr . fromIntegral) upack++digestLBS :: SSL.MessageDigest -> L.ByteString -> IO [Word8]+digestLBS mdType xs = {-# SCC "sha1_digestLBS" #-}+ SSL.mkDigest mdType $ evalStateT (updateLBS xs >> SSL.final)++updateBS :: B.ByteString -> SSL.Digest ()+updateBS bs = {-# SCC "sha1_updateBS" #-} 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+ where chunked = {-# SCC "sha1_updateLBS_chunked" #-} L.toChunks lbs
+ src/FS.hs view
@@ -0,0 +1,197 @@+-- | Filesystem routines. These are used for working with and+-- manipulating files in the filesystem.+module FS (PieceInfo(..),+ PieceMap,+ Handles,+ readPiece,+ readBlock,+ writeBlock,+ mkPieceMap,+ checkFile,+ checkPiece,+ openAndCheckFile,+ canSeed)+where++import Control.Monad.State++import qualified Data.ByteString.Char8 as B+import qualified Data.ByteString.Lazy as L+import qualified Data.Map as M++import System.IO+import System.Directory (createDirectoryIfMissing)+import System.FilePath (joinPath)++import Protocol.BCode as BCode+import qualified Digest as D+import Torrent++-- | For multi-file torrents we've got to maintain multiple file+-- handles. The data structure may as well be a Map Range Handle,+-- but that's detailto only @projectHandles@. More importantly,+-- functions operating on the files must be aware that a+-- piece/block can span multiple files.+--+-- FIXME: Replace this with a handle cache later. Many peers & many+-- tiny files will make us overstep the fd limit (usually+-- 1024).+newtype Handles = Handles [(Handle, Integer)] -- ^[(fileHandle, fileLength)]++projectHandles :: Handles+ -> Integer -- ^Torrent offset+ -> Integer -- ^Torrent size+ -> [(Handle -- ^File handle+ ,Integer -- ^File chunk offset+ ,Integer -- ^File chunk size+ )]+{-+projectHandles handles offset size = let r = projectHandles' handles offset size+ in trace ("projectHandles " +++ show handles ++ " " +++ show offset ++ " " +++ show size ++ " = " +++ show r+ ) $+ r+-}+projectHandles (Handles handles@((h1, length1):handles')) offset 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+ | otherwise =+ let size1 = length1 - offset -- ^How much of h1 to take?+ in if size1 >= size+ then [(h1, offset, size)]+ else (h1, offset, size1) :+ projectHandles (Handles handles') 0 (size - size1)++pInfoLookup :: PieceNum -> PieceMap -> IO PieceInfo+pInfoLookup pn mp = case M.lookup pn mp of+ Nothing -> fail "FS: Error lookup in PieceMap"+ Just i -> return i++-- | FIXME: minor code duplication with @readBlock@+readPiece :: PieceNum -> Handles -> PieceMap -> IO L.ByteString+readPiece pn handles mp =+ {-# SCC "readPiece" #-}+ do pInfo <- pInfoLookup pn mp+ bs <- L.concat `fmap`+ forM (projectHandles handles (offset pInfo) (len pInfo))+ (\(h, offset, size) ->+ do hSeek h AbsoluteSeek offset+ L.hGet h (fromInteger size)+ )+ if L.length bs == (fromInteger . len $ pInfo)+ then return bs+ else fail "FS: Wrong number of bytes read"++-- | FIXME: concatenating strict ByteStrings may turn out+-- expensive. Returning lazy ones may be more appropriate.+readBlock :: PieceNum -> Block -> Handles -> PieceMap -> IO B.ByteString+readBlock pn blk handles mp =+ {-# SCC "readBlock" #-}+ do pInfo <- pInfoLookup pn mp+ B.concat `fmap`+ forM (projectHandles handles (offset pInfo + (fromIntegral $ blockOffset blk))+ (fromIntegral $ blockSize blk))+ (\(h, offset, size) ->+ do hSeek h AbsoluteSeek offset+ B.hGet h $ fromInteger size+ )++-- | The call @writeBlock h n blk pm blkData@ will write the contents of @blkData@+-- to the file pointed to by handle at the correct position in the file. If the+-- block is of a wrong length, the call will fail.+writeBlock :: Handles -> PieceNum -> Block -> PieceMap -> B.ByteString -> IO ()+writeBlock handles n blk pm blkData =+ {-# SCC "writeBlock" #-}+ do when lenFail $ fail "Writing block of wrong length"+ pInfo <- pInfoLookup n pm+ foldM_ (\blkData (h, offset, size) ->+ do let size' = fromInteger size+ hSeek h AbsoluteSeek offset+ B.hPut h $ B.take size' blkData+ hFlush h+ return $ B.drop size' blkData+ ) blkData (projectHandles handles (position pInfo) (fromIntegral $ B.length blkData))+ return ()+ where+ position :: PieceInfo -> Integer+ position pinfo = (offset pinfo) + fromIntegral (blockOffset blk)+ lenFail = B.length blkData /= blockSize blk++-- | The @checkPiece h inf@ checks the file system for correctness of a given piece, namely if+-- the piece described by @inf@ is correct inside the file pointed to by @h@.+checkPiece :: PieceInfo -> Handles -> IO Bool+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+ L.hGet h (fromInteger size)+ )+ dgs <- liftIO $ D.digest bs+ return (dgs == digest inf)++-- | Create a MissingMap from a file handle and a piecemap. The system will read each part of+-- the file and then check it against the digest. It will create a map of what we are missing+-- in the file as a missing map. We could alternatively choose a list of pieces missing rather+-- then creating the data structure here. This is perhaps better in the long run.+checkFile :: Handles -> PieceMap -> IO PiecesDoneMap+checkFile handles pm = do l <- mapM checkP pieces+ return $ M.fromList l+ where pieces = M.toAscList pm+ checkP :: (PieceNum, PieceInfo) -> IO (PieceNum, Bool)+ checkP (pn, pInfo) = do b <- checkPiece pInfo handles+ return (pn, b)++-- | Extract the PieceMap from a bcoded structure+-- Needs some more defense in the long run.+mkPieceMap :: BCode -> Maybe PieceMap+mkPieceMap bc = fetchData+ where fetchData = do pLen <- infoPieceLength bc+ pieceData <- infoPieces bc+ tLen <- infoLength bc+ let pm = M.fromList . zip [0..] . extract pLen tLen 0 $ pieceData+ when ( tLen /= (sum $ map len $ M.elems pm) )+ (error "PieceMap construction size assertion failed")+ return pm+ extract :: Integer -> Integer -> Integer -> [B.ByteString] -> [PieceInfo]+ extract _ 0 _ [] = []+ extract plen tlen offst (p : ps) | tlen < plen = PieceInfo { offset = offst,+ len = tlen,+ digest = B.unpack p } : extract plen 0 (offst + plen) ps+ | otherwise = inf : extract plen (tlen - plen) (offst + plen) ps+ where inf = PieceInfo { offset = offst,+ len = plen,+ digest = B.unpack p }+ extract _ _ _ _ = error "mkPieceMap: the impossible happened!"++-- | Predicate function. True if nothing is missing from the map.+canSeed :: PiecesDoneMap -> Bool+canSeed = M.fold (&&) True++-- | Process a BCoded torrent file. Create directories, open the files+-- in question, check it and return Handles plus a haveMap for the+-- file+openAndCheckFile :: BCode -> IO (Handles, PiecesDoneMap, PieceMap)+openAndCheckFile bc =+ do+ handles <- Handles `fmap`+ forM files+ (\(path, length) ->+ do let dir = joinPath $ init path+ when (dir /= "") $+ createDirectoryIfMissing True dir+ let fpath = joinPath path+ h <- openBinaryFile fpath ReadWriteMode+ return (h, length)+ )+ have <- checkFile handles pieceMap+ return (handles, have, pieceMap)+ where Just files = BCode.infoFiles bc+ Just pieceMap = mkPieceMap bc
+ src/Process.hs view
@@ -0,0 +1,117 @@+-- | Core Process code+{-# LANGUAGE ExistentialQuantification, FlexibleInstances,+ GeneralizedNewtypeDeriving,+ MultiParamTypeClasses, CPP #-}+-- required for deriving Typeable+{-# OPTIONS_GHC -fglasgow-exts #-}++module Process (+ -- * Types+ Process+ -- * Interface+ , runP+ , spawnP+ , catchP+ , cleanupP+ , foreverP+ , stopP+ -- * Log Interface+ , Logging(..)+ , logP+ , infoP+ , debugP+ , warningP+ , criticalP+ , errorP+ )+where++import Control.Concurrent+import Control.Exception++import Control.Monad.Reader+import Control.Monad.State++import Data.Typeable++import Prelude hiding (catch, log)++import System.Log.Logger++-- | A @Process a b c@ is the type of processes with access to configuration data @a@, state @b@+-- returning values of type @c@. Usually, the read-only data are configuration parameters and+-- channels, and the state the internal process state. It is implemented by means of a transformer+-- stack on top of IO.+newtype Process a b c = Process (ReaderT a (StateT b IO) c)+#ifndef __HADDOCK__+ deriving (Functor, Monad, MonadIO, MonadState b, MonadReader a, Typeable)+#endif++data StopException = StopException+ deriving (Show, Typeable)++instance Exception StopException++stopP :: Process a b c+stopP = throw StopException++-- | Run the process monad given a configuation of type @a@ and a initial state of type @b@+runP :: a -> b -> Process a b c -> IO (c, b)+runP c st (Process p) = runStateT (runReaderT p c) st++-- | Spawn and run a process monad+spawnP :: a -> b -> Process a b () -> IO ThreadId+spawnP c st p = forkIO proc+ where proc = do _ <- runP c st p+ return ()++-- | Run the process monad for its side effect, with a stopHandler if exceptions+-- are raised in the process+catchP :: Logging a => Process a b () -> Process a b () -> Process a b ()+catchP proc stopH = cleanupP proc stopH (return ())++-- | Run the process monad for its side effect. @cleanupP p sh ch@ describes to+-- run @p@. If @p@ dies by a kill from a supervisor, run @ch@. Otherwise it runs+-- @ch >> sh@ on death.+cleanupP :: Logging a => Process a b () -> Process a b () -> Process a b () -> Process a b ()+cleanupP proc stopH cleanupH = do+ st <- get+ c <- ask+ (a, s') <- liftIO $ runP c st proc `catches`+ [ Handler (\ThreadKilled ->+ runP c st ( do infoP $ "Process Terminated by Supervisor"+ cleanupH ))+ , Handler (\StopException ->+ runP c st (do infoP $ "Process Terminating gracefully"+ cleanupH >> stopH)) -- This one is ok+ , Handler (\(ex :: SomeException) ->+ runP c st (do criticalP $ "Process exiting due to ex: " ++ show ex+ cleanupH >> stopH))+ ]+ put s'+ return a++-- | Run a process forever in a loop+foreverP :: Process a b c -> Process a b c+foreverP p = p >> foreverP p++------ LOGGING++--+-- | The class of types where we have a logger inside them somewhere+class Logging a where+ -- | Returns a channel for logging and an Identifying string to use+ logName :: a -> String++logP :: Logging a => Priority -> String -> Process a b ()+logP prio msg = do+ n <- asks logName+ liftIO $ logM n prio (n ++ ":\t" ++ msg)++infoP, debugP, criticalP, warningP, errorP :: Logging a => String -> Process a b ()+infoP = logP INFO+debugP = logP DEBUG+criticalP = logP CRITICAL+warningP = logP WARNING+errorP = logP ERROR+
+ src/Process/ChokeMgr.hs view
@@ -0,0 +1,349 @@+{-# LANGUAGE FlexibleContexts, TupleSections #-}+module Process.ChokeMgr (+ -- * Types, Channels+ ChokeMgrChannel+ , RateTVar+ , ChokeMgrMsg(..)+ -- * Interface+ , start+ )+where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Exception (assert)+import Control.Monad.Reader+import Control.Monad.State++import Data.Function+import Data.List+import qualified Data.Map as M+import qualified Data.Set as S+import Data.Traversable as T++import Prelude hiding (catch, log)++import System.Random++import Channels hiding (Peer)+import Process+import Process.Timer+import Supervisor+import Torrent hiding (infoHash)++-- DATA STRUCTURES+----------------------------------------------------------------------++-- | Messages to the Choke Manager+data ChokeMgrMsg = Tick -- ^ Request that we run another round+ | RemovePeer ThreadId -- ^ Request that this peer is removed+ | AddPeer InfoHash ThreadId PeerChannel -- ^ Request that this peer is added+ | PieceDone InfoHash PieceNum -- ^ Note that a given piece is done+ | BlockComplete InfoHash PieceNum Block -- ^ Note that a block is complete (endgame)+ | TorrentComplete InfoHash -- ^ Note that the torrent in question is complete++type ChokeMgrChannel = TChan ChokeMgrMsg+type RateTVar = TVar [(ThreadId, (Double, Double, Bool, Bool, Bool))]++data CF = CF { mgrCh :: ChokeMgrChannel+ , rateTV :: RateTVar }++instance Logging CF where+ logName _ = "Process.ChokeMgr"++-- PeerDB described below+type ChokeMgrProcess a = Process CF PeerDB a++-- INTERFACE+----------------------------------------------------------------------++roundTickSecs :: Int+roundTickSecs = 11++start :: ChokeMgrChannel -> RateTVar -> Int -> SupervisorChan+ -> IO ThreadId+start ch rtv ur supC = do+ _ <- registerSTM roundTickSecs ch Tick+ spawnP (CF ch rtv) (initPeerDB $ calcUploadSlots ur Nothing)+ (catchP (forever pgm)+ (defaultStopHandler supC))+ where+ initPeerDB slots = PeerDB 2 slots S.empty M.empty []+ pgm = do+ msg <- liftIO . atomically $ readTChan ch+ case msg of+ Tick -> tick+ RemovePeer t -> removePeer t+ AddPeer ih t pCh -> do+ debugP $ "Adding peer " ++ show (ih, t)+ addPeer pCh ih t+ BlockComplete ih pn blk -> informBlockComplete ih pn blk+ PieceDone ih pn -> informDone ih pn+ TorrentComplete ih -> modify (\s -> s { seeding = S.insert ih $ seeding s })+ tick = do debugP "Ticked"+ c <- asks mgrCh+ _ <- registerSTM roundTickSecs c Tick+ updateDB+ runRechokeRound+ removePeer tid = do debugP $ "Removing peer " ++ show tid+ modify (\db -> db { chain = filter (not . isPeer tid) (chain db)+ , rateMap = M.delete tid (rateMap db) })+ isPeer tid pr | tid == pThreadId pr = True+ | otherwise = False++-- INTERNAL FUNCTIONS+----------------------------------------------------------------------++-- The data structure is split into pieces so it is easier to manipulate.+-- The PeerDB is the state we thread around in the process. The PChain contains all+-- the important information about processes.+type PChain = [Peer]++-- | Main data for a peer+data Peer = Peer+ { pThreadId :: ThreadId+ , pInfoHash :: InfoHash+ , pChannel :: PeerChannel+ }++-- | Peer upload and download ratio+data PRate = PRate { pUpRate :: Double,+ pDownRate :: Double }+-- | 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+ }++type RateMap = M.Map ThreadId (PRate, PState)+data PeerDB = PeerDB+ { chokeRound :: Int -- ^ Counted down by one from 2. If 0 then we should+ -- advance the peer chain. (Optimistic Unchoking)+ , uploadSlots :: Int -- ^ Current number of upload slots+ , seeding :: S.Set InfoHash -- ^ Set of torrents we seed+ , rateMap :: RateMap -- ^ Map from Peer ThreadIds to state+ , chain :: PChain -- ^ The order in which peers are optimistically unchoked+ }++-- | Update the Peer Database with the newest information from peers+updateDB :: ChokeMgrProcess ()+updateDB = do+ rc <- asks rateTV+ rateUpdate <- liftIO . atomically $ do+ q <- readTVar rc+ writeTVar rc []+ 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+ nm m = foldl f m $ reverse updates+ in do+ debugP $ "Rate updates since last round: " ++ show updates+ modify (\db -> db { rateMap = nm (rateMap db) })++addPeer :: PeerChannel -> InfoHash -> ThreadId -> ChokeMgrProcess ()+addPeer ch ih t = do+ chn <- gets chain+ pt <- liftIO $ getStdRandom (\gen -> randomR (0, length chn - 1) gen)+ let (front, back) = splitAt pt chn+ modify (\db -> db { chain = (front ++ initPeer : back) })+ where initPeer = Peer t ih ch++runRechokeRound :: ChokeMgrProcess ()+runRechokeRound = do+ cRound <- gets chokeRound+ if (cRound == 0)+ then do advancePeerChain+ modify (\db -> db { chokeRound = 2 })+ else modify (\db -> db { chokeRound = (chokeRound db) - 1 })+ rechoke++-- | Advance the peer chain to the next peer eligible for optimistic+-- unchoking. That is, skip peers which are not interested in our pieces+-- and peers which are not choking us. The former we can't send any data to,+-- so we can't get better speeds at them. The latter are already sending us data,+-- so we know how good they are as peers.+advancePeerChain :: ChokeMgrProcess ()+advancePeerChain = do+ peers <- gets chain+ rd <- gets rateMap+ let (front, back) = break (breakPoint rd) peers+ modify (\db -> db { chain = back ++ front })+ where+ breakPoint rd peer =+ case M.lookup (pThreadId peer) rd of+ Nothing -> True -- Really new peer, give it the chance :)+ Just (_, st) -> pInterestedInUs st && pChokingUs st++rechoke :: ChokeMgrProcess ()+rechoke = do+ us <- gets uploadSlots+ chn <- gets chain+ sd <- gets seeding+ rm <- gets rateMap+ debugP $ "Chain is: " ++ show (map pThreadId chn)+ let (seed, down) = splitSeedLeech sd rm chn+ electedPeers <- selectPeers us down seed+ performChokingUnchoking electedPeers chn++-- | Function to split peers into those where we are seeding and those where we are leeching.+-- also prunes the list for peers which are not interesting.+-- TODO: Snubbed peers+splitSeedLeech :: S.Set InfoHash -> RateMap -> [Peer] -> ([Peer], [Peer])+splitSeedLeech seeders rm ps = partition (\p -> S.member (pInfoHash p) seeders) $ filter picker ps+ where+ -- TODO: pIsASeeder is always false at the moment+ picker :: Peer -> Bool+ picker p = case M.lookup (pThreadId p) rm of+ Nothing -> False -- Don't know anything about the peer yet+ Just (_, st) -> not (pIsASeeder st) && (pInterestedInUs st)++-- | Comparison with inverse ordering+compareInv :: Ord a => a -> a -> Ordering+compareInv x y =+ case compare x y of+ LT -> GT+ EQ -> EQ+ GT -> LT++comparingWith :: Ord a => (a -> a -> Ordering) -> (b -> a) -> b -> b -> Ordering+comparingWith comp project x y =+ comp (project x) (project y)++-- | Leechers are sorted by their current download rate. We want to keep fast peers around.+sortLeech :: [(Peer, (PRate, PState))] -> [(Peer, (PRate, PState))]+sortLeech = sortBy $ comparingWith compareInv (pDownRate . fst . snd)++-- | Seeders are sorted by their current upload rate.+sortSeeds :: [(Peer, (PRate, PState))] -> [(Peer, (PRate, PState))]+sortSeeds = sortBy $ comparingWith compareInv (pUpRate . fst . snd)+++-- | Calculate the amount of upload slots we have available. If the+-- number of slots is explicitly given, use that. Otherwise we+-- choose the slots based the current upload rate set. The faster+-- the rate, the more slots we allow.+calcUploadSlots :: Int -> Maybe Int -> Int+calcUploadSlots _ (Just n) = n+calcUploadSlots rate Nothing | rate <= 0 = 7 -- This is just a guess+ | rate < 9 = 2+ | rate < 15 = 3+ | rate < 42 = 4+ | otherwise = calcRate $ fromIntegral rate+ where calcRate :: Double -> Int+ calcRate x = round $ sqrt (x * 0.6)++-- | The call @assignUploadSlots c ds ss@ will assume that we have @c@+-- slots for uploading at our disposal. The list @ds@ will be peers+-- that we would like to upload to among the torrents we are+-- currently downloading. The list @ss@ is the same thing but for+-- torrents that we seed. The function returns a pair @(kd,ks)@+-- where @kd@ is the number of downloader slots and @ks@ is the+-- number of seeder slots.+--+-- The function will move surplus slots around so all of them gets used.+assignUploadSlots :: Int -> Int -> Int -> (Int, Int)+assignUploadSlots slots numDownPeers numSeedPeers =+ -- Shuffle surplus slots around so all gets used+ shuffleSeeders . shuffleDownloaders $ (downloaderSlots, seederSlots)+ where+ -- Calculate the slots available for the downloaders and seeders+ -- We allocate 70% of them to leeching and 30% of the to seeding+ -- though we assign at least one slot to both+ slotRound :: Double -> Double -> Int+ slotRound ss fraction = max 1 $ round $ ss * fraction++ downloaderSlots = slotRound (fromIntegral slots) 0.7+ seederSlots = slotRound (fromIntegral slots) 0.3++ -- If there is a surplus of downloader slots, then assign them to+ -- the seeder slots+ shuffleDownloaders (dSlots, sSlots) =+ case max 0 (dSlots - numDownPeers) of+ 0 -> (dSlots, sSlots)+ k -> (dSlots - k, sSlots + k)++ -- If there is a surplus of seeder slots, then assign these to+ -- the downloader slots. Limit the downloader slots to the number+ -- of downloaders, however+ shuffleSeeders (dSlots, sSlots) =+ case max 0 (sSlots - numSeedPeers) of+ 0 -> (dSlots, sSlots)+ k -> (min (dSlots + k) numDownPeers, sSlots - k)++-- | @selectPeers upSlots d s@ selects peers from a list of downloader peers @d@ and a list of seeder+-- peers @s@. The value of @upSlots@ defines the number of upload slots available+selectPeers :: Int -> [Peer] -> [Peer] -> ChokeMgrProcess (S.Set ThreadId)+selectPeers ups downPeers seedPeers = do+ rm <- gets rateMap+ let selector p = maybe (p, (PRate 0.0 0.0, PState True False False)) (p,)+ (M.lookup (pThreadId p) rm)+ dp = map selector downPeers+ sp = map selector seedPeers+ (nDownSlots, nSeedSlots) = assignUploadSlots ups (length downPeers) (length seedPeers)+ downPids = S.fromList $ map (pThreadId . fst) $ take nDownSlots $ sortLeech dp+ seedPids = S.fromList $ map (pThreadId . fst) $ take nSeedSlots $ sortSeeds sp+ debugP $ "Leechers: " ++ show (length downPeers) ++ ", Seeders: " ++ show (length seedPeers)+ debugP $ "Slots: " ++ show nDownSlots ++ " downloads, " ++ show nSeedSlots ++ " seeders"+ debugP $ "Electing peers - leechers: " ++ show downPids ++ "; seeders: " ++ show seedPids+ return $ assertSlots (nDownSlots + nSeedSlots) (S.union downPids seedPids)+ where assertSlots slots = assert (ups >= slots)++-- | Send a message to the peer process at PeerChannel. Message is sent asynchronously+-- to the peer in question. If the system is really loaded, this might+-- actually fail since the order in which messages arrive might be inverted.+msgPeer :: PeerChannel -> PeerMessage -> ChokeMgrProcess ()+msgPeer ch = liftIO . atomically . writeTChan ch++-- | This function performs the choking and unchoking of peers in a round.+performChokingUnchoking :: S.Set ThreadId -> [Peer] -> ChokeMgrProcess ()+performChokingUnchoking elected peers =+ do _ <- T.mapM unchoke electedPeers+ rm <- gets rateMap+ optChoke rm defaultOptimisticSlots nonElectedPeers+ 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++ -- 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+ -- not interested (ensuring fast start should they become interested); or+ -- they will be choked to avoid TCP/IP congestion.+ optChoke _rm _ [] = return ()+ optChoke rm 0 (p : ps) =+ case M.lookup (pThreadId p) rm of+ Nothing -> choke p >> optChoke rm 0 ps+ Just (_, st) ->+ if pInterestedInUs st+ then choke p >> optChoke rm 0 ps+ else unchoke p >> optChoke rm 0 ps+ optChoke rm k (p : ps) =+ case M.lookup (pThreadId p) rm of+ Nothing -> unchoke p >> optChoke rm (k-1) ps+ Just (_, st) ->+ if pInterestedInUs st+ then unchoke p >> optChoke rm (k-1) ps+ else unchoke p >> optChoke rm k ps++informDone :: InfoHash -> PieceNum -> ChokeMgrProcess ()+informDone ih pn = do+ chn <- gets chain+ T.mapM inform chn >> return ()+ where inform p | (pInfoHash p) == ih = sendDone p >> return ()+ | otherwise = return ()+ sendDone p = msgPeer (pChannel p) (PieceCompleted pn)++informBlockComplete :: InfoHash -> PieceNum -> Block -> ChokeMgrProcess ()+informBlockComplete ih pn blk = do+ chn <- gets chain+ T.mapM inform chn >> return ()+ where inform p | (pInfoHash p) == ih = sendComp p >> return ()+ | otherwise = return ()+ sendComp p = msgPeer (pChannel p) (CancelBlock pn blk)+
+ src/Process/Console.hs view
@@ -0,0 +1,86 @@+-- | The Console process has two main purposes. It is a telnet-like+-- interface with the user and it is our first simple logging device+-- for what happens inside the system.+{-# LANGUAGE ScopedTypeVariables #-}+module Process.Console+ ( start+ )+where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad.Reader++import Prelude hiding (catch)+++import Process+import qualified Process.Status as St+import Supervisor+++data Cmd = Quit -- ^ Quit the program+ | Show -- ^ Show current state+ | Help -- ^ Print Help message+ | Unknown String -- ^ Unknown command+ deriving (Eq, Show)++type CmdChannel = TChan Cmd++data CF = CF { cmdCh :: CmdChannel+ , wrtCh :: TChan String }++instance Logging CF where+ logName _ = "Process.Console"++-- | Start the logging process and return a channel to it. Sending on this+-- Channel means writing stuff out on stdOut+start :: TMVar () -> St.StatusChannel -> SupervisorChan -> IO ThreadId+start waitC statusC supC = do+ cmdC <- readerP -- We shouldn't be doing this in the long run+ wrtC <- writerP+ spawnP (CF cmdC wrtC) () (catchP (forever lp) (defaultStopHandler supC))+ where+ lp = do+ c <- asks cmdCh+ o <- asks wrtCh+ m <- liftIO . atomically $ readTChan c+ case m of+ Quit -> liftIO . atomically $ putTMVar waitC ()+ Help -> liftIO . atomically $ writeTChan o helpMessage+ (Unknown n) -> liftIO . atomically $ writeTChan o $ "Uknown command: " ++ n+ Show -> do+ v <- liftIO newEmptyTMVarIO+ liftIO . atomically $ writeTChan statusC (St.RequestAllTorrents v)+ sts <- liftIO . atomically $ takeTMVar v+ liftIO . atomically $ writeTChan o (show sts)++helpMessage :: String+helpMessage = concat+ [ "Command Help:\n"+ , "\n"+ , " help - Show this help\n"+ , " quit - Quit the program\n"+ , " show - Show the current downloading status\n"+ ]++writerP :: IO (TChan String)+writerP = do wrt <- newTChanIO+ _ <- forkIO $ lp wrt+ return wrt+ where lp wCh = forever (do m <- atomically $ readTChan wCh+ putStrLn m)++readerP :: IO CmdChannel+readerP = do cmd <- newTChanIO+ _ <- forkIO $ lp cmd+ return cmd+ where lp cmd = forever $+ do l <- getLine+ atomically $ writeTChan cmd+ (case l of+ "help" -> Help+ "quit" -> Quit+ "show" -> Show+ c -> Unknown c)+
+ src/Process/DirWatcher.hs view
@@ -0,0 +1,62 @@+-- | The DirWatcher Process runs a watcher over a directory. It will tell about any change+-- happening inside that directory.+module Process.DirWatcher (+ -- * Interface+ start+ )+where++import Control.Concurrent+import Control.Concurrent.STM++import Control.Monad.Reader+import Control.Monad.State++import qualified Data.Set as S++import System.Directory+import System.FilePath++import Prelude hiding (log)+import Process+import Process.TorrentManager hiding (start)+import Supervisor+++data CF = CF { reportCh :: TorrentMgrChan -- ^ Channel for reporting directory changes+ , dirToWatch :: FilePath }++type ST = S.Set FilePath++instance Logging CF where+ logName _ = "Process.DirWatcher"++start :: FilePath -- ^ Path to watch+ -> TorrentMgrChan -- ^ Channel to return answers on+ -> SupervisorChan+ -> IO ThreadId+start fp chan supC = do+ spawnP (CF chan fp) S.empty+ (catchP (foreverP pgm) (defaultStopHandler supC))+ where pgm = do+ q <- liftIO $ registerDelay (5 * 1000000)+ liftIO . atomically $ do+ b <- readTVar q+ if b then return () else retry+ processDirectory++processDirectory :: Process CF ST ()+processDirectory = do+ watchDir <- asks dirToWatch+ files <- liftIO $ getDirectoryContents watchDir+ let torrents = S.fromList $ filter (\fp -> (== ".torrent") $ snd . splitExtension $ fp) files+ running <- get+ let (added, removed) = (S.toList $ S.difference torrents running,+ S.toList $ S.difference running torrents)+ msg = (map AddedTorrent added ++ map RemovedTorrent removed)+ when (msg /= [])+ (do rc <- asks reportCh+ liftIO . atomically $ writeTChan rc msg+ -- Make ready for next iteration+ put torrents)+
+ src/Process/FS.hs view
@@ -0,0 +1,73 @@+-- | File system process. Acts as a maintainer for the filesystem in+-- question and can only do single-file torrents. It should be+-- fairly easy to add Multi-file torrents by altering this file and+-- the FS module.+{-# LANGUAGE ScopedTypeVariables #-}+module Process.FS+ ( FSPChannel+ , FSPMsg(..)+ , start+ )+where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad.Reader+import Control.Monad.State++import qualified Data.ByteString as B+import qualified Data.Map as M++import Process+import Torrent+import qualified FS+import Supervisor++data FSPMsg = CheckPiece PieceNum (TMVar (Maybe Bool))+ | WriteBlock PieceNum Block B.ByteString+ | ReadBlock PieceNum Block (TMVar B.ByteString)++type FSPChannel = TChan FSPMsg++data CF = CF+ { fspCh :: FSPChannel -- ^ Channel on which to receive messages+ }++instance Logging CF where+ logName _ = "Process.FS"++data ST = ST+ { fileHandles :: !FS.Handles -- ^ The file we are working on+ , pieceMap :: FS.PieceMap -- ^ Map of where the pieces reside+ }+++-- INTERFACE+----------------------------------------------------------------------++start :: FS.Handles -> FS.PieceMap -> FSPChannel -> SupervisorChan-> IO ThreadId+start handles pm fspC supC =+ spawnP (CF fspC) (ST handles pm) (catchP (forever lp) (defaultStopHandler supC))+ where+ lp = do+ c <- asks fspCh+ msg <- liftIO . atomically $ readTChan c+ case msg of+ CheckPiece n v -> do+ pmap <- gets pieceMap+ case M.lookup n pmap of+ Nothing -> liftIO . atomically $ putTMVar v Nothing+ Just p -> do r <- gets fileHandles >>= (liftIO . FS.checkPiece p)+ liftIO . atomically $ putTMVar v (Just r)+ ReadBlock n blk v -> do+ debugP $ "Reading block #" ++ show n+ ++ "(" ++ show (blockOffset blk) ++ ", " ++ show (blockSize blk) ++ ")"+ -- TODO: Protection, either here or in the Peer code+ h <- gets fileHandles+ bs <- gets pieceMap >>= (liftIO . FS.readBlock n blk h)+ liftIO . atomically $ putTMVar v bs+ WriteBlock pn blk bs -> {-# SCC "FS_WriteBlock" #-} do+ -- TODO: Protection, either here or in the Peer code+ fh <- gets fileHandles+ pmap <- gets pieceMap+ liftIO $ FS.writeBlock fh pn blk pmap bs
+ src/Process/Listen.hs view
@@ -0,0 +1,32 @@+module Process.Listen+ ( start+ )+where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad+import Control.Monad.Reader++import Network++import Process+import Process.PeerMgr hiding (start)+import Supervisor++data CF = CF { peerMgrCh :: PeerMgrChannel }++instance Logging CF where+ logName _ = "Process.Listen"++start :: PortID -> PeerMgrChannel -> SupervisorChan -> IO ThreadId+start port peerMgrC supC = do+ spawnP (CF peerMgrC) () (catchP (openListen >>= pgm)+ (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)+
+ src/Process/Peer.hs view
@@ -0,0 +1,422 @@+-- | Peer proceeses+{-# LANGUAGE ScopedTypeVariables #-}+module Process.Peer (+ -- * Types+ PeerMessage(..)+ -- * Interface+ , Process.Peer.start+ )+where++import Control.Applicative+import Control.Concurrent+import Control.Concurrent.STM++import Control.Monad.State+import Control.Monad.Reader++import Prelude hiding (catch, log)++import Data.Bits+import qualified Data.ByteString as B+import qualified Data.ByteString.Lazy as L++import qualified Data.Map as M+import qualified Data.PieceSet as PS+import Data.Maybe++import Data.Set as S hiding (map)+import Data.Time.Clock+import Data.Word++import System.IO++import Channels+import Process+import Process.FS+import Process.PieceMgr+import RateCalc as RC+import Process.Status+import Process.ChokeMgr (RateTVar)+import Process.Timer+import Supervisor+import Torrent+import Protocol.Wire++import qualified Process.Peer.Sender as Sender+import qualified Process.Peer.SenderQ as SenderQ+import qualified Process.Peer.Receiver as Receiver++-- INTERFACE+----------------------------------------------------------------------++start :: Handle -> MgrChannel -> RateTVar -> PieceMgrChannel+ -> FSPChannel -> TVar [PStat] -> PieceMap -> Int -> InfoHash+ -> IO Children+start handle pMgrC rtv pieceMgrC fsC stv pm nPieces ih = do+ queueC <- newTChanIO+ senderMV <- newEmptyTMVarIO+ receiverC <- newTChanIO+ sendBWC <- newTChanIO+ return [Worker $ Sender.start handle senderMV,+ Worker $ SenderQ.start queueC senderMV sendBWC,+ Worker $ Receiver.start handle receiverC,+ Worker $ peerP pMgrC rtv pieceMgrC fsC pm nPieces+ queueC receiverC sendBWC stv ih]++-- INTERNAL FUNCTIONS+----------------------------------------------------------------------++data PCF = PCF { inCh :: TChan (Message, Integer)+ , outCh :: TChan SenderQ.SenderQMsg+ , peerMgrCh :: MgrChannel+ , pieceMgrCh :: PieceMgrChannel+ , fsCh :: FSPChannel+ , peerCh :: PeerChannel+ , sendBWCh :: BandwidthChannel+ , timerCh :: TChan ()+ , statTV :: TVar [PStat]+ , rateTV :: RateTVar+ , pcInfoHash :: InfoHash+ , pieceMap :: PieceMap+ }++instance Logging PCF where+ logName _ = "Process.Peer"++data PST = PST { weChoke :: !Bool -- ^ True if we are choking the peer+ , weInterested :: !Bool -- ^ True if we are interested in the peer+ , blockQueue :: !(S.Set (PieceNum, Block)) -- ^ Blocks queued at the peer+ , peerChoke :: !Bool -- ^ Is the peer choking us? True if yes+ , peerInterested :: !Bool -- ^ True if the peer is interested+ , peerPieces :: !(PS.PieceSet) -- ^ List of pieces the peer has access to+ , upRate :: !Rate -- ^ Upload rate towards the peer (estimated)+ , downRate :: !Rate -- ^ Download rate from the peer (estimated)+ , runningEndgame :: !Bool -- ^ True if we are in endgame+ }+++peerP :: MgrChannel -> RateTVar -> PieceMgrChannel -> FSPChannel -> PieceMap -> Int+ -> TChan SenderQ.SenderQMsg -> TChan (Message, Integer) -> BandwidthChannel+ -> TVar [PStat] -> InfoHash+ -> SupervisorChan -> IO ThreadId+peerP pMgrC rtv pieceMgrC fsC pm nPieces outBound inBound sendBWC stv ih supC = do+ ch <- newTChanIO+ tch <- newTChanIO+ ct <- getCurrentTime+ pieceSet <- PS.new nPieces+ spawnP (PCF inBound outBound pMgrC pieceMgrC fsC ch sendBWC tch stv rtv ih pm)+ (PST True False S.empty True False pieceSet (RC.new ct) (RC.new ct) False)+ (cleanupP startup (defaultStopHandler supC) cleanup)+ where startup = do+ tid <- liftIO $ myThreadId+ debugP "Syncing a connectBack"+ asks peerCh >>= (\ch -> do+ c <- asks peerMgrCh+ liftIO . atomically $ writeTChan c $ Connect ih tid ch)+ pieces <- getPiecesDone+ outChan $ SenderQ.SenderQM $ BitField (constructBitField nPieces pieces)+ -- Install the StatusP timer+ c <- asks timerCh+ _ <- registerSTM 5 c ()+ foreverP eventLoop++ cleanup = do+ t <- liftIO myThreadId+ pieces <- gets peerPieces >>= PS.toList+ ch <- asks pieceMgrCh+ ch2 <- asks peerMgrCh+ liftIO . atomically $ writeTChan ch (PeerUnhave pieces)+ liftIO . atomically $ writeTChan ch2 (Disconnect t)++ readOp :: Process PCF PST Operation+ readOp = do+ inb <- asks inCh+ chk <- asks peerCh+ tch <- asks timerCh+ bwc <- asks sendBWCh+ liftIO . atomically $+ (readTChan inb >>= return . PeerMsgEvt) `orElse`+ (readTChan chk >>= return . ChokeMgrEvt) `orElse`+ (readTChan tch >> return TimerEvent) `orElse`+ (readTChan bwc >>= return . UpRateEvent)++ eventLoop = do+ 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})+ TimerEvent -> timerTick++data Operation = PeerMsgEvt (Message, Integer)+ | ChokeMgrEvt PeerMessage+ | TimerEvent+ | UpRateEvent Integer+ ++-- | Return a list of pieces which are currently done by us+getPiecesDone :: Process PCF PST [PieceNum]+getPiecesDone = do+ ch <- asks pieceMgrCh+ liftIO $ do+ c <- newEmptyTMVarIO+ atomically $ writeTChan ch (GetDone c)+ atomically $ takeTMVar c++-- | Process an event from the Choke Manager+chokeMsg :: PeerMessage -> Process PCF PST ()+chokeMsg msg = do+ debugP "ChokeMgrEvent"+ case msg of+ PieceCompleted pn -> outChan $ SenderQ.SenderQM $ Have pn+ ChokePeer -> do choking <- gets weChoke+ when (not choking)+ (do outChan $ SenderQ.SenderOChoke+ debugP "ChokePeer"+ modify (\s -> s {weChoke = True}))+ UnchokePeer -> do choking <- gets weChoke+ when choking+ (do outChan $ SenderQ.SenderQM Unchoke+ debugP "UnchokePeer"+ modify (\s -> s {weChoke = False}))+ CancelBlock pn blk -> do+ modify (\s -> s { blockQueue = S.delete (pn, blk) $ blockQueue s })+ outChan $ SenderQ.SenderQRequestPrune pn blk++-- True if the peer is a seeder+-- Optimization: Don't calculate this all the time. It only changes once and then it keeps+-- being there.+isASeeder :: Process PCF PST Bool+isASeeder = liftM2 (==) (gets peerPieces >>= PS.size) (M.size <$> asks pieceMap)++-- A Timer event handles a number of different status updates. One towards the+-- Choke Manager so it has a information about whom to choke and unchoke - and+-- one towards the status process to keep track of uploaded and downloaded+-- stuff.+timerTick :: Process PCF PST ()+timerTick = do+ mTid <- liftIO myThreadId+ debugP "TimerEvent"+ tch <- asks timerCh+ _ <- registerSTM 5 tch ()+ -- Tell the ChokeMgr about our progress+ ur <- gets upRate+ dr <- gets downRate+ t <- liftIO $ getCurrentTime+ let (up, nur) = RC.extractRate t ur+ (down, ndr) = RC.extractRate t dr+ infoP $ "Peer has rates up/down: " ++ show up ++ "/" ++ show down+ i <- gets peerInterested+ seed <- isASeeder+ pchoke <- gets peerChoke+ rtv <- asks rateTV+ liftIO . atomically $ do+ q <- readTVar rtv+ writeTVar rtv ((mTid, (up, down, i, seed, 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+ q <- readTVar stv+ writeTVar stv (PStat { pInfoHash = ih+ , pUploaded = upCnt+ , pDownloaded = downCnt } : q)+ modify (\s -> s { upRate = nuRate, downRate = ndRate })+++-- | Process an Message from the peer in the other end of the socket.+peerMsg :: Message -> Integer -> Process PCF PST ()+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+ 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+ fillBlocks+ Cancel pn blk -> cancelMsg pn blk+ Port _ -> return () -- No DHT yet, silently ignore++-- | 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 PCF PST ()+putbackBlocks = do+ blks <- gets blockQueue+ pmch <- asks pieceMgrCh+ liftIO . atomically $ writeTChan pmch (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+haveMsg :: PieceNum -> Process PCF PST ()+haveMsg pn = do+ pm <- asks pieceMap+ if M.member pn pm+ then do PS.insert pn =<< gets peerPieces+ pmch <- asks pieceMgrCh+ liftIO . atomically $ writeTChan pmch (PeerHave [pn])+ considerInterest+ else do warningP "Unknown Piece"+ stopP++-- | Process a BITFIELD message from the peer. Side effect: Consider Interest.+bitfieldMsg :: BitField -> Process PCF PST ()+bitfieldMsg bf = do+ pieces <- gets peerPieces+ piecesNull <- PS.null pieces+ if piecesNull+ -- TODO: Don't trust the bitfield+ then do nPieces <- M.size <$> asks pieceMap+ pp <- createPeerPieces nPieces bf+ modify (\s -> s { peerPieces = pp })+ peerLs <- PS.toList pp+ pmch <- asks pieceMgrCh+ liftIO . atomically $ writeTChan pmch (PeerHave peerLs)+ considerInterest+ else do infoP "Got out of band Bitfield request, dying"+ stopP++-- | Process a request message from the Peer+requestMsg :: PieceNum -> Block -> Process PCF PST ()+requestMsg pn blk = do+ choking <- gets weChoke+ unless (choking)+ (do+ bs <- readBlock pn blk -- TODO: Pushdown to send process+ outChan $ SenderQ.SenderQM $ Piece pn (blockOffset blk) bs)++-- | Read a block from the filesystem for sending+readBlock :: PieceNum -> Block -> Process PCF PST B.ByteString+readBlock pn blk = do+ v <- liftIO $ newEmptyTMVarIO+ fch <- asks fsCh+ liftIO $ do+ atomically $ writeTChan fch (ReadBlock pn blk v)+ atomically $ takeTMVar v++-- | Handle a Piece Message incoming from the peer+pieceMsg :: PieceNum -> Int -> B.ByteString -> Process PCF PST ()+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 (Block os sz) bs+ modify (\s -> s { blockQueue = S.delete e (blockQueue s)}))++-- | Handle a cancel message from the peer+cancelMsg :: PieceNum -> Block -> Process PCF PST ()+cancelMsg n blk = outChan $ SenderQ.SenderQCancel n blk++-- | Update our interest state based on the pieces the peer has.+-- Obvious optimization: Do less work, there is no need to consider all pieces most of the time+considerInterest :: Process PCF PST ()+considerInterest = do+ c <- liftIO newEmptyTMVarIO+ pcs <- gets peerPieces+ pmch <- asks pieceMgrCh+ interested <- liftIO $ do+ atomically $ writeTChan pmch (AskInterested pcs c)+ atomically $ takeTMVar c+ if interested+ then do modify (\s -> s { weInterested = True })+ outChan $ SenderQ.SenderQM Interested+ else modify (\s -> s { weInterested = False})++-- | 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 PCF PST ()+fillBlocks = do+ choked <- gets peerChoke+ unless choked 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+-- long as the peer is interested in talking to us.+-- TODO: Decide on a queue size based on the current download rate.+checkWatermark :: Process PCF PST ()+checkWatermark = do+ q <- gets blockQueue+ eg <- gets runningEndgame+ 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)++-- These three values are chosen rather arbitrarily at the moment.+loMark :: Int+loMark = 10++hiMark :: Int+hiMark = 15++-- Low mark when running in endgame mode+endgameLoMark :: Int+endgameLoMark = 1+++-- | Queue up pieces for retrieval at the Peer+queuePieces :: [(PieceNum, Block)] -> Process PCF PST ()+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 PCF PST ()+pushRequest pn blk = outChan $ SenderQ.SenderQM $ Request pn blk++-- | Tell the PieceManager to store the given block+storeBlock :: PieceNum -> Block -> B.ByteString -> Process PCF PST ()+storeBlock n blk bs = do+ pmch <- asks pieceMgrCh+ liftIO . atomically $ writeTChan pmch (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.+grabBlocks :: Int -> Process PCF PST [(PieceNum, Block)]+grabBlocks n = do+ c <- liftIO newEmptyTMVarIO+ ps <- gets peerPieces+ pmch <- asks pieceMgrCh+ blks <- liftIO $ do+ atomically $ writeTChan pmch (GrabBlocks n ps c)+ 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+createPeerPieces nPieces =+ PS.fromList nPieces . map fromIntegral . concat . decodeBytes 0 . L.unpack+ where decodeByte :: Int -> Word8 -> [Maybe Int]+ decodeByte soFar w =+ let dBit n = if testBit w (7-n)+ then Just (n+soFar)+ else Nothing+ in fmap dBit [0..7]+ decodeBytes _ [] = []+ decodeBytes soFar (w : ws) = catMaybes (decodeByte soFar w) : decodeBytes (soFar + 8) ws++-- | Send a message on a chan from the process queue+outChan :: SenderQ.SenderQMsg -> Process PCF PST ()+outChan qm = do+ ch <- asks outCh+ liftIO . atomically $ writeTChan ch qm
+ src/Process/Peer/Receiver.hs view
@@ -0,0 +1,56 @@+module Process.Peer.Receiver+ ( start )+where++import Control.Concurrent+import Control.Concurrent.STM++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 Data.Word++import System.IO++import Process+import Supervisor+import Protocol.Wire+++data CF = CF { rpMsgCh :: TChan (Message, Integer) }++instance Logging CF where+ logName _ = "Process.Peer.Receiver"++start :: Handle -> TChan (Message, Integer)+ -> SupervisorChan -> IO ThreadId+start h ch supC = spawnP (CF ch) h+ (catchP (foreverP readSend)+ (defaultStopHandler supC))++readSend :: Process CF Handle ()+readSend = do+ h <- 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)++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
+ src/Process/Peer/Sender.hs view
@@ -0,0 +1,52 @@+module Process.Peer.Sender+ ( start )+where++import Control.Applicative+import Control.Concurrent+import Control.Concurrent.STM++import Control.Monad.Reader+import Control.Monad.State++import qualified Data.ByteString as B++import System.IO++import Process+import Protocol.Wire+import Supervisor++data CF = CF { chan :: TMVar B.ByteString }++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 B.ByteString -> SupervisorChan -> IO ThreadId+start h ch supC = spawnP (CF ch) h (catchP (foreverP pgm)+ (do t <- liftIO $ myThreadId+ liftIO . atomically $ writeTChan supC $ IAmDying t+ liftIO $ hClose h))++pgm :: Process CF Handle ()+pgm = {-# SCC "Peer.Sender" #-} do+ ch <- asks chan+ tout <- liftIO $ registerDelay defaultTimeout+ r <- liftIO . atomically $ do+ t <- readTVar tout+ if t+ then return Nothing+ else Just <$> takeTMVar ch+ h <- get+ case r of+ Nothing -> putMsg (encodePacket KeepAlive)+ Just m -> putMsg m+ liftIO $ hFlush h+ where+ putMsg m = do+ h <- get+ liftIO $ B.hPut h m++defaultTimeout :: Int+defaultTimeout = 120 * 1000000
+ src/Process/Peer/SenderQ.hs view
@@ -0,0 +1,101 @@+module Process.Peer.SenderQ+ ( SenderQMsg(..)+ , start+ )+where++import Control.Concurrent+import Control.Concurrent.STM++import Control.Monad.Reader+import Control.Monad.State++import Prelude hiding (catch, log)++import qualified Data.ByteString as B++import Channels+import Process+import qualified Data.Queue as Q+import Supervisor+import Torrent+import Protocol.Wire++-- | Messages we can send to the Send Queue+data SenderQMsg = SenderQCancel PieceNum Block -- ^ Peer requested that we cancel a piece+ | SenderQM Message -- ^ We want to send the Message to the peer+ | SenderOChoke -- ^ We want to choke the peer+ | SenderQRequestPrune PieceNum Block -- ^ Prune SendQueue of this (pn, blk) pair++data CF = CF { sqIn :: TChan SenderQMsg+ , sqOut :: TMVar B.ByteString+ , bandwidthCh :: BandwidthChannel+ }++data ST = ST { outQueue :: Q.Queue Message+ , bytesTransferred :: Integer+ }++instance Logging CF where+ logName _ = "Process.Peer.SendQueue"++-- | sendQueue Process, simple version.+-- TODO: Split into fast and slow.+start :: TChan SenderQMsg -> TMVar B.ByteString -> BandwidthChannel+ -> SupervisorChan+ -> IO ThreadId+start inC outC bandwC supC = spawnP (CF inC outC bandwC) (ST Q.empty 0)+ (catchP (foreverP pgm)+ (defaultStopHandler supC))++pgm :: Process CF ST ()+pgm = {-# SCC "Peer.SendQueue" #-} do+ q <- gets outQueue+ l <- gets bytesTransferred+ -- Gather together events which may trigger+ when (l > 0) rateUpdateEvent+ ic <- asks sqIn+ ov <- asks sqOut+ r <- case Q.first q of+ Nothing -> liftIO $ atomically (readTChan ic >>= return . Right)+ Just p -> do let bs = encodePacket p+ sz = fromIntegral $ B.length bs+ liftIO . atomically $+ (putTMVar ov bs >> return (Left sz)) `orElse`+ (readTChan ic >>= return . Right)+ case r of+ Left sz ->+ modify (\s -> s { bytesTransferred = bytesTransferred s + sz+ , outQueue = Q.remove (outQueue s)})+ Right m ->+ case m of+ SenderQM msg -> modifyQ (Q.push msg)+ SenderQCancel n blk -> modifyQ (Q.filter (filterPiece n (blockOffset blk)))+ SenderOChoke -> do modifyQ (Q.filter filterAllPiece)+ modifyQ (Q.push Choke)+ SenderQRequestPrune n blk ->+ modifyQ (Q.filter (filterRequest n blk))++rateUpdateEvent :: Process CF ST ()+rateUpdateEvent = {-# SCC "Peer.SendQ.rateUpd" #-} do+ l <- gets bytesTransferred+ bwc <- asks bandwidthCh+ liftIO . atomically $ writeTChan bwc l+ modify (\s -> s { bytesTransferred = 0 })++filterAllPiece :: Message -> Bool+filterAllPiece (Piece _ _ _) = True+filterAllPiece _ = False++filterPiece :: PieceNum -> Int -> Message -> Bool+filterPiece n off (Piece n1 off1 _) | n == n1 && off == off1 = False+ | otherwise = True+filterPiece _ _ _ = True++filterRequest :: PieceNum -> Block -> Message -> Bool+filterRequest n blk (Request n1 blk1) | n == n1 && blk == blk1 = False+ | otherwise = True+filterRequest _ _ _ = True++modifyQ :: (Q.Queue Message -> Q.Queue Message) -> Process CF ST ()+modifyQ f = modify (\s -> s { outQueue = f (outQueue s) })
+ src/Process/PeerMgr.hs view
@@ -0,0 +1,217 @@+{-# LANGUAGE TupleSections #-}+module Process.PeerMgr (+ -- * Types+ Peer(..)+ , PeerMgrMsg(..)+ , PeerMgrChannel+ , TorrentLocal(..)+ -- * Interface+ , Process.PeerMgr.start+)+where++import qualified Data.Map as M++import Control.Concurrent+import Control.Concurrent.STM++import Control.Monad.State+import Control.Monad.Reader++import Network+import System.IO+import System.Log.Logger++import Channels+import Process+import Process.Peer as Peer+import Process.ChokeMgr hiding (start)+import Process.FS hiding (start)+import Process.PieceMgr hiding (start)+import Process.Status hiding (start)+import Protocol.Wire++import Supervisor+import Torrent hiding (infoHash)++data PeerMgrMsg = PeersFromTracker InfoHash [Peer]+ | NewIncoming (Handle, HostName, PortNumber)+ | NewTorrent InfoHash TorrentLocal+ | StopTorrent InfoHash++data TorrentLocal = TorrentLocal+ { tcPcMgrCh :: PieceMgrChannel+ , tcFSCh :: FSPChannel+ , tcStatTV :: TVar [PStat]+ , tcPM :: PieceMap+ }++++type PeerMgrChannel = TChan PeerMgrMsg++data CF = CF { peerCh :: PeerMgrChannel+ , mgrCh :: MgrChannel+ , peerPool :: SupervisorChan+ , chokeMgrCh :: ChokeMgrChannel+ , chokeRTV :: RateTVar+ }++instance Logging CF where+ logName _ = "Process.PeerMgr"+++type ChanManageMap = M.Map InfoHash TorrentLocal++data ST = ST { peersInQueue :: [(InfoHash, Peer)]+ , peers :: M.Map ThreadId PeerChannel+ , peerId :: PeerId+ , cmMap :: ChanManageMap+ }++start :: PeerMgrChannel -> PeerId+ -> ChokeMgrChannel -> RateTVar -> SupervisorChan+ -> IO ThreadId+start ch pid chokeMgrC rtv supC =+ do mgrC <- newTChanIO+ fakeChan <- newTChanIO+ pool <- liftM snd $ oneForOne "PeerPool" [] fakeChan+ spawnP (CF ch mgrC pool chokeMgrC rtv)+ (ST [] M.empty pid cmap) (catchP (forever lp)+ (defaultStopHandler supC))+ where+ cmap = M.empty+ lp = do+ pc <- asks peerCh+ mc <- asks mgrCh+ q <- liftIO . atomically $+ (readTChan pc >>= return . Left) `orElse`+ (readTChan mc >>= return . Right)+ case q of+ Left msg -> incomingPeers msg+ Right msg -> peerEvent msg+ fillPeers++incomingPeers :: PeerMgrMsg -> Process CF ST ()+incomingPeers msg =+ case msg of+ PeersFromTracker ih ps -> do+ debugP "Adding peers to queue"+ modify (\s -> s { peersInQueue = (map (ih,) ps) ++ peersInQueue s })+ NewIncoming conn@(h, _, _) -> 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+ NewTorrent ih tl -> do+ modify (\s -> s { cmMap = M.insert ih tl (cmMap s)})+ StopTorrent _ih -> do+ errorP "Not implemented stopping yet"++peerEvent :: MgrMessage -> Process CF ST ()+peerEvent msg = case msg of+ Connect ih tid c -> newPeer ih tid c+ Disconnect tid -> removePeer tid+ where+ newPeer ih tid c = do debugP $ "Adding new peer " ++ show tid+ cch <- asks chokeMgrCh+ liftIO . atomically $ writeTChan cch (AddPeer ih tid c)+ modify (\s -> s { peers = M.insert tid c (peers s)})+ removePeer tid = do debugP $ "Removing peer " ++ show tid+ cch <- asks chokeMgrCh+ liftIO . atomically $ writeTChan cch (RemovePeer tid)+ modify (\s -> s { peers = M.delete tid (peers s)})++numPeers :: Int+numPeers = 40++fillPeers :: Process CF ST ()+fillPeers = do+ sz <- liftM M.size $ gets peers+ when (sz < numPeers)+ (do q <- gets peersInQueue+ let (toAdd, rest) = splitAt (numPeers - sz) q+ debugP $ "Filling with up to " ++ show (numPeers - sz) ++ " peers"+ mapM_ addPeer toAdd+ modify (\s -> s { peersInQueue = rest }))++addPeer :: (InfoHash, Peer) -> Process CF ST ThreadId+addPeer (ih, (Peer hn prt)) = 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++addIncoming :: (Handle, HostName, PortNumber) -> Process CF ST ThreadId+addIncoming conn = do+ ppid <- gets peerId+ pool <- asks peerPool+ mgrC <- asks mgrCh+ v <- asks chokeRTV+ cm <- gets cmMap+ liftIO $ acceptor conn pool ppid mgrC v cm++type ConnectRecord = (HostName, PortID, PeerId, InfoHash)++connect :: ConnectRecord -> SupervisorChan -> MgrChannel -> RateTVar -> ChanManageMap+ -> IO ThreadId+connect (host, port, 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+ debugM "Process.PeerMgr.connect" "Connected, initiating handShake"+ r <- initiateHandshake h pid ih+ debugM "Process.PeerMgr.connect" "Handshake run"+ case r of+ Left err -> do debugM "Process.PeerMgr.connect"+ ("Peer handshake failure at host " ++ host+ ++ " with error " ++ err)+ return ()+ 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) (M.size (tcPM tc)) ihsh+ atomically $ writeTChan pool $+ SpawnNew (Supervisor $ allForOne "PeerSup" children)+ return ()++acceptor :: (Handle, HostName, PortNumber) -> SupervisorChan+ -> PeerId -> MgrChannel -> RateTVar -> ChanManageMap+ -> IO ThreadId+acceptor (h,hn,pn) pool pid mgrC rtv cmmap =+ forkIO (connector >> return ())+ where ihTst k = M.member k cmmap+ connector = do+ debugLog "Handling incoming connection"+ r <- receiveHandshake h 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) ->+ 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)+ (tcStatTV tc) (tcPM tc) (M.size (tcPM tc)) ih+ atomically $ writeTChan pool $+ SpawnNew (Supervisor $ allForOne "PeerSup" children)+ return ()+ debugLog = debugM "Process.PeerMgr.acceptor"++showPort :: PortID -> String+showPort (PortNumber pn) = show pn+showPort _ = "N/A"
+ src/Process/PieceMgr.hs view
@@ -0,0 +1,548 @@+module Process.PieceMgr+ ( PieceMgrMsg(..)+ , PieceMgrChannel+ , Blocks(..)+ , start+ , createPieceDb+ )+where++import Control.Applicative+import Control.Concurrent+import Control.Concurrent.STM+import Control.Exception (assert)++import Control.Monad.Reader+import Control.Monad.State+import Data.List+import qualified Data.PendingSet as PendS+import qualified Data.ByteString as B+import qualified Data.Map as M+import qualified Data.Set as S+import qualified Data.PieceSet as PS++import Prelude hiding (log)++import System.Random+import System.Random.Shuffle++import Process.FS hiding (start)+import Process.Status as STP hiding (start) +import Process.ChokeMgr (ChokeMgrMsg(..), ChokeMgrChannel)+import Supervisor+import Torrent+import Tracer+import Process++----------------------------------------------------------------------++-- | The Piece Database tracks the current state of the Torrent with respect to pieces.+-- In the database, we book-keep what pieces are missing, what are done and what are+-- currently in the process of being downloaded. The crucial moment is when we think+-- we have a full piece: we check it against its SHA1 and if it is good, we can mark+-- that piece as done.+--+-- Better implementations for selecting among the pending Pieces is probably crucial+-- to an effective client, but we keep it simple for now.+data ST = ST+ { pendingPieces :: !PS.PieceSet -- ^ Pieces currently pending download+ , donePiece :: !PS.PieceSet -- ^ Pieces that are done+ , donePush :: ![ChokeMgrMsg] -- ^ Pieces that should be pushed to the Choke Mgr.+ , inProgress :: !(M.Map PieceNum InProgressPiece) -- ^ Pieces in progress+ , downloading :: ![(PieceNum, Block)] -- ^ Blocks we are currently downloading+ , infoMap :: !PieceMap -- ^ Information about pieces+ , endGaming :: !Bool -- ^ If we have done any endgame work this is true+ , histogram :: !PendS.PendingSet -- ^ Track the rarity of pieces+ , assertCount :: !Int -- ^ When to next check the database for consistency+ , traceBuffer :: !(Tracer String)+ }++sizeReport :: Process CF ST String+sizeReport = do+ (ST pend done dpush progress down _ _ histo _ _) <- get+ p1sz <- PS.size pend+ p2sz <- PS.size done+ return $ show+ [ ("Pending", p1sz)+ , ("Done" , p2sz)+ , ("DonePush", length dpush)+ , ("InProgress", M.size progress)+ , ("Downloading", length down)+ , ("Histo", PendS.size histo) ]++-- | The InProgressPiece data type describes pieces in progress of being downloaded.+-- we keep track of blocks which are pending as well as blocks which are done. We+-- also keep track of a count of the blocks. When a block is done, we cons it unto+-- @ipHaveBlocks@. When @ipHave == ipDone@, we check the piece for correctness. The+-- field @ipHaveBlocks@ could in principle be omitted, but for now it is kept since+-- we can use it for asserting implementation correctness. We note that both the+-- check operations are then O(1) and probably fairly fast.+data InProgressPiece = InProgressPiece+ { ipDone :: Int -- ^ Number of blocks when piece is done+ , ipHaveBlocks :: S.Set Block -- ^ The blocks we have+ , ipPendingBlocks :: [Block] -- ^ Blocks still pending+ } deriving Show++-- INTERFACE+----------------------------------------------------------------------++-- | When the PieceMgrP returns blocks to a peer, it will return them in either+-- "Leech Mode" or in "Endgame mode". The "Leech mode" is when the client is+-- leeching like normal. The "Endgame mode" is when the client is entering the+-- endgame. This means that the Peer should act differently to the blocks.+data Blocks = Leech [(PieceNum, Block)]+ | Endgame [(PieceNum, Block)]++-- | Messages for RPC towards the PieceMgr.+data PieceMgrMsg = GrabBlocks Int PS.PieceSet (TMVar Blocks)+ -- ^ Ask for grabbing some blocks+ | StoreBlock PieceNum Block B.ByteString+ -- ^ Ask for storing a block on the file system+ | PutbackBlocks [(PieceNum, Block)]+ -- ^ Put these blocks back for retrieval+ | AskInterested PS.PieceSet (TMVar Bool)+ -- ^ Ask if any of these pieces are interesting+ | GetDone (TMVar [PieceNum])+ -- ^ Get the pieces which are already done+ | PeerHave [PieceNum]+ -- ^ A peer has the given piece(s)+ | PeerUnhave [PieceNum]+ -- ^ A peer relinquished the given piece Indexes++instance Show PieceMgrMsg where+ show (GrabBlocks x _ _) = "GrabBlocks " ++ show x+ show (StoreBlock pn blk _) = "StoreBlock " ++ show pn ++ " " ++ show blk+ show (PutbackBlocks x) = "PutbackBlocks " ++ show x+ show (AskInterested _ _) = "AskInterested"+ show (GetDone _) = "GetDone"+ show (PeerHave xs) = "PeerHave " ++ show xs+ show (PeerUnhave xs) = "PeerUnhave " ++ show xs++type PieceMgrChannel = TChan PieceMgrMsg++data CF = CF+ { pieceMgrCh :: PieceMgrChannel+ , fspCh :: FSPChannel+ , chokeCh :: ChokeMgrChannel+ , statusCh :: StatusChannel+ , pMgrInfoHash :: InfoHash+ }++instance Logging CF where+ logName _ = "Process.PieceMgr"++type PieceMgrProcess v = Process CF ST v++start :: PieceMgrChannel -> FSPChannel -> ChokeMgrChannel -> StatusChannel -> ST -> InfoHash+ -> SupervisorChan -> IO ThreadId+start mgrC fspC chokeC statC db ih supC =+ {-# SCC "PieceMgr" #-}+ spawnP (CF mgrC fspC chokeC statC ih) db+ (catchP (forever pgm)+ (defaultStopHandler supC))+ where pgm = do+ assertST+ rpcMessage+ drainSend++drainSend :: Process CF ST ()+drainSend = do+ 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) })++traceMsg :: PieceMgrMsg -> Process CF ST ()+traceMsg m = modify (\db -> db { traceBuffer = trace (show m) (traceBuffer db) })++rpcMessage :: Process CF ST ()+rpcMessage = do+ ch <- asks pieceMgrCh+ m <- liftIO . atomically $ readTChan ch+ traceMsg m+ case m of+ GrabBlocks n eligible c ->+ do blocks <- grabBlocks n eligible+ liftIO . atomically $ do putTMVar c blocks -- Is never supposed to block+ StoreBlock pn blk d -> storeBlock pn blk d+ PutbackBlocks blks -> mapM_ putbackBlock blks+ GetDone c -> do+ done <- PS.toList =<< gets donePiece+ liftIO . atomically $ do putTMVar c done -- Is never supposed to block either+ PeerHave idxs -> peerHave idxs+ PeerUnhave idxs -> peerUnhave idxs+ AskInterested pieces retC -> do+ 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)+ fch <- asks fspCh+ liftIO . atomically $ writeTChan fch $ WriteBlock pn blk d+ modify (\s -> s { downloading = downloading s \\ [(pn, blk)] })+ endgameBroadcast pn blk+ done <- updateProgress pn blk+ when done+ (do assertPieceComplete pn+ pend <- gets pendingPieces+ iprog <- gets inProgress+ pendSz <- PS.size pend+ infoP $ "Piece #" ++ show pn+ ++ " completed, there are "+ ++ (show pendSz) ++ " pending "+ ++ (show $ M.size iprog) ++ " in progress"+ l <- gets infoMap >>=+ (\pm -> case M.lookup pn pm of+ Nothing -> fail "Storeblock: M.lookup"+ Just x -> return $ len x)+ ih <- asks pMgrInfoHash+ c <- asks statusCh+ liftIO . atomically $ writeTChan c (CompletedPiece ih l)+ pieceOk <- checkPiece pn+ case pieceOk of+ Nothing ->+ do fail "PieceMgrP: Piece Nonexisting!"+ Just True -> do completePiece pn+ markDone pn+ checkFullCompletion+ Just False -> putbackPiece pn)++askInterested :: PS.PieceSet -> Process CF ST Bool+askInterested pieces = do+ inProg <- M.keys <$> gets inProgress+ amongProg <- filterM (flip PS.member pieces) inProg+ if (not $ null amongProg)+ then return True+ else do+ pend <- gets pendingPieces+ intsct <- PS.intersection pieces pend+ return (not $ null intsct)++peerHave :: [PieceNum] -> Process CF ST ()+peerHave idxs = modify (\db -> db { histogram = PendS.haves idxs (histogram db)})++peerUnhave :: [PieceNum] -> Process CF ST ()+peerUnhave idxs = modify (\db -> db { histogram = PendS.unhaves idxs (histogram db)})++endgameBroadcast :: PieceNum -> Block -> Process CF ST ()+endgameBroadcast pn blk = do+ ih <- asks pMgrInfoHash+ gets endGaming >>=+ flip when (modify (\db -> db { donePush = (BlockComplete ih pn blk) : donePush db }))++markDone :: PieceNum -> Process CF ST ()+markDone pn = do+ ih <- asks pMgrInfoHash+ modify (\db -> db { donePush = (PieceDone ih pn) : donePush db })++checkPiece :: PieceNum -> Process CF ST (Maybe Bool)+checkPiece n = do+ v <- liftIO newEmptyTMVarIO+ fch <- asks fspCh+ liftIO $ do+ atomically $ writeTChan fch $ CheckPiece n v+ atomically $ takeTMVar v++-- HELPERS+----------------------------------------------------------------------++createPieceDb :: MonadIO m => PiecesDoneMap -> PieceMap -> m ST+createPieceDb mmap pmap = do+ pending <- filt (==False)+ done <- filt (==True)+ return $ ST pending done [] M.empty [] pmap False PendS.empty 0 (Tracer.new 20)+ where+ filt f = PS.fromList (M.size pmap) . M.keys $ M.filter f mmap++----------------------------------------------------------------------++-- | The call @completePiece db pn@ will mark that the piece @pn@ is completed+completePiece :: PieceNum -> PieceMgrProcess ()+completePiece pn = do+ PS.insert pn =<< gets donePiece+ modify (\db -> db { inProgress = M.delete pn (inProgress db) })++-- | Handle torrent completion+checkFullCompletion :: PieceMgrProcess ()+checkFullCompletion = do+ doneP <- gets donePiece+ im <- gets infoMap+ donePSz <- PS.size doneP+ when (M.size im == donePSz)+ (do liftIO $ putStrLn "Torrent Completed; to honor the torrent-gods thou must now sacrifice a goat!"+ ih <- asks pMgrInfoHash+ asks statusCh >>= (\ch -> liftIO . atomically $ writeTChan ch (STP.TorrentCompleted ih))+ c <- asks chokeCh+ liftIO . atomically $ writeTChan c (TorrentComplete ih))++-- | The call @putBackPiece db pn@ will mark the piece @pn@ as not being complete+-- and put it back into the download queue again.+putbackPiece :: PieceNum -> PieceMgrProcess ()+putbackPiece pn = do+ PS.insert pn =<< gets pendingPieces+ modify (\db -> db { inProgress = M.delete pn (inProgress db) })++-- | Put back a block for downloading.+-- TODO: This is rather slow, due to the (\\) call, but hopefully happens rarely.+putbackBlock :: (PieceNum, Block) -> PieceMgrProcess ()+putbackBlock (pn, blk) = do+ done <- gets donePiece+ doneMember <- PS.member pn done+ unless (doneMember) -- Happens at endgame, stray block+ $ modify (\db -> db { inProgress = ndb (inProgress db)+ , downloading = downloading db \\ [(pn, blk)]})+ where ndb db = M.alter f pn db+ -- In endgame, the first will never happen. If it is done, the doneMember+ -- check above should take care of the problem. If the block has been downloaded+ -- by another peer in endgame, there is nothing to do.+ --+ -- Otherwise, we put the block back as pending. If in endgame, the next request+ -- will pull it into downloading again for endgaming.+ f Nothing = fail "The 'impossible' happened"+ f (Just ipp)+ | S.member blk (ipHaveBlocks ipp) = Just ipp+ | otherwise = Just ipp {ipPendingBlocks = blk : ipPendingBlocks ipp}++-- | Assert that a Piece is Complete. Can be omitted when we know it works+-- and we want a faster client.+assertPieceComplete :: PieceNum -> PieceMgrProcess ()+assertPieceComplete pn = do+ inprog <- gets inProgress+ ipp <- case M.lookup pn inprog of+ Nothing -> fail "assertPieceComplete: Could not lookup piece number"+ Just x -> return x+ dl <- gets downloading+ pm <- gets infoMap+ sz <- case M.lookup pn pm of+ Nothing -> fail "assertPieceComplete: Could not lookup piece in piecemap"+ Just x -> return $ len x+ unless (assertAllDownloaded dl pn)+ (fail "Could not assert that all pieces were downloaded when completing a piece")+ unless (assertComplete ipp sz)+ (fail $ "Could not assert completion of the piece #" ++ show pn+ ++ " with block state " ++ show ipp)+ where assertComplete ip sz = checkContents 0 (fromIntegral sz) (S.toAscList (ipHaveBlocks ip))+ -- Check a single block under assumptions of a cursor at offs+ checkBlock (offs, l, state) blk = (offs + blockSize blk,+ l - blockSize blk,+ state && offs == blockOffset blk)+ checkContents os l blks = case foldl checkBlock (os, l, True) blks of+ (_, 0, True) -> True+ _ -> False+ assertAllDownloaded blocks p = all (\(p', _) -> p /= p') blocks++-- | Update the progress on a Piece. When we get a block from the piece, we will+-- track this in the Piece Database. This function returns @complete@+-- where @complete@ is @True@ if the piece is percieved to be complete and @False@+-- otherwise.+updateProgress :: PieceNum -> Block -> PieceMgrProcess Bool+updateProgress pn blk = do+ ipdb <- gets inProgress+ case M.lookup pn ipdb of+ Nothing -> do warningP "updateProgress can't find progress block, error?"+ return False+ Just pg ->+ let blkSet = ipHaveBlocks pg+ in if blk `S.member` blkSet+ then return False -- Stray block download.+ -- Will happen without FAST extension+ -- at times+ else checkComplete pg { ipHaveBlocks = S.insert blk blkSet }+ where checkComplete pg = do+ modify (\db -> db { inProgress = M.adjust (const pg) pn (inProgress db) })+ debugP $ "Iphave : " ++ show (ipHave pg) ++ " ipDone: " ++ show (ipDone pg)+ return (ipHave pg == ipDone pg)+ ipHave = S.size . ipHaveBlocks++blockPiece :: BlockSize -> PieceSize -> [Block]+blockPiece blockSz pieceSize = build pieceSize 0 []+ where build 0 _os accum = reverse accum+ build leftBytes os accum | leftBytes >= blockSz =+ build (leftBytes - blockSz)+ (os + blockSz)+ $ Block os blockSz : accum+ | otherwise = build 0 (os + leftBytes) $ Block os leftBytes : accum++-- | The call @grabBlocks n eligible@ tries to pick off up to @n@ pieces from+-- to download. In doing so, it will only consider pieces in @eligible@. It+-- 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' []+ 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 })+ debugP $ "PieceMgr entered endgame."+ return $ Endgame blks+ else do modify (\s -> s { downloading = blocks ++ (downloading s) })+ 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)]+ -> 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++-- The Piece @p@ was found, grab it+grabFromProgress :: PieceNum -> PS.PieceSet -> PieceNum -> [(PieceNum, Block)]+ -> Process CF ST [(PieceNum, Block)]+grabFromProgress n ps p captured = do+ inprog <- gets inProgress+ ipp <- case M.lookup p inprog of+ Nothing -> fail "grabFromProgress: could not lookup piece"+ Just x -> return x+ let (grabbed, rest) = splitAt n (ipPendingBlocks ipp)+ 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)++-- Try grabbing pieces from the pending blocks+tryGrabPending :: PieceNum -> PS.PieceSet -> [(PieceNum, Block)] -> Process CF ST [(PieceNum, Block)]+tryGrabPending n ps captured = 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+ case culprits of+ Nothing -> do+ isn <- PS.intersection ps pending+ assert (null isn) (return ())+ return captured+ Just pieces -> do+ h <- pickRandom pieces+ inProg <- gets inProgress+ blockList <- createBlock h+ let sz = length blockList+ ipp = InProgressPiece sz S.empty blockList+ PS.delete h =<< gets pendingPieces+ modify (\db -> db { inProgress = M.insert h ipp inProg })+ tryGrabProgress n ps captured++grabEndGame :: PieceNum -> PS.PieceSet -> Process CF ST [(PieceNum, Block)]+grabEndGame n ps = 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++-- | Pick a random element among a finite list af them.+pickRandom :: MonadIO m => [a] -> m a+pickRandom ls = do+ n <- liftIO $ getStdRandom (\gen -> randomR (0, length ls - 1) gen)+ return $ ls !! n++-- | If given a Piece number, convert said number into its list of blocks to+-- download at peers.+createBlock :: PieceNum -> PieceMgrProcess [Block]+createBlock pn = do+ gets infoMap >>= (\im -> case M.lookup pn im of+ Nothing -> fail "createBlock: could not lookup piece"+ Just ipp -> return $ cBlock ipp)+ where cBlock = blockPiece defaultBlockSize . fromInteger . len++anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool+anyM f l = do r <- mapM f l+ return (any (==True) r)++assertST :: PieceMgrProcess ()+assertST = {-# SCC "assertST" #-} do+ c <- gets assertCount+ if c == 0+ then do modify (\db -> db { assertCount = 10 })+ assertSets >> assertInProgress >> assertDownloading+ sizes <- sizeReport+ debugP sizes+ else modify (\db -> db { assertCount = assertCount db - 1 })+ where+ -- If a piece is pending in the database, we have the following rules:+ --+ -- - It is not done.+ -- - It is not being downloaded+ -- - It is not in progresss.+ --+ -- If a piece is done, we have the following rules:+ --+ -- - It is not in progress.+ -- - There are no more downloading blocks.+ assertSets = do+ pending <- gets pendingPieces+ done <- gets donePiece+ down <- return . map fst =<< gets downloading+ iprog <- return . M.keys =<< gets inProgress+ pdownis <- anyM (flip PS.member pending) down+ donedownis <- anyM (flip PS.member done) down+ pdis <- PS.intersection pending done+ piprogis <- anyM (flip PS.member pending) iprog+ doneprogis <- anyM (flip PS.member done) iprog++ when (not $ null pdis)+ (do trb <- gets traceBuffer+ liftIO $ print trb+ return $ assert False ())+ when pdownis+ (do trb <- gets traceBuffer+ liftIO $ print trb+ return $ assert False ())+ when piprogis+ (do trb <- gets traceBuffer+ liftIO $ print trb+ return $ assert False ())+ when doneprogis+ (do trb <- gets traceBuffer+ liftIO $ print trb+ return $ assert False ())+ when donedownis+ (do trb <- gets traceBuffer+ liftIO $ print trb+ return $ assert False ())++ -- If a piece is in Progress, we have:+ --+ -- - There is a relationship with what pieces are downloading+ -- - If a block is ipPending, it is not in the downloading list+ -- - If a block is ipHave, it is not in the downloading list+ assertInProgress = do+ inProg <- gets inProgress+ mapM_ checkInProgress $ M.toList inProg+ checkInProgress (pn, ipp) = do+ when ( (S.size $ ipHaveBlocks ipp) >= ipDone ipp)+ (fail $ "Piece in progress " ++ show pn+ ++ " has downloaded more blocks than the piece has")+ assertDownloading = do+ down <- gets downloading+ mapM_ checkDownloading down+ checkDownloading (pn, blk) = do+ prog <- gets inProgress+ tr <- gets traceBuffer+ case M.lookup pn prog of+ Nothing -> fail $ "Piece " ++ show pn ++ " not in progress while We think it was"+ Just ipp -> do+ when (blk `elem` ipPendingBlocks ipp)+ (fail $ "P/Blk " ++ show (pn, blk) ++ " is in the Pending Block list")+ when (S.member blk $ ipHaveBlocks ipp)+ (fail $ "P/Blk " ++ show (pn, blk) ++ " is in the HaveBlocks set" +++ "Trace: " ++ show tr)+
+ src/Process/Status.hs view
@@ -0,0 +1,156 @@+-- | The status code runs a Status Process. This process keeps track+-- of a number of interval valies for a given torrent file and it+-- periodically updates the tracker process with the relevant+-- information about data uploaded, downloaded and how much is+-- left. The tracker is then responsible for using this data+-- correctly to tell the tracker what to do+{-# LANGUAGE FlexibleInstances #-}+module Process.Status (+ -- * Types+ StatusMsg(..)+ , PStat(..)+ -- * Channels+ , StatusChannel+ -- * State+ , StatusState(uploaded, downloaded, left)+ -- * Interface+ , start+ )+where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Exception (assert)++import Control.Monad.Reader+import Control.Monad.State++import Data.IORef+import qualified Data.Map as M++import Prelude hiding (log)++import Channels+import Process+import Supervisor+import Torrent+import Version++data StatusMsg = TrackerStat { trackInfoHash :: InfoHash+ , trackIncomplete :: Maybe Integer+ , trackComplete :: Maybe Integer }+ | CompletedPiece InfoHash Integer+ | InsertTorrent InfoHash Integer TrackerChannel+ | RemoveTorrent InfoHash+ | TorrentCompleted InfoHash+ | RequestStatus InfoHash (TMVar StatusState)+ | RequestAllTorrents (TMVar [(InfoHash, StatusState)])++data PStat = PStat { pInfoHash :: InfoHash+ , pUploaded :: Integer+ , pDownloaded :: Integer }++type StatusChannel = TChan StatusMsg++data CF = CF { statusCh :: StatusChannel,+ statusTV :: TVar [PStat] }++instance Logging CF where+ logName _ = "Process.Status"++type ST = M.Map InfoHash StatusState++data StatusState = SState+ { uploaded :: Integer+ , downloaded :: Integer+ , left :: Integer+ , incomplete :: Maybe Integer+ , complete :: Maybe Integer+ , state :: TorrentState+ , trackerMsgCh :: TrackerChannel+ }++gatherStats :: (Integer, Integer) -> [(String, String)]+gatherStats (upload, download) =+ [("uploaded", show upload), ("downloaded", show download),+ ("version", version)]++instance Show StatusState where+ show (SState up down l inc comp st _) = concat+ ["{ Uploaded: " ++ show up ++ "\n"+ ," Downloaded: " ++ show down ++ "\n"+ ," Left: " ++ show l ++ "\n"+ ," State: " ++ show st ++ "\n"+ ," Complete: " ++ show comp ++ "\n"+ ," Incomplete: " ++ show inc ++ " }"]++-- | Start a new Status process with an initial torrent state and a+-- channel on which to transmit status updates to the tracker.+start :: Maybe FilePath -> StatusChannel -> TVar [PStat] -> SupervisorChan -> IO ThreadId+start fp statusC tv supC = do+ r <- newIORef (0,0)+ spawnP (CF statusC tv) M.empty+ (cleanupP (foreverP (pgm r)) (defaultStopHandler supC) (cleanup r))+ where+ cleanup r = do+ st <- liftIO $ readIORef r+ case fp of+ Nothing -> return ()+ Just fpath -> liftIO $ writeFile fpath (show . gatherStats $ st)+ pgm r = {-# SCC "StatusP" #-} do+ fetchUpdates r+ d <- liftIO $ registerDelay (5 * 1000000)+ ch <- asks statusCh+ x <- liftIO . atomically $ do+ q <- readTVar d+ if q+ then return Nothing+ else return . Just =<< readTChan ch+ case x of+ Nothing -> return ()+ Just msg -> recvMsg msg++newMap :: Integer -> TrackerChannel -> StatusState+newMap l trackerMsgC =+ SState 0 0 l Nothing Nothing (if l == 0 then Seeding else Leeching) trackerMsgC++recvMsg :: StatusMsg -> Process CF ST ()+recvMsg msg = + case msg of+ TrackerStat ih ic c -> do+ modify (\s -> M.adjust (\st -> st { incomplete = ic, complete = c }) ih s)+ CompletedPiece ih bytes -> do+ modify (\s -> M.adjust (\st -> st { left = (left st) - bytes }) ih s)+ InsertTorrent ih l trackerMsgC ->+ modify (\s -> M.insert ih (newMap l trackerMsgC) s)+ RemoveTorrent ih -> modify (\s -> M.delete ih s)+ RequestStatus ih v -> do+ s <- get+ case M.lookup ih s of+ Nothing -> fail "Unknown InfoHash"+ Just st -> liftIO . atomically $ putTMVar v st+ RequestAllTorrents v -> do+ s <- get+ liftIO . atomically $ putTMVar v (M.toList s)+ TorrentCompleted ih -> do+ mp <- get+ let q = M.lookup ih mp+ ns <- maybe (fail "Unknown Torrent") return q+ assert (left ns == 0) (return ())+ 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+ updates <- liftIO . atomically $ do+ updates <- readTVar tv+ writeTVar tv []+ return updates+ mapM_ (\(PStat ih up down) -> do+ (u, d) <- liftIO $ readIORef r+ liftIO $ writeIORef r (u+up, d+down)+ modify (\s -> M.adjust (\st ->+ st { uploaded = (uploaded st) + up+ , downloaded = (downloaded st) + down }) ih s)) updates+
+ src/Process/Timer.hs view
@@ -0,0 +1,20 @@+-- | The timer module is responsible for timing in the project. With+-- the timer system, we can register timers in the future and then+-- we can get a tick triggered at the point in time where we need to+-- act. This allows us to postpone events into the future at a+-- designated time.+--+module Process.Timer+ ( registerSTM+ )++where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad.Trans++registerSTM :: MonadIO m => Int -> TChan a -> a -> m ThreadId+registerSTM secs c m = liftIO $ forkIO $ do+ threadDelay (secs * 1000000)+ atomically $ writeTChan c m
+ src/Process/TorrentManager.hs view
@@ -0,0 +1,116 @@+-- | The Manager Process - Manages the torrents and controls them+module Process.TorrentManager (+ -- * Types+ TorrentManagerMsg(..)+ -- * Channels+ , TorrentMgrChan+ -- * Interface+ , start+ )+where++import Control.Concurrent+import Control.Concurrent.STM++import Control.Monad.State+import Control.Monad.Reader++import qualified Data.ByteString as B+import Prelude hiding (log)++import Protocol.BCode as BCode+import Process+import qualified Process.Status as Status+import qualified Process.PeerMgr as PeerMgr+import qualified Process.FS as FSP+import qualified Process.PieceMgr as PieceMgr (start, createPieceDb)+import qualified Process.ChokeMgr as ChokeMgr (ChokeMgrChannel)+import qualified Process.Tracker as Tracker+import Channels+import FS+import Supervisor+import Torrent++data TorrentManagerMsg = AddedTorrent FilePath+ | RemovedTorrent FilePath+ deriving (Eq, Show)++type TorrentMgrChan = TChan [TorrentManagerMsg]++data CF = CF { tCh :: TorrentMgrChan+ , tStatusCh :: Status.StatusChannel+ , tStatusTV :: TVar [Status.PStat]+ , tPeerId :: PeerId+ , tPeerMgrCh :: PeerMgr.PeerMgrChannel+ , tChokeCh :: ChokeMgr.ChokeMgrChannel+ }++instance Logging CF where+ logName _ = "Process.TorrentManager"++data ST = ST { workQueue :: [TorrentManagerMsg] }+start :: TorrentMgrChan -- ^ Channel to watch for changes to torrents+ -> Status.StatusChannel+ -> TVar [Status.PStat]+ -> ChokeMgr.ChokeMgrChannel+ -> PeerId+ -> PeerMgr.PeerMgrChannel+ -> SupervisorChan+ -> IO ThreadId+start chan statusC stv chokeC pid peerC supC =+ spawnP (CF chan statusC stv pid peerC chokeC) (ST [])+ (catchP (forever pgm) (defaultStopHandler supC))+ where pgm = do startStop >> dirMsg+ dirMsg = do+ c <- asks tCh+ ls <- liftIO . atomically $ readTChan c+ modify (\s -> s { workQueue = ls ++ workQueue s })+ startStop = do+ q <- gets workQueue+ case q of+ [] -> return ()+ (AddedTorrent fp : rest) -> do+ debugP $ "Adding torrent file: " ++ fp+ _ <- startTorrent fp+ modify (\s -> s { workQueue = rest })+ startStop+ (RemovedTorrent _ : _) -> do+ errorP "Removal of torrents not yet supported :P"+ stopP++readTorrent :: FilePath -> Process CF ST BCode+readTorrent fp = do+ torrent <- liftIO $ B.readFile fp+ let bcoded = BCode.decode torrent+ case bcoded of+ Left err -> do liftIO $ print err+ stopP+ Right bc -> return bc++startTorrent :: FilePath -> Process CF ST ThreadId+startTorrent fp = do+ bc <- readTorrent fp+ fspC <- liftIO newTChanIO+ trackerC <- liftIO newTChanIO+ supC <- liftIO newTChanIO+ pieceMgrC <- liftIO newTChanIO+ chokeC <- asks tChokeCh+ statusC <- asks tStatusCh+ stv <- asks tStatusTV+ pid <- asks tPeerId+ pmC <- asks tPeerMgrCh+ (handles, haveMap, pieceMap) <- liftIO $ openAndCheckFile bc+ let left = bytesLeft haveMap pieceMap+ ti <- liftIO $ mkTorrentInfo bc+ pieceDb <- PieceMgr.createPieceDb haveMap pieceMap+ tid <- liftIO $ allForOne ("TorrentSup - " ++ fp)+ [ Worker $ FSP.start handles pieceMap fspC+ , Worker $ PieceMgr.start pieceMgrC fspC chokeC statusC pieceDb (infoHash ti)+ , Worker $ Tracker.start (infoHash ti) ti pid defaultPort statusC trackerC pmC+ ] supC+ liftIO . atomically $ writeTChan statusC $ Status.InsertTorrent (infoHash ti) left trackerC+ c <- asks tPeerMgrCh+ liftIO . atomically $ writeTChan c $ PeerMgr.NewTorrent (infoHash ti)+ (PeerMgr.TorrentLocal pieceMgrC fspC stv pieceMap)+ liftIO . atomically $ writeTChan trackerC Start+ return tid
+ src/Process/Tracker.hs view
@@ -0,0 +1,278 @@+-- | The TrackerP module is responsible for keeping in touch with the Tracker of a torrent.+-- The tracker is contacted periodically, and we exchange information with it. Specifically,+-- we tell the tracker how much we have downloaded, uploaded and what is left. We also+-- tell it about our current state (i.e., are we a seeder or a leecher?).+--+-- The tracker responds to us with a new set of Peers and general information about the+-- torrent in question. It may also respond with an error in which case we should present+-- it to the user.+--+module Process.Tracker+ ( start+ )+where++import Control.Applicative+import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad.Reader+import Control.Monad.State++import Data.Char (ord)+import Data.List (intersperse)+import qualified Data.ByteString as B++import Network+import Network.HTTP hiding (port)+import Network.URI hiding (unreserved)++import Numeric (showHex)+++import Protocol.BCode as BCode hiding (encode)+import Process+import Channels+import Supervisor+import Torrent+import qualified Process.Status as Status+import qualified Process.PeerMgr as PeerMgr+import qualified Process.Timer as Timer++++-- | The tracker state is used to tell the tracker our current state. In order+-- to output it correctly, we override the default show instance with the+-- version below. This may be wrong to do in the long run, but for now it works+-- fine.+--+-- The state is either started or stopped upon the client starting. The+-- tracker will create an entry for us if we tell it that we started, and it+-- will tear down this entry if we tell it that we stopped. It will know if+-- we are a seeder or a leecher based on how much data is left for us to+-- download.+--+-- the 'Completed' entry is used once in the lifetime of a torrent. It+-- explains to the tracker that we completed the torrent in question.+data TrackerEvent = Started | Stopped | Completed | Running+ deriving Eq++-- | The tracker will in general respond with a BCoded dictionary. In our world, this is+-- not the data structure we would like to work with. Hence, we parse the structure into+-- the ADT below.+data TrackerResponse = ResponseOk { newPeers :: [PeerMgr.Peer],+ completeR :: Maybe Integer,+ incompleteR :: Maybe Integer,+ timeoutInterval :: Integer,+ timeoutMinInterval :: Maybe Integer }+ | ResponseDecodeError B.ByteString+ | ResponseWarning B.ByteString+ | ResponseError B.ByteString++-- | If we fail to contact the tracker, we will wait for 15 minutes. The number is quite arbitrarily chosen+failTimerInterval :: Integer+failTimerInterval = 15 * 60++-- | Configuration of the tracker process+data CF = CF {+ statusPCh :: Status.StatusChannel+ , trackerMsgCh :: TrackerChannel+ , peerMgrCh :: PeerMgr.PeerMgrChannel+ , cfInfoHash :: InfoHash+ }++instance Logging CF where+ logName _ = "Process.Tracker"++-- | Internal state of the tracker CHP process+data ST = ST {+ torrentInfo :: TorrentInfo+ , peerId :: PeerId+ , state :: TrackerEvent+ , localPort :: PortID+ , nextTick :: Integer+ }++start :: InfoHash -> TorrentInfo -> PeerId -> PortID+ -> Status.StatusChannel -> TrackerChannel -> PeerMgr.PeerMgrChannel+ -> SupervisorChan -> IO ThreadId+start ih ti pid port statusC msgC pc supC =+ spawnP (CF statusC msgC pc ih) (ST ti pid Stopped port 0)+ (cleanupP (forever loop)+ (defaultStopHandler supC)+ stopEvent)+ where+ stopEvent :: Process CF ST ()+ stopEvent = do+ debugP "Stopping... telling tracker"+ modify (\s -> s { state = Stopped }) >> talkTracker+ loop :: Process CF ST ()+ loop = do+ ch <- asks trackerMsgCh+ msg <- liftIO . atomically $ readTChan ch+ debugP $ "Got tracker event"+ case msg of+ TrackerTick x ->+ do t <- gets nextTick+ when (x+1 == t) talkTracker+ Stop ->+ modify (\s -> s { state = Stopped }) >> talkTracker+ Start ->+ modify (\s -> s { state = Started }) >> talkTracker+ Complete ->+ modify (\s -> s { state = Completed }) >> talkTracker+ talkTracker = pokeTracker >>= timerUpdate++eventTransition :: Process CF ST ()+eventTransition = do+ st <- gets state+ modify (\s -> s { state = newS st})+ where newS st =+ case st of+ Running -> Running+ Stopped -> Stopped+ Completed -> Running+ Started -> Running++-- | Poke the tracker. It returns the new timer intervals to use+pokeTracker :: Process CF ST (Integer, Maybe Integer)+pokeTracker = do+ v <- liftIO $ newEmptyTMVarIO+ ih <- asks cfInfoHash+ asks statusPCh >>= (\ch -> liftIO . atomically $ writeTChan ch (Status.RequestStatus ih v))+ upDownLeft <- liftIO . atomically $ takeTMVar v+ url <- buildRequestURL upDownLeft+ debugP $ "Request URL: " ++ url+ uri <- case parseURI url of+ Nothing -> fail $ "Could not parse the url " ++ url+ Just u -> return u+ resp <- trackerRequest uri+ case resp of+ Left err -> do infoP $ "Tracker HTTP Error: " ++ err+ return (failTimerInterval, Just failTimerInterval)+ Right (ResponseWarning wrn) ->+ do infoP $ "Tracker Warning Response: " ++ fromBS wrn+ return (failTimerInterval, Just failTimerInterval)+ Right (ResponseError err) ->+ do infoP $ "Tracker Error Response: " ++ fromBS err+ return (failTimerInterval, Just failTimerInterval)+ Right (ResponseDecodeError err) ->+ do infoP $ "Response Decode error: " ++ fromBS err+ return (failTimerInterval, Just failTimerInterval)+ Right bc -> do+ c <- asks peerMgrCh+ liftIO . atomically $ writeTChan c (PeerMgr.PeersFromTracker ih $ newPeers bc)+ let trackerStats = Status.TrackerStat+ { Status.trackInfoHash = ih+ , Status.trackComplete = completeR bc+ , Status.trackIncomplete = incompleteR bc }+ asks statusPCh >>= \ch -> liftIO . atomically $ writeTChan ch trackerStats+ eventTransition+ return (timeoutInterval bc, timeoutMinInterval bc)++timerUpdate :: (Integer, Maybe Integer) -> Process CF ST ()+timerUpdate (timeout, _minTimeout) = do+ st <- gets state+ when (st == Running)+ (do t <- tick+ ch <- asks trackerMsgCh+ _ <- Timer.registerSTM (fromIntegral timeout) ch (TrackerTick t)+ debugP $ "Set timer to: " ++ show timeout)+ where tick = do t <- gets nextTick+ modify (\s -> s { nextTick = t + 1 })+ return t++-- Process a result dict into a tracker response object.+processResultDict :: BCode -> TrackerResponse+processResultDict d =+ case BCode.trackerError d of+ Just err -> ResponseError err+ Nothing -> case BCode.trackerWarning d of+ Just warn -> ResponseWarning warn+ Nothing -> case decodeOk of+ Nothing -> ResponseDecodeError . toBS $ "Could not decode response properly"+ Just rok -> rok+ where decodeOk =+ ResponseOk <$> (decodeIps <$> BCode.trackerPeers d)+ <*> (pure $ BCode.trackerComplete d)+ <*> (pure $ BCode.trackerIncomplete d)+ <*> (BCode.trackerInterval d)+ <*> (pure $ BCode.trackerMinInterval d)+++decodeIps :: B.ByteString -> [PeerMgr.Peer]+decodeIps = decodeIps' . fromBS++-- 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 = []++trackerRequest :: URI -> Process CF ST (Either String TrackerResponse)+trackerRequest uri =+ do resp <- liftIO $ simpleHTTP request+ case resp of+ Left x -> return $ Left ("Error connecting: " ++ show x)+ Right r ->+ case rspCode r of+ (2,_,_) ->+ case BCode.decode . toBS . rspBody $ r of+ Left pe -> return $ Left (show pe)+ Right bc -> do debugP $ "Response: " ++ BCode.prettyPrint bc+ return $ Right $ processResultDict bc+ (3,_,_) ->+ case findHeader HdrLocation r of+ Nothing -> return $ Left (show r)+ Just newURL -> case parseURI newURL of+ Nothing -> return $ Left (show newURL)+ Just u -> trackerRequest u+ _ -> return $ Left (show r)+ where request = Request {rqURI = uri,+ rqMethod = GET,+ rqHeaders = [],+ rqBody = ""}++-- Construct a new request URL. Perhaps this ought to be done with the HTTP client library+buildRequestURL :: Status.StatusState -> Process CF ST String+buildRequestURL ss = do ti <- gets torrentInfo+ hdrs <- headers+ let hl = concat $ hlist hdrs+ return $ concat [fromBS $ announceURL ti, "?", hl]+ where hlist x = intersperse "&" $ map (\(k,v) -> k ++ "=" ++ v) x+ headers = do+ s <- get+ p <- prt+ return $ [("info_hash", rfc1738Encode $ infoHash $ torrentInfo s),+ ("peer_id", rfc1738Encode $ peerId s),+ ("uploaded", show $ Status.uploaded ss),+ ("downloaded", show $ Status.downloaded ss),+ ("left", show $ Status.left ss),+ ("port", show p),+ ("compact", "1")] +++ (trackerfyEvent $ state s)+ prt :: Process CF ST Integer+ prt = do lp <- gets localPort+ case lp of+ PortNumber pnum -> return $ fromIntegral pnum+ _ -> do fail "Unknown port type"+ trackerfyEvent ev =+ case ev of+ Running -> []+ Completed -> [("event", "completed")]+ Started -> [("event", "started")]+ Stopped -> [("event", "stopped")]++-- Carry out URL-encoding of a string. Note that the clients seems to do it the wrong way+-- so we explicitly code it up here in the same wrong way, jlouis.+rfc1738Encode :: String -> String+rfc1738Encode = concatMap (\c -> if unreserved c then [c] else encode c)+ where unreserved = (`elem` chars)+ -- I killed ~ from this list as the Mainline client doesn't announce it - jlouis+ chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_./"+ encode :: Char -> String+ encode c = '%' : pHex c+ pHex c =+ let p = (showHex . ord $ c) ""+ in if length p == 1 then '0' : p else p
+ src/Protocol/BCode.hs view
@@ -0,0 +1,372 @@+-- | Add a module description here+-- also add descriptions to each function.+module Protocol.BCode +(+ BCode,+ Path(..),+ encode,+ -- encodeBS,+ decode,+ search,+ announce,+ comment,+ creationDate,+ info,+ hashInfoDict,+ infoLength,+ infoName,+ infoPieceLength,+ infoPieces,+ numberPieces,+ infoFiles,+ prettyPrint,++ trackerComplete,+ trackerIncomplete,+ trackerInterval,+ trackerMinInterval,+ trackerPeers,+ trackerWarning,+ trackerError,+ toBS,+ fromBS,+ --Tests+ testSuite)+where++import Control.Monad+import Control.Applicative hiding (many)+import qualified Data.ByteString.Lazy as L+import qualified Data.ByteString as B+import Data.Char++import Data.List+import qualified Data.Map as M+import Text.PrettyPrint.HughesPJ hiding (char)++import Data.Serialize+import Data.Serialize.Put+import Data.Serialize.Get++import Test.QuickCheck+import Test.Framework+import Test.Framework.Providers.QuickCheck2+import Test.Framework.Providers.HUnit+import Test.HUnit hiding (Path, Test)++import Digest+import TestInstance() -- for instances only++-- | BCode represents the structure of a bencoded file+data BCode = BInt Integer -- ^ An integer+ | BString B.ByteString -- ^ A string of bytes+ | BArray [BCode] -- ^ An array+ | BDict (M.Map B.ByteString BCode) -- ^ A key, value map+ deriving (Show, Eq)++instance Arbitrary BCode where+ arbitrary = sized bc'+ where bc' :: Int -> Gen BCode+ bc' 0 = oneof [BInt <$> arbitrary,+ BString <$> arbitrary]+ bc' n | n > 0 =+ oneof [BInt <$> arbitrary,+ BString <$> arbitrary,+ BArray <$> sequence (replicate n $ bc' (n `div` 8)),+ do keys <- vectorOf n arbitrary+ values <- sequence (replicate n $ bc' (n `div` 8))+ return $ BDict $ M.fromList $ zip keys values]++data Path = PString B.ByteString+ | PInt Int++toW8 :: Char -> Word8+toW8 = fromIntegral . ord++fromW8 :: Word8 -> Char+fromW8 = chr . fromIntegral++toBS :: String -> B.ByteString+toBS = B.pack . map toW8++fromBS :: B.ByteString -> String+fromBS = map fromW8 . B.unpack+++instance Serialize BCode where+ put (BInt i) = wrap 'i' 'e' $ putShow i+ put (BArray arr) = wrap 'l' 'e' . mapM_ put $ arr+ put (BDict mp) = wrap 'd' 'e' dict+ where dict = mapM_ encPair . M.toList $ mp+ encPair (k, v) = put (BString k) >> put v+ put (BString s) = do+ putShow (B.length s)+ putWord8 (toW8 ':')+ putByteString s+ + get = getBInt <|> getBArray <|> getBDict <|> getBString++-- | Get something wrapped in two Chars+getWrapped :: Char -> Char -> Get a -> Get a+getWrapped a b p = char a *> p <* char b++-- | Parses a BInt+getBInt :: Get BCode+getBInt = BInt . read <$> getWrapped 'i' 'e' intP+ where intP = ((:) <$> char '-' <*> getDigits) <|> getDigits++-- | Parses a BArray+getBArray :: Get BCode+getBArray = BArray <$> getWrapped 'l' 'e' (many get)+++-- | Parses a BDict+getBDict :: Get BCode+getBDict = BDict . M.fromList <$> getWrapped 'd' 'e' (many getPairs)+ where getPairs = do+ (BString s) <- getBString+ x <- get+ return (s,x)++-- | Parses a BString+getBString :: Get BCode+getBString = do+ count <- getDigits+ BString <$> ( char ':' *> getStr (read count :: Integer))+ where maxInt = fromIntegral (maxBound :: Int) :: Integer+ + getStr n | n >= 0 = B.concat <$> (sequence $ getStr' n)+ | otherwise = fail $ "read a negative length string, length: " ++ show n+ + getStr' n | n > maxInt = getByteString maxBound : getStr' (n-maxInt)+ | otherwise = [getByteString . fromIntegral $ n]+++-- | Get one or more digit characters+getDigits :: Get String+getDigits = many1 digit++-- | Returns a character if it is a digit, fails otherwise. uses isDigit.+digit :: Get Char+digit = do+ x <- getCharG+ if isDigit x+ then return x+ else fail $ "Expected digit, got: " ++ show x+++-- * Put helper functions++-- | Put an element, wrapped by two characters+wrap :: Char -> Char -> Put -> Put+wrap a b m = do+ putWord8 (toW8 a)+ m+ putWord8 (toW8 b)++-- | Put something as it is shown using @show@+putShow :: Show a => a -> Put+putShow = mapM_ put . show++-- * Get Helper functions++-- | Parse zero or items using a given parser+many :: Get a -> Get [a]+many p = many1 p `mplus` return []++-- | Parse one or more items using a given parser+many1 :: Get a -> Get [a]+many1 p = (:) <$> p <*> many p++-- | Parse a given character+char :: Char -> Get Char+char c = do+ x <- getCharG+ if x == c+ then return c+ else fail $ "Expected char: '" ++ c:"' got: '" ++ [x,'\'']++-- | Get a Char. Only works with single byte characters+getCharG :: Get Char+getCharG = fromW8 <$> getWord8++-- BCode helper functions++-- | Return the hash of the info-dict in a torrent file+hashInfoDict :: BCode -> IO Digest+hashInfoDict bc =+ do ih <- case info bc of+ Nothing -> fail "Could not find infoHash"+ Just x -> return x+ let encoded = encode ih+ digest $ L.fromChunks $ [encoded]+++toPS :: String -> Path+toPS = PString . toBS++{- Simple search function over BCoded data structures, general case. In practice, we+ will prefer some simpler mnemonics -}+search :: [Path] -> BCode -> Maybe BCode+search [] bc = Just bc+search (PInt i : rest) (BArray bs) | i < 0 || i > length bs = Nothing+ | otherwise = search rest (bs!!i)+search (PString s : rest) (BDict mp) = M.lookup s mp >>= search rest+search _ _ = Nothing++search' :: String -> BCode -> Maybe B.ByteString+search' str b = case search [toPS str] b of+ Nothing -> Nothing+ Just (BString s) -> Just s+ _ -> Nothing++searchStr :: String -> BCode -> Maybe B.ByteString+searchStr = search'++searchInt :: String -> BCode -> Maybe Integer+searchInt str b = case search [toPS str] b of+ Just (BInt i) -> Just i+ _ -> Nothing++searchInfo :: String -> BCode -> Maybe BCode+searchInfo str = search [toPS "info", toPS str]++{- Various accessors -}+announce, comment, creationDate :: BCode -> Maybe B.ByteString+announce = search' "announce"+comment = search' "comment"+creationDate = search' "creation date"++{- Tracker accessors -}+trackerComplete, trackerIncomplete, trackerInterval :: BCode -> Maybe Integer+trackerMinInterval :: BCode -> Maybe Integer+trackerComplete = searchInt "complete"+trackerIncomplete = searchInt "incomplete"+trackerInterval = searchInt "interval"+trackerMinInterval = searchInt "min interval"++trackerError, trackerWarning :: BCode -> Maybe B.ByteString+trackerError = searchStr "failure reason"+trackerWarning = searchStr "warning mesage"++trackerPeers :: BCode -> Maybe B.ByteString+trackerPeers = searchStr "peers"++info :: BCode -> Maybe BCode+info = search [toPS "info"]++infoName :: BCode -> Maybe B.ByteString+infoName bc = case search [toPS "info", toPS "name"] bc of+ Just (BString s) -> Just s+ _ -> Nothing++infoPieceLength ::BCode -> Maybe Integer+infoPieceLength bc = do BInt i <- search [toPS "info", toPS "piece length"] bc+ return i++infoLength :: BCode -> Maybe Integer+infoLength bc = maybe length2 Just length1+ where+ -- |info/length key for single-file torrent+ length1 = do BInt i <- search [toPS "info", toPS "length"] bc+ return i+ -- |length summed from files of multi-file torrent+ length2 = sum `fmap`+ map snd `fmap`+ infoFiles bc++infoPieces :: BCode -> Maybe [B.ByteString]+infoPieces b = do t <- searchInfo "pieces" b+ case t of+ BString str -> return $ sha1Split str+ _ -> mzero+ where sha1Split r | r == B.empty = []+ | otherwise = block : sha1Split rest+ where (block, rest) = B.splitAt 20 r++numberPieces :: BCode -> Maybe Int+numberPieces = fmap length . infoPieces++infoFiles :: BCode -> Maybe [([String], Integer)] -- ^[(filePath, fileLength)]+infoFiles bc = let mbFpath = fromBS `fmap` infoName bc+ mbLength = infoLength bc+ mbFiles = do BArray fileList <- searchInfo "files" bc+ return $ do fileDict@(BDict _) <- fileList+ let Just (BInt length) = search [toPS "length"] fileDict+ Just (BArray path) = search [toPS "path"] fileDict+ path' = map (\(BString s) -> fromBS s) path+ return (path', length)+ in case (mbFpath, mbLength, mbFiles) of+ (Just fpath, _, Just files) ->+ Just $+ map (\(path, length) ->+ (fpath:path, length)+ ) files+ (Just fpath, Just length, _) ->+ Just [([fpath], length)]+ (_, _, Just files) ->+ Just files+ _ ->+ Nothing++pp :: BCode -> Doc+pp bc =+ case bc of+ BInt i -> integer i+ BString s -> text (show s)+ BArray arr -> text "[" <+> (cat $ intersperse comma al) <+> text "]"+ where al = map pp arr+ BDict mp -> text "{" <+> cat (intersperse comma mpl) <+> text "}"+ where mpl = map (\(s, bc') -> text (fromBS s) <+> text "->" <+> pp bc') $ M.toList mp++prettyPrint :: BCode -> String+prettyPrint = render . pp+++toBDict :: [(String,BCode)] -> BCode+toBDict = BDict . M.fromList . map (\(k,v) -> ((toBS k),v))++toBString :: String -> BCode+toBString = BString . toBS+++-- TESTS+++testSuite :: Test+testSuite = testGroup "Protocol/BCode"+ [ testProperty "QC encode-decode/id" propEncodeDecodeId,+ testCase "HUnit encode-decode/id" testDecodeEncodeProp1 ]++propEncodeDecodeId :: BCode -> Bool+propEncodeDecodeId bc =+ let encoded = encode bc+ decoded = decode encoded+ in+ Right bc == decoded++testDecodeEncodeProp1 :: Assertion+testDecodeEncodeProp1 =+ let encoded = encode testData+ decoded = decode encoded+ in+ assertEqual "for encode/decode identify" (Right testData) decoded++testData :: [BCode]+testData = [BInt 123,+ BInt (-123),+ BString (toBS "Hello"),+ BString (toBS ['\NUL'..'\255']),+ BArray [BInt 1234567890+ ,toBString "a longer string with eieldei stuff to mess things up"+ ],+ toBDict [+ ("hello",BInt 3)+ ,("a key",toBString "and a value")+ ,("a sub dict",toBDict [+ ("some stuff",BInt 1)+ ,("some more stuff", toBString "with a string")+ ])+ ]+ ]+
+ src/Protocol/Wire.hs view
@@ -0,0 +1,263 @@+{-+ A parser and encoder for the BitTorrent wire protocol using the+ cereal package.+-}++{-# LANGUAGE EmptyDataDecls #-}++module Protocol.Wire+ ( Message(..)+ , encodePacket+ , decodeMsg+ , BitField+ , constructBitField+ -- Handshaking+ , initiateHandshake+ , receiveHandshake+ -- Tests+ , testSuite+ )+where++import Control.Applicative hiding (empty)+import Control.Monad++import Data.Monoid+import qualified Data.ByteString as B+import qualified Data.ByteString.Lazy as L++import Data.Serialize+import Data.Serialize.Put+import Data.Serialize.Get++import Data.Char+import System.IO+import System.Log.Logger++import Test.Framework+import Test.Framework.Providers.QuickCheck2+import Test.QuickCheck++import Torrent++------------------------------------------------------------++type BitField = L.ByteString++data Message = KeepAlive+ | Choke+ | Unchoke+ | Interested+ | NotInterested+ | Have PieceNum -- Int+ | BitField BitField+ | Request PieceNum Block+ | Piece PieceNum Int B.ByteString+ | Cancel PieceNum Block+ | Port Integer+ deriving (Eq, Show)++instance Arbitrary Message where+ arbitrary = oneof [return KeepAlive, return Choke, return Unchoke, return Interested,+ return NotInterested,+ Have <$> pos,+ BitField <$> arbitrary,+ Request <$> pos <*> arbitrary,+ Piece <$> pos <*> pos <*> arbitrary,+ Cancel <$> pos <*> arbitrary,+ Port <$> choose (0,16383)]+ where+ pos :: Gen Int+ pos = choose (0, 4294967296 - 1)+++-- | The Protocol header for the Peer Wire Protocol+protocolHeader :: String+protocolHeader = "BitTorrent protocol"++extensionBasis :: Word64+extensionBasis = 0++p8 :: Word8 -> Put+p8 = putWord8++p32be :: Integral a => a -> Put+p32be = putWord32be . fromIntegral++decodeMsg :: Get Message+decodeMsg = {-# SCC "decodeMsg" #-} get++encodePacket :: Message -> B.ByteString+encodePacket m = mconcat [szEnc, mEnc]+ where mEnc = encode m+ sz = B.length mEnc+ szEnc = runPut . p32be $ sz++instance Serialize Message where+ put KeepAlive = return ()+ put Choke = p8 0+ put Unchoke = p8 1+ put Interested = p8 2+ put NotInterested = p8 3+ put (Have pn) = p8 4 *> p32be pn+ put (BitField bf) = p8 5 *> putLazyByteString 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)++ get = getKA <|> getChoke+ <|> getUnchoke <|> getIntr+ <|> getNI <|> getHave+ <|> getBF <|> getReq+ <|> getPiece <|> getCancel+ <|> getPort++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)+getKA = do+ empty <- isEmpty+ if empty+ then return KeepAlive+ else fail "Non empty message - not a KeepAlive"++gw32 :: Integral a => Get a+gw32 = fromIntegral <$> getWord32be++byte :: Word8 -> Get Word8+byte w = do+ x <- lookAhead getWord8+ if x == w+ then getWord8+ else fail $ "Expected byte: " ++ show w ++ " got: " ++ show x++-- | Size of the protocol header+protocolHeaderSize :: Int+protocolHeaderSize = length protocolHeader++-- | Protocol handshake code. This encodes the protocol handshake part+protocolHandshake :: L.ByteString+protocolHandshake = L.fromChunks . map runPut $+ [putWord8 $ fromIntegral protocolHeaderSize,+ mapM_ (putWord8 . fromIntegral . ord) protocolHeader,+ putWord64be extensionBasis]++toBS :: String -> B.ByteString+toBS = B.pack . map toW8++toLBS :: String -> L.ByteString+toLBS = L.pack . map toW8++fromLBS :: L.ByteString -> String+fromLBS = map (chr . fromIntegral) . L.unpack++toW8 :: Char -> Word8+toW8 = fromIntegral . ord+++-- | Receive the header parts from the other end+receiveHeader :: Handle -> Int -> (InfoHash -> Bool)+ -> IO (Either String ([Capabilities], L.ByteString, InfoHash))+receiveHeader h sz ihTst = parseHeader `fmap` B.hGet h sz+ where parseHeader = runGet (headerParser ihTst)++headerParser :: (InfoHash -> Bool) -> Get ([Capabilities], L.ByteString, InfoHash)+headerParser ihTst = do+ hdSz <- getWord8+ when (fromIntegral hdSz /= protocolHeaderSize) $ fail "Wrong header size"+ protoString <- getByteString protocolHeaderSize+ when (protoString /= toBS protocolHeader) $ fail "Wrong protocol header"+ caps <- getWord64be+ ihR <- liftM fromLBS $ getLazyByteString 20+ unless (ihTst ihR) $ fail "Unknown InfoHash"+ pid <- getLazyByteString 20+ return (decodeCapabilities caps, pid, ihR)++data Capabilities+decodeCapabilities :: Word64 -> [Capabilities]+decodeCapabilities _ = []++-- | Initiate a handshake on a socket+initiateHandshake :: Handle -> PeerId -> InfoHash+ -> IO (Either String ([Capabilities], L.ByteString, InfoHash))+initiateHandshake handle peerid infohash = do+ debugM "Protocol.Wire" "Sending off handshake message"+ L.hPut handle msg+ hFlush handle+ debugM "Protocol.Wire" "Receiving handshake from other end"+ receiveHeader handle sz (== infohash)+ where msg = handShakeMessage peerid infohash+ sz = fromIntegral (L.length msg)++-- | Construct a default handshake message from a PeerId and an InfoHash+handShakeMessage :: PeerId -> InfoHash -> L.ByteString+handShakeMessage pid ih =+ L.fromChunks . map runPut $ [putLazyByteString protocolHandshake,+ putLazyByteString $ toLBS ih,+ putByteString . toBS $ pid]++-- | Receive a handshake on a socket+receiveHandshake :: Handle -> PeerId -> (InfoHash -> Bool)+ -> IO (Either String ([Capabilities], L.ByteString, InfoHash))+receiveHandshake h pid ihTst = do+ debugM "Protocol.Wire" "Receiving handshake from other end"+ r <- receiveHeader h 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+ return $ Right (caps, rpid, ih)+ where msg ih = handShakeMessage pid ih+ sz = fromIntegral (L.length $ msg "12345678901234567890") -- Dummy value+++-- | 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+ where m = map (`elem` pieces) [0..sz-1 + pad]+ pad = 8 - (sz `mod` 8)+ build [] = []+ build l = let (first, rest) = splitAt 8 l+ in if length first /= 8+ then error "Wront bitfield"+ else bytify first : build rest+ bytify [b7,b6,b5,b4,b3,b2,b1,b0] = sum [if b0 then 1 else 0,+ if b1 then 2 else 0,+ if b2 then 4 else 0,+ if b3 then 8 else 0,+ if b4 then 16 else 0,+ if b5 then 32 else 0,+ if b6 then 64 else 0,+ if b7 then 128 else 0]+ bytify _ = error "Bitfield construction failed"++--+-- -- TESTS+testSuite :: Test+testSuite = testGroup "Protocol/Wire"+ [ testProperty "QC encode-decode/id" propEncodeDecodeId]+++propEncodeDecodeId :: Message -> Bool+propEncodeDecodeId m =+ let encoded = encode m+ decoded = decode encoded+ in+ Right m == decoded++
+ src/RateCalc.hs view
@@ -0,0 +1,73 @@+-- | Rate calculation.+module RateCalc (+ -- * Types+ Rate+ -- * Interface+ , new+ , update+ , extractCount+ , extractRate+ )++where++import Data.Time.Clock++-- | A Rate is a record of information used for calculating the rate+data Rate = Rate+ { rate :: !Double -- ^ The current rate+ , bytes :: !Integer -- ^ The amount of bytes transferred since last rate extraction+ , count :: !Integer -- ^ The amount of bytes transferred since last count extraction+ , nextExpected :: !UTCTime -- ^ When is the next rate update expected+ , lastExt :: !UTCTime -- ^ When was the last rate update+ , rateSince :: !UTCTime -- ^ From where is the rate measured+ }++fudge :: NominalDiffTime+fudge = fromInteger 5 -- Seconds++maxRatePeriod :: NominalDiffTime+maxRatePeriod = fromInteger 20 -- Seconds++new :: UTCTime -> Rate+new t = Rate { rate = 0.0+ , bytes = 0+ , count = 0+ , nextExpected = addUTCTime fudge t+ , lastExt = addUTCTime (-fudge) t+ , rateSince = addUTCTime (-fudge) t+ }++-- | The call @update n rt@ updates the rate structure @rt@ with @n@ new bytes+update :: Integer -> Rate -> Rate+update n rt = rt { bytes = bytes rt + n+ , count = count rt + n}++-- | The call @extractRate t rt@ extracts the current rate from the rate structure and updates the rate+-- structures internal book-keeping+extractRate :: UTCTime -> Rate -> (Double, Rate)+extractRate t rt =+ let oldWindow :: Double+ oldWindow = realToFrac $ diffUTCTime (lastExt rt) (rateSince rt)+ newWindow :: Double+ newWindow = realToFrac $ diffUTCTime t (rateSince rt)+ n = bytes rt+ r = (rate rt * oldWindow + (fromIntegral n)) / newWindow+ expectN = min 5 (round $ (fromIntegral n / (max r 0.0001)))+ 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)+ })++-- | The call @extractCount rt@ extract the bytes transferred since last extraction+extractCount :: Rate -> (Integer, Rate)+extractCount rt = (count rt, rt { count = 0 })+
+ src/Supervisor.hs view
@@ -0,0 +1,184 @@+-- | Erlang style supervisors for Haskell.+-- Note that yet, these are not really good enough for using in other projects.+-- are currently subject to change until I figure out how a nice interface will+-- look like. At that moment they could be split off into their own package.+{-# LANGUAGE ScopedTypeVariables #-}+module Supervisor (+ -- * Types+ Child(..)+ , Children+ , SupervisorMsg(..)+ , SupervisorChan+ -- * Supervisor Initialization+ , allForOne+ , oneForOne+ -- * helper calls+ , pDie+ , defaultStopHandler+ )+where++import Control.Applicative+import Control.Concurrent+import Control.Concurrent.STM+import Control.Monad.State+import Control.Monad.Reader++import Prelude hiding (catch)++import Process++data Child = Supervisor (SupervisorChan -> IO ThreadId)+ | Worker (SupervisorChan -> IO ThreadId)++data SupervisorMsg = IAmDying ThreadId+ | PleaseDie ThreadId+ | SpawnNew Child++type SupervisorChan = TChan SupervisorMsg+type Children = [Child]++data ChildInfo = HSupervisor ThreadId+ | HWorker ThreadId+++pDie :: SupervisorChan -> IO ()+pDie supC = do+ tid <- myThreadId+ atomically $ writeTChan supC (IAmDying tid)++class SupervisorConf a where+ getParent :: a -> SupervisorChan+ getChan :: a -> SupervisorChan++data CFOFA = CFOFA { name :: String+ , chan :: SupervisorChan+ , parent :: SupervisorChan+ }++instance SupervisorConf CFOFA where+ getParent = parent+ getChan = chan++instance Logging CFOFA where+ logName = name++data STOFA = STOFA { childInfo :: [ChildInfo] }++-- | Run a set of processes and do it once in the sense that if someone dies,+-- no restart is attempted. We will just kill off everybody without any kind+-- of prejudice.+allForOne :: String -> Children -> SupervisorChan -> IO ThreadId+allForOne n children parentC = do+ c <- newTChanIO+ spawnP (CFOFA n c parentC) (STOFA []) (catchP startup+ (defaultStopHandler parentC))+ where+ startup = do+ childs <- mapM spawnChild children+ modify (\_ -> STOFA (reverse childs))+ forever eventLoop+ eventLoop = do+ mTid <- liftIO myThreadId+ pc <- asks parent+ ch <- asks chan+ m <- liftIO . atomically $+ (readTChan ch >>= return . Left) `orElse`+ (readTChan pc >>= return . Right)+ case m of+ Left ev -> case ev of+ IAmDying _tid -> do+ gets childInfo >>= mapM_ finChild+ t <- liftIO myThreadId+ asks parent >>= \c -> liftIO . atomically $ writeTChan c (IAmDying t)+ SpawnNew chld -> do+ nc <- spawnChild chld+ modify (\(STOFA cs) -> STOFA (nc : cs))+ _ -> fail "Impossible"+ Right ev -> case ev of+ PleaseDie tid | tid == mTid -> do+ gets childInfo >>= mapM_ finChild+ stopP+ _ -> return ()++data CFOFO = CFOFO { oName :: String+ , oChan :: SupervisorChan+ , oparent :: SupervisorChan+ }++instance SupervisorConf CFOFO where+ getParent = oparent+ getChan = oChan++instance Logging CFOFO where+ logName = oName++data STOFO = STOFO { oChildInfo :: [ChildInfo] }++-- | A One-for-one supervisor is called with @oneForOne children parentCh@. It will spawn and run+-- @children@ and be linked into the supervisor structure on @parentCh@. It returns the ThreadId+-- of the supervisor itself and the Channel of which it is the controller.+--+-- Should a process die, the one-for-one supervisor will do nothing about it. It will just record+-- the death and let the other processes keep running.+--+-- TODO: Restart policies.+oneForOne :: String -> Children -> SupervisorChan -> IO (ThreadId, SupervisorChan)+oneForOne n children parentC = do+ c <- newTChanIO+ t <- spawnP (CFOFO n c parentC) (STOFO []) (catchP startup+ (defaultStopHandler parentC))+ return (t, c)+ where+ startup :: Process CFOFO STOFO ()+ startup = do+ childs <- mapM spawnChild children+ modify (\_ -> STOFO (reverse childs))+ forever eventLoop+ eventLoop :: Process CFOFO STOFO ()+ eventLoop = do+ mTid <- liftIO myThreadId+ pc <- asks oparent+ ch <- asks oChan+ m <- liftIO . atomically $+ (readTChan ch >>= return . Left) `orElse`+ (readTChan pc >>= return . Right)+ case m of+ Left ev -> case ev of+ IAmDying tid -> pruneChild tid+ SpawnNew chld -> do nc <- spawnChild chld+ modify (\(STOFO cs) -> STOFO (nc : cs))+ _ -> fail "Impossible (2)"+ Right ev -> case ev of+ PleaseDie tid | tid == mTid -> do+ gets oChildInfo >>= mapM_ finChild+ stopP+ _ -> return ()+ pruneChild tid = modify (\(STOFO cs) -> STOFO (filter chk cs))+ where chk (HSupervisor t) = t == tid+ chk (HWorker t) = t == tid+++finChild :: SupervisorConf a => ChildInfo -> Process a b ()+finChild (HWorker tid) = liftIO $ killThread tid -- Make this call killP in Process?+finChild (HSupervisor tid) = do+ c <- getChan <$> ask+ liftIO . atomically $ writeTChan c (PleaseDie tid)++spawnChild :: SupervisorConf a => Child -> Process a b ChildInfo+spawnChild (Worker proc) = do+ c <- getChan <$> ask+ nc <- liftIO . atomically $ dupTChan c+ tid <- liftIO $ proc nc+ return $ HWorker tid+spawnChild (Supervisor proc) = do+ c <- getChan <$> ask+ nc <- liftIO . atomically $ dupTChan c+ tid <- liftIO $ proc nc+ return $ HSupervisor tid++defaultStopHandler :: SupervisorChan -> Process a b ()+defaultStopHandler supC = do+ t <- liftIO $ myThreadId+ liftIO . atomically $ writeTChan supC $ IAmDying t+
+ src/Test.hs view
@@ -0,0 +1,36 @@+{- The Test module provides an interface to test-framework. It is called when --tests+ - is supplied on the command line. This file only gathers together various test suites+ - from all over the rest of the code and then executes them via test-framework.+ -}+module Test (runTests) where++import System.Environment ( getArgs )+import Test.Framework+import Test.Framework.Providers.QuickCheck2++import qualified Data.PieceSet (testSuite)+import qualified Data.Queue (testSuite)+import qualified Protocol.BCode (testSuite)+import qualified Protocol.Wire (testSuite)+++runTests :: IO ()+runTests =+ do args <- filter (/= "--tests") `fmap` getArgs+ flip defaultMainWithArgs args+ [ testSuite+ , Data.Queue.testSuite+ , Data.PieceSet.testSuite+ , Protocol.BCode.testSuite+ , Protocol.Wire.testSuite+ ]++testSuite :: Test+testSuite = testGroup "Test test-framework"+ [ testProperty "reverse-reverse/id" prop_reversereverse ]++-- reversing twice a finite list, is the same as identity+prop_reversereverse :: [Int] -> Bool+prop_reversereverse s = (reverse . reverse) s == id s+ where _ = s :: [Int]+
+ src/TestInstance.hs view
@@ -0,0 +1,51 @@+-- Define a set of test instances of common types+-- Portions of this code is taken from "Real World Haskell"+module TestInstance+ ()+where++import Data.Word+import qualified Data.ByteString as B+import qualified Data.ByteString.Lazy as L++import System.Random+import Test.QuickCheck++integralRandomR :: (Integral a, Integral b, RandomGen g, Num b) => (a, b) -> g -> (b, g)+integralRandomR (a,b) g = case randomR (c,d) g of+ (x,h) -> (fromIntegral x, h)+ where (c,d) = (fromIntegral a :: Integer,+ fromIntegral b :: Integer)++instance Random Word32 where+ randomR = integralRandomR+ random = randomR (minBound, maxBound)++instance Arbitrary Word32 where+ arbitrary = arbitraryBoundedRandom++instance CoArbitrary Word32 where+ coarbitrary = coarbitraryIntegral++instance Random Word8 where+ randomR = integralRandomR+ random = randomR (minBound, maxBound)++instance Arbitrary Word8 where+ arbitrary = arbitraryBoundedRandom++instance CoArbitrary Word8 where+ coarbitrary = coarbitraryIntegral++instance Arbitrary L.ByteString where+ arbitrary = L.pack `fmap` arbitrary++instance CoArbitrary L.ByteString where+ coarbitrary = coarbitrary . L.unpack++instance Arbitrary B.ByteString where+ arbitrary = B.pack `fmap` arbitrary++instance CoArbitrary B.ByteString where+ coarbitrary = coarbitrary . B.unpack+
+ src/Torrent.hs view
@@ -0,0 +1,145 @@+-- | The following module is responsible for general types used+-- throughout the system.+module Torrent (+ -- * Types+ InfoHash+ , PeerId+ , AnnounceURL+ , TorrentState(..)+ , TorrentInfo(..)+ , PieceNum+ , PieceSize+ , PieceMap+ , PiecesDoneMap+ , PieceInfo(..)+ , BlockSize+ , Block(..)+ -- * Interface+ , determineState+ , bytesLeft+ , defaultBlockSize+ , defaultOptimisticSlots+ , defaultPort+ , haskellTorrentVersion+ , mkPeerId+ , mkTorrentInfo+ )+where++import Control.Applicative+import qualified Data.Foldable as F+import qualified Data.ByteString as B+import Data.List+import qualified Data.Map as M++import Network+import Numeric++import System.Random+import Test.QuickCheck++import Protocol.BCode+import Digest++-- | The type of Infohashes as used in torrents. These are identifiers+-- of torrents+type InfoHash = Digest++-- | The peerId is the ID of a client. It is used to identify clients+-- from each other+type PeerId = String++-- | The internal type of Announce URLs+type AnnounceURL = B.ByteString++-- | Internal type for a torrent. It identifies a torrent in various places of the system.+data TorrentInfo = TorrentInfo {+ infoHash :: InfoHash,+ pieceCount :: Int, -- Number of pieces in torrent+ announceURL :: AnnounceURL } deriving Show++data TorrentState = Seeding | Leeching+ deriving Show++-- PIECES+----------------------------------------------------------------------+type PieceNum = Int+type PieceSize = Int++data PieceInfo = PieceInfo {+ offset :: Integer, -- ^ Offset of the piece, might be greater than Int+ len :: Integer, -- ^ Length of piece; usually a small value+ digest :: String -- ^ Digest of piece; taken from the .torret file+ } deriving (Eq, Show)++type PieceMap = M.Map PieceNum PieceInfo++-- | The PiecesDoneMap is a map which is true if we have the piece and false otherwise+type PiecesDoneMap = M.Map PieceNum Bool++-- | Given what pieces that are done, return the current state of the client.+determineState :: PiecesDoneMap -> TorrentState+determineState pd | F.all (==True) pd = Seeding+ | otherwise = Leeching++-- | Return the amount of bytes left on a torrent given what pieces are done and the+-- map of the shape of the torrent in question.+bytesLeft :: PiecesDoneMap -> PieceMap -> Integer+bytesLeft done pm =+ M.foldWithKey (\k v accu ->+ case M.lookup k done of+ Just False -> (len v) + accu+ _ -> accu) 0 pm++-- BLOCKS+----------------------------------------------------------------------+type BlockSize = Int++data Block = Block { blockOffset :: Int -- ^ offset of this block within the piece+ , blockSize :: BlockSize -- ^ size of this block within the piece+ } deriving (Eq, Ord, Show)++instance Arbitrary Block where+ arbitrary = Block <$> pos <*> pos+ where pos = choose (0, 4294967296 - 1)+++defaultBlockSize :: BlockSize+defaultBlockSize = 16384 -- Bytes++-- | Default number of optimistic slots+defaultOptimisticSlots :: Int+defaultOptimisticSlots = 2++-- | Default port to communicate on+defaultPort :: PortID+defaultPort = PortNumber $ 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.+-- It provides a way for trackers to disallow versions of the client which are misbehaving.+haskellTorrentVersion :: String+haskellTorrentVersion = "d001"++-- | Convert a BCode block into its corresponding TorrentInfo block, perhaps+-- failing in the process.+mkTorrentInfo :: BCode -> IO TorrentInfo+mkTorrentInfo bc = do+ (ann, np) <- case queryInfo bc of Nothing -> fail "Could not create torrent info"+ Just x -> return x+ ih <- hashInfoDict bc+ return TorrentInfo { infoHash = ih, announceURL = ann, pieceCount = np }+ where+ queryInfo bc =+ do ann <- announce bc+ np <- numberPieces bc+ return (ann, np)++-- | Create a new PeerId for this client+mkPeerId :: StdGen -> PeerId+mkPeerId gen = header ++ take (20 - length header) ranString+ where randomList :: Int -> StdGen -> [Int]+ randomList n = take n . unfoldr (Just . random)+ rs = randomList 10 gen+ ranString = concatMap (\i -> showHex (abs i) "") rs+ header = "-HT" ++ haskellTorrentVersion ++ "-"
+ src/Tracer.hs view
@@ -0,0 +1,20 @@+module Tracer+ ( Tracer+ , new+ , trace+ )+where++data Tracer a = Tracer [a] [a] Int Int++instance Show a => Show (Tracer a) where+ show (Tracer cur old _ _) = show (cur ++ old)++new :: Int -> Tracer a+new x = Tracer [] [] 0 x++trace :: a -> Tracer a -> Tracer a+trace msg (Tracer cur old sz l)+ | sz == l = Tracer [msg] cur 0 l+ | otherwise = Tracer (msg : cur) old (sz+1) l+
+ src/Version.hs.in view
@@ -0,0 +1,8 @@+-- | Combinatorrent version+module Version (version) where++githead :: String+githead = "@GITHEAD@"++version :: String+version = githead