diff --git a/AUTHORS b/AUTHORS new file mode 100644 --- /dev/null +++ b/AUTHORS @@ -0,0 +1,5 @@ +Alex Mason +Astro +Jesper Louis Andersen +John Gunderman +Thomas Christensen diff --git a/Combinatorrent.cabal b/Combinatorrent.cabal new file mode 100644 --- /dev/null +++ b/Combinatorrent.cabal @@ -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 + + + diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -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 diff --git a/Setup.lhs b/Setup.lhs new file mode 100644 --- /dev/null +++ b/Setup.lhs @@ -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" diff --git a/configure b/configure new file mode 100644 --- /dev/null +++ b/configure @@ -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 + + diff --git a/src/Channels.hs b/src/Channels.hs new file mode 100644 --- /dev/null +++ b/src/Channels.hs @@ -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 diff --git a/src/Combinatorrent.hs b/src/Combinatorrent.hs new file mode 100644 --- /dev/null +++ b/src/Combinatorrent.hs @@ -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 () + diff --git a/src/Data/PendingSet.hs b/src/Data/PendingSet.hs new file mode 100644 --- /dev/null +++ b/src/Data/PendingSet.hs @@ -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 + diff --git a/src/Data/PieceSet.hs b/src/Data/PieceSet.hs new file mode 100644 --- /dev/null +++ b/src/Data/PieceSet.hs @@ -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) + + diff --git a/src/Data/Queue.hs b/src/Data/Queue.hs new file mode 100644 --- /dev/null +++ b/src/Data/Queue.hs @@ -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 diff --git a/src/Digest.hs b/src/Digest.hs new file mode 100644 --- /dev/null +++ b/src/Digest.hs @@ -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 diff --git a/src/FS.hs b/src/FS.hs new file mode 100644 --- /dev/null +++ b/src/FS.hs @@ -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 diff --git a/src/Process.hs b/src/Process.hs new file mode 100644 --- /dev/null +++ b/src/Process.hs @@ -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 + diff --git a/src/Process/ChokeMgr.hs b/src/Process/ChokeMgr.hs new file mode 100644 --- /dev/null +++ b/src/Process/ChokeMgr.hs @@ -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) + diff --git a/src/Process/Console.hs b/src/Process/Console.hs new file mode 100644 --- /dev/null +++ b/src/Process/Console.hs @@ -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) + diff --git a/src/Process/DirWatcher.hs b/src/Process/DirWatcher.hs new file mode 100644 --- /dev/null +++ b/src/Process/DirWatcher.hs @@ -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) + diff --git a/src/Process/FS.hs b/src/Process/FS.hs new file mode 100644 --- /dev/null +++ b/src/Process/FS.hs @@ -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 diff --git a/src/Process/Listen.hs b/src/Process/Listen.hs new file mode 100644 --- /dev/null +++ b/src/Process/Listen.hs @@ -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) + diff --git a/src/Process/Peer.hs b/src/Process/Peer.hs new file mode 100644 --- /dev/null +++ b/src/Process/Peer.hs @@ -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 diff --git a/src/Process/Peer/Receiver.hs b/src/Process/Peer/Receiver.hs new file mode 100644 --- /dev/null +++ b/src/Process/Peer/Receiver.hs @@ -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 diff --git a/src/Process/Peer/Sender.hs b/src/Process/Peer/Sender.hs new file mode 100644 --- /dev/null +++ b/src/Process/Peer/Sender.hs @@ -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 diff --git a/src/Process/Peer/SenderQ.hs b/src/Process/Peer/SenderQ.hs new file mode 100644 --- /dev/null +++ b/src/Process/Peer/SenderQ.hs @@ -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) }) diff --git a/src/Process/PeerMgr.hs b/src/Process/PeerMgr.hs new file mode 100644 --- /dev/null +++ b/src/Process/PeerMgr.hs @@ -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" diff --git a/src/Process/PieceMgr.hs b/src/Process/PieceMgr.hs new file mode 100644 --- /dev/null +++ b/src/Process/PieceMgr.hs @@ -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) + diff --git a/src/Process/Status.hs b/src/Process/Status.hs new file mode 100644 --- /dev/null +++ b/src/Process/Status.hs @@ -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 + diff --git a/src/Process/Timer.hs b/src/Process/Timer.hs new file mode 100644 --- /dev/null +++ b/src/Process/Timer.hs @@ -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 diff --git a/src/Process/TorrentManager.hs b/src/Process/TorrentManager.hs new file mode 100644 --- /dev/null +++ b/src/Process/TorrentManager.hs @@ -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 diff --git a/src/Process/Tracker.hs b/src/Process/Tracker.hs new file mode 100644 --- /dev/null +++ b/src/Process/Tracker.hs @@ -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 diff --git a/src/Protocol/BCode.hs b/src/Protocol/BCode.hs new file mode 100644 --- /dev/null +++ b/src/Protocol/BCode.hs @@ -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") + ]) + ] + ] + diff --git a/src/Protocol/Wire.hs b/src/Protocol/Wire.hs new file mode 100644 --- /dev/null +++ b/src/Protocol/Wire.hs @@ -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 + + diff --git a/src/RateCalc.hs b/src/RateCalc.hs new file mode 100644 --- /dev/null +++ b/src/RateCalc.hs @@ -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 }) + diff --git a/src/Supervisor.hs b/src/Supervisor.hs new file mode 100644 --- /dev/null +++ b/src/Supervisor.hs @@ -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 + diff --git a/src/Test.hs b/src/Test.hs new file mode 100644 --- /dev/null +++ b/src/Test.hs @@ -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] + diff --git a/src/TestInstance.hs b/src/TestInstance.hs new file mode 100644 --- /dev/null +++ b/src/TestInstance.hs @@ -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 + diff --git a/src/Torrent.hs b/src/Torrent.hs new file mode 100644 --- /dev/null +++ b/src/Torrent.hs @@ -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 ++ "-" diff --git a/src/Tracer.hs b/src/Tracer.hs new file mode 100644 --- /dev/null +++ b/src/Tracer.hs @@ -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 + diff --git a/src/Version.hs.in b/src/Version.hs.in new file mode 100644 --- /dev/null +++ b/src/Version.hs.in @@ -0,0 +1,8 @@ +-- | Combinatorrent version +module Version (version) where + +githead :: String +githead = "@GITHEAD@" + +version :: String +version = githead