-- | Peer proceeses
{-# LANGUAGE ScopedTypeVariables #-}
module Process.Peer (
-- * Types
PeerMessage(..)
-- * Interface
, Process.Peer.start
-- * Tests
, Process.Peer.testSuite
)
where
import Control.Applicative
import Control.Concurrent
import Control.Concurrent.STM
import Control.DeepSeq
import Control.Exception
import Control.Monad.State
import Control.Monad.Reader
import Prelude hiding (catch, log)
import Data.Array
import Data.Bits
import qualified Data.ByteString as B
import Data.Function (on)
import qualified Data.PieceSet as PS
import Data.Maybe
import Data.Monoid(Monoid(..), Last(..))
import Data.Set as S hiding (map)
import Data.Time.Clock
import Data.Word
import Network.Socket hiding (KeepAlive)
import Test.Framework
import Test.Framework.Providers.HUnit
import Test.HUnit hiding (Path, Test, assert)
import Channels
import Digest
import Process
import Process.FS
import Process.PieceMgr
import RateCalc as RC
import Process.Status
import qualified Process.ChokeMgr as ChokeMgr (RateTVar, PeerRateInfo(..))
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 :: Socket -> [Capabilities] -> MgrChannel -> ChokeMgr.RateTVar -> PieceMgrChannel
-> FSPChannel -> TVar [PStat] -> PieceMap -> Int -> InfoHash
-> IO Children
start s caps pMgrC rtv pieceMgrC fsC stv pm nPieces ih = do
queueC <- newTChanIO
senderMV <- newEmptyTMVarIO
receiverC <- newTChanIO
sendBWC <- newTChanIO
return [Worker $ Sender.start s senderMV,
Worker $ SenderQ.start caps queueC senderMV sendBWC fsC,
Worker $ Receiver.start s receiverC,
Worker $ peerP caps pMgrC rtv pieceMgrC pm nPieces
queueC receiverC sendBWC stv ih]
-- INTERNAL FUNCTIONS
----------------------------------------------------------------------
data CF = CF { inCh :: TChan (Message, Integer)
, outCh :: TChan SenderQ.SenderQMsg
, peerMgrCh :: MgrChannel
, pieceMgrCh :: PieceMgrChannel
, peerCh :: PeerChannel
, sendBWCh :: BandwidthChannel
, timerCh :: TChan ()
, statTV :: TVar [PStat]
, rateTV :: ChokeMgr.RateTVar
, pcInfoHash :: InfoHash
, pieceMap :: !PieceMap
, piecesDoneTV :: TMVar [PieceNum]
, interestTV :: TMVar Bool
, grabBlockTV :: TMVar Blocks
, extConf :: ExtensionConfig
}
instance Logging CF where
logName _ = "Process.Peer"
data ST = ST { 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
, missingPieces :: !Int -- ^ Tracks the number of pieces the peer misses before seeding
, upRate :: !Rate -- ^ Upload rate towards the peer (estimated)
, downRate :: !Rate -- ^ Download rate from the peer (estimated)
, runningEndgame :: !Bool -- ^ True if we are in endgame
, lastMessage :: !Int
}
data ExtensionConfig = ExtensionConfig
{ handleHaveAll :: Last (Process CF ST ())
, handleHaveNone :: Last (Process CF ST ())
, handleBitfield :: Last (BitField -> Process CF ST ())
, handleSuggest :: Last (PieceNum -> Process CF ST ())
, handleAllowedFast :: Last (PieceNum -> Process CF ST ())
, handleRejectRequest :: Last (PieceNum -> Block -> Process CF ST ())
, handleExtendedMsg :: Last (Int -> B.ByteString -> Process CF ST ())
, handleRequestMsg :: Last (PieceNum -> Block -> Process CF ST ())
, handleChokeMsg :: Last (Process CF ST ())
, handleCancelMsg :: Last (PieceNum -> Block -> Process CF ST ())
, handlePieceMsg :: Last (PieceNum -> Int -> B.ByteString -> Process CF ST ())
}
emptyExtensionConfig :: ExtensionConfig
emptyExtensionConfig = ExtensionConfig
mempty mempty mempty mempty mempty mempty mempty mempty mempty mempty mempty
appendEConfig :: ExtensionConfig -> ExtensionConfig -> ExtensionConfig
appendEConfig a b =
ExtensionConfig {
handleHaveAll = app handleHaveAll a b
, handleHaveNone = app handleHaveNone a b
, handleBitfield = app handleBitfield a b
, handleSuggest = app handleSuggest a b
, handleAllowedFast = app handleAllowedFast a b
, handleRejectRequest = app handleRejectRequest a b
, handleExtendedMsg = app handleExtendedMsg a b
, handleRequestMsg = app handleRequestMsg a b
, handleChokeMsg = app handleChokeMsg a b
, handleCancelMsg = app handleCancelMsg a b
, handlePieceMsg = app handlePieceMsg a b
}
where app f = mappend `on` f
instance Monoid ExtensionConfig where
mempty = emptyExtensionConfig
mappend = appendEConfig
-- | Constructor for 'Last' values.
ljust :: a -> Last a
ljust = Last . Just
-- | Deconstructor for 'Last' values.
fromLJ :: (ExtensionConfig -> Last a) -- ^ Field to access.
-> ExtensionConfig -- ^ Default to use.
-> a
fromLJ f cfg = case f cfg of
Last Nothing -> fromLJ f extensionBase
Last (Just a) -> a
extensionBase :: ExtensionConfig
extensionBase = ExtensionConfig
(ljust errorHaveAll)
(ljust errorHaveNone)
(ljust bitfieldMsg)
(ljust errorSuggest)
(ljust errorAllowedFast)
(ljust errorRejectRequest)
(ljust errorExtendedMsg)
(ljust requestMsg)
(ljust chokeMsg)
(ljust cancelBlock)
(ljust pieceMsg)
fastExtension :: ExtensionConfig
fastExtension = ExtensionConfig
(ljust haveAllMsg)
(ljust haveNoneMsg)
(ljust bitfieldMsg)
(ljust ignoreSuggest)
(ljust ignoreAllowedFast)
(ljust rejectMsg)
mempty
(ljust requestFastMsg)
(ljust fastChokeMsg)
(ljust fastCancelBlock)
(ljust fastPieceMsg)
ignoreSuggest :: PieceNum -> Process CF ST ()
ignoreSuggest _ = debugP "Ignoring SUGGEST message"
ignoreAllowedFast :: PieceNum -> Process CF ST ()
ignoreAllowedFast _ = debugP "Ignoring ALLOWEDFAST message"
errorHaveAll :: Process CF ST ()
errorHaveAll = do
errorP "Received a HAVEALL, but the extension is not enabled"
stopP
errorHaveNone :: Process CF ST ()
errorHaveNone = do
errorP "Received a HAVENONE, but the extension is not enabled"
stopP
errorSuggest :: PieceNum -> Process CF ST ()
errorSuggest _ = do
errorP "Received a SUGGEST PIECE, but the extension is not enabled"
stopP
errorAllowedFast :: PieceNum -> Process CF ST ()
errorAllowedFast _ = do
errorP "Received a ALLOWEDFAST, but the extension is not enabled"
stopP
errorRejectRequest :: PieceNum -> Block -> Process CF ST ()
errorRejectRequest _ _ = do
errorP "Received a REJECT REQUEST, but the extension is not enabled"
stopP
errorExtendedMsg :: Int -> B.ByteString -> Process CF ST ()
errorExtendedMsg _ _ = do
errorP "Received an EXTENDED MESSAGE, but the extension is not enabled"
stopP
peerP :: [Capabilities] -> MgrChannel -> ChokeMgr.RateTVar -> PieceMgrChannel -> PieceMap -> Int
-> TChan SenderQ.SenderQMsg -> TChan (Message, Integer) -> BandwidthChannel
-> TVar [PStat] -> InfoHash
-> SupervisorChannel -> IO ThreadId
peerP caps pMgrC rtv pieceMgrC pm nPieces outBound inBound sendBWC stv ih supC = do
ch <- newTChanIO
tch <- newTChanIO
ct <- getCurrentTime
pdtmv <- newEmptyTMVarIO
intmv <- newEmptyTMVarIO
gbtmv <- newEmptyTMVarIO
pieceSet <- PS.new nPieces
let cs = configCapabilities caps
spawnP (CF inBound outBound pMgrC pieceMgrC ch sendBWC tch stv rtv ih pm
pdtmv intmv gbtmv cs)
(ST True False S.empty True False pieceSet nPieces (RC.new ct) (RC.new ct) False 0)
({-# SCC "PeerControl" #-}
cleanupP (startup nPieces) (defaultStopHandler supC) cleanup)
configCapabilities :: [Capabilities] -> ExtensionConfig
configCapabilities caps =
mconcat [mempty, if Fast `elem` caps then fastExtension else mempty]
startup :: Int -> Process CF ST ()
startup nPieces = do
tid <- liftIO $ myThreadId
ch <- asks peerCh
pmc <- asks peerMgrCh
ih <- asks pcInfoHash
liftIO . atomically $ writeTChan pmc $ Connect ih tid ch
pieces <- getPiecesDone
outChan $ SenderQ.SenderQM $ BitField (constructBitField nPieces pieces)
-- Install the StatusP timer
c <- asks timerCh
_ <- registerSTM 5 c ()
eventLoop
cleanup :: Process CF ST ()
cleanup = do
t <- liftIO myThreadId
pieces <- gets peerPieces >>= PS.toList
ch2 <- asks peerMgrCh
msgPieceMgr (PeerUnhave pieces)
liftIO . atomically $ writeTChan ch2 (Disconnect t)
readOp :: Process CF ST 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 :: Process CF ST ()
eventLoop = do
op <- readOp
case op of
PeerMsgEvt (m, sz) -> peerMsg m sz
ChokeMgrEvt m -> chokeMgrMsg m
UpRateEvent up -> do s <- get
u <- return $ RC.update up $ upRate s
put $! s { upRate = u }
TimerEvent -> timerTick
eventLoop
data Operation = PeerMsgEvt (Message, Integer)
| ChokeMgrEvt PeerMessage
| TimerEvent
| UpRateEvent Integer
-- | Return a list of pieces which are currently done by us
getPiecesDone :: Process CF ST [PieceNum]
getPiecesDone = do
c <- asks piecesDoneTV
msgPieceMgr (GetDone c)
liftIO $ do atomically $ takeTMVar c
-- | Process an event from the Choke Manager
chokeMgrMsg :: PeerMessage -> Process CF ST ()
chokeMgrMsg msg = do
case msg of
PieceCompleted pn -> do
debugP "Telling about Piece Completion"
outChan $ SenderQ.SenderQM $ Have pn
considerInterest
ChokePeer -> do choking <- gets weChoke
when (not choking)
(do t <- liftIO myThreadId
debugP $ "Pid " ++ show t ++ " choking"
outChan $ SenderQ.SenderOChoke
modify (\s -> s {weChoke = True}))
UnchokePeer -> do choking <- gets weChoke
when choking
(do t <- liftIO myThreadId
debugP $ "Pid " ++ show t ++ " unchoking"
outChan $ SenderQ.SenderQM Unchoke
modify (\s -> s {weChoke = False}))
CancelBlock pn blk -> do
cf <- asks extConf
fromLJ handleCancelMsg cf pn blk
cancelBlock :: PieceNum -> Block -> Process CF ST ()
cancelBlock pn blk = do
s <- get
put $! s { blockQueue = S.delete (pn, blk) $ blockQueue s }
outChan $ SenderQ.SenderQRequestPrune pn blk
fastCancelBlock :: PieceNum -> Block -> Process CF ST ()
fastCancelBlock pn blk = do
bq <- gets blockQueue
when (S.member (pn, blk) bq)
$ outChan $ SenderQ.SenderQRequestPrune pn blk
checkKeepAlive :: Process CF ST ()
checkKeepAlive = do
lm <- gets lastMessage
if lm >= 24
then do outChan $ SenderQ.SenderQM KeepAlive
else let inc = succ lm
in inc `seq` modify (\st -> st { lastMessage = inc })
-- 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 CF ST ()
timerTick = do
mTid <- liftIO myThreadId
checkKeepAlive
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, ChokeMgr.PRI {
ChokeMgr.peerUpRate = up,
ChokeMgr.peerDownRate = down,
ChokeMgr.peerInterested = i,
ChokeMgr.peerSeeding = seed,
ChokeMgr.peerChokingUs = pchoke }) : q)
-- Tell the Status Process about our progress
let (upCnt, nuRate) = RC.extractCount $ nur
(downCnt, ndRate) = RC.extractCount $ ndr
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 })
chokeMsg :: Process CF ST ()
chokeMsg = do
putbackBlocks
s <- get
put $! s { peerChoke = True }
fastChokeMsg :: Process CF ST ()
fastChokeMsg = do
s <- get
put $! s { peerChoke = True}
unchokeMsg :: Process CF ST ()
unchokeMsg = do
s <- get
put $! s { peerChoke = False }
fillBlocks
-- | Process an Message from the peer in the other end of the socket.
peerMsg :: Message -> Integer -> Process CF ST ()
peerMsg msg sz = do
modify (\s -> s { downRate = RC.update sz $ downRate s})
case msg of
KeepAlive -> return ()
Choke -> asks extConf >>= fromLJ handleChokeMsg
Unchoke -> unchokeMsg
Interested -> modify (\s -> s { peerInterested = True })
NotInterested -> modify (\s -> s { peerInterested = False })
Have pn -> haveMsg pn
BitField bf -> bitfieldMsg bf
Request pn blk -> do cf <- asks extConf
fromLJ handleRequestMsg cf pn blk
Piece n os bs -> do cf <- asks extConf
fromLJ handlePieceMsg cf n os bs
fillBlocks
Cancel pn blk -> cancelMsg pn blk
Port _ -> return () -- No DHT yet, silently ignore
HaveAll -> fromLJ handleHaveAll =<< asks extConf
HaveNone -> fromLJ handleHaveNone =<< asks extConf
Suggest pn -> do cf <- asks extConf
fromLJ handleSuggest cf pn
AllowedFast pn -> do cf <- asks extConf
fromLJ handleAllowedFast cf pn
RejectRequest pn blk -> do cf <- asks extConf
fromLJ handleRejectRequest cf pn blk
ExtendedMsg idx bs -> do cf <- asks extConf
fromLJ handleExtendedMsg cf idx bs
-- | Put back blocks for other peer processes to grab. This is done whenever
-- the peer chokes us, or if we die by an unknown cause.
putbackBlocks :: Process CF ST ()
putbackBlocks = do
blks <- gets blockQueue
msgPieceMgr (PutbackBlocks (S.toList blks))
modify (\s -> s { blockQueue = S.empty })
-- | Process a HAVE message from the peer. Note we also update interest as a side effect
haveMsg :: PieceNum -> Process CF ST ()
haveMsg pn = do
pm <- asks pieceMap
let (lo, hi) = bounds pm
if pn >= lo && pn <= hi
then do PS.insert pn =<< gets peerPieces
debugP $ "Peer has pn: " ++ show pn
msgPieceMgr (PeerHave [pn])
decMissingCounter 1
considerInterest
fillBlocks
else do warningP "Unknown Piece"
stopP
-- True if the peer is a seeder
isASeeder :: Process CF ST Bool
isASeeder = do sdr <- gets missingPieces
return $! sdr == 0
-- Decrease the counter of missing pieces for the peer
decMissingCounter :: Int -> Process CF ST ()
decMissingCounter n = do
modify (\s -> s { missingPieces = missingPieces s - n})
m <- gets missingPieces
when (m == 0) assertSeeder
-- Assert that the peer is a seeder
assertSeeder :: Process CF ST ()
assertSeeder = do
ok <- liftM2 (==) (gets peerPieces >>= PS.size) (succ . snd . bounds <$> asks pieceMap)
assert ok (return ())
-- | Process a BITFIELD message from the peer. Side effect: Consider Interest.
bitfieldMsg :: BitField -> Process CF ST ()
bitfieldMsg bf = do
pieces <- gets peerPieces
piecesNull <- PS.null pieces
if piecesNull
-- TODO: Don't trust the bitfield
then do nPieces <- succ . snd . bounds <$> asks pieceMap
pp <- createPeerPieces nPieces bf
modify (\s -> s { peerPieces = pp })
peerLs <- PS.toList pp
msgPieceMgr (PeerHave peerLs)
decMissingCounter (length peerLs)
considerInterest
else do infoP "Got out of band Bitfield request, dying"
stopP
haveAllNoneMsg :: String -> Bool -> Process CF ST ()
haveAllNoneMsg ty a = do
pieces <- gets peerPieces
piecesNull <- PS.null pieces
if piecesNull
then do nPieces <- succ . snd . bounds <$> asks pieceMap
pp <- createAllPieces nPieces a
modify (\s -> s { peerPieces = pp})
peerLs <- PS.toList pp
msgPieceMgr (PeerHave peerLs)
decMissingCounter (length peerLs)
considerInterest
else do infoP $ "Got out of band " ++ ty ++ " request, dying"
stopP
haveNoneMsg :: Process CF ST ()
haveNoneMsg = haveAllNoneMsg "HaveNone" False
haveAllMsg :: Process CF ST ()
haveAllMsg = haveAllNoneMsg "HaveAll" True
-- | Process a request message from the Peer
requestMsg :: PieceNum -> Block -> Process CF ST ()
requestMsg pn blk = do
choking <- gets weChoke
unless choking
(do debugP $ "Peer requested: " ++ show pn ++ "(" ++ show blk ++ ")"
outChan $ SenderQ.SenderQPiece pn blk)
requestFastMsg :: PieceNum -> Block -> Process CF ST ()
requestFastMsg pn blk = do
choking <- gets weChoke
debugP $ "Peer fastRequested: " ++ show pn ++ "(" ++ show blk ++ ")"
if choking
then outChan $ SenderQ.SenderQM (RejectRequest pn blk)
else outChan $ SenderQ.SenderQPiece pn blk
-- | Handle a Piece Message incoming from the peer
pieceMsg :: PieceNum -> Int -> B.ByteString -> Process CF ST ()
pieceMsg pn offs bs = pieceMsg' pn offs bs >> return ()
fastPieceMsg :: PieceNum -> Int -> B.ByteString -> Process CF ST ()
fastPieceMsg pn offs bs = do
r <- pieceMsg' pn offs bs
unless r
(do infoP "Peer sent out-of-band piece we did not request, closing"
stopP)
pieceMsg' :: PieceNum -> Int -> B.ByteString -> Process CF ST Bool
pieceMsg' n os bs = do
let sz = B.length bs
blk = Block os sz
e = (n, blk)
q <- gets blockQueue
-- When e is not a member, the piece may be stray, so ignore it.
-- Perhaps print something here.
if S.member e q
then do storeBlock n blk bs
bq <- gets blockQueue >>= return . S.delete e
s <- get
bq `deepseq` put $! s { blockQueue = bq }
return True
else return False
rejectMsg :: PieceNum -> Block -> Process CF ST ()
rejectMsg pn blk = do
let e = (pn, blk)
q <- gets blockQueue
if S.member e q
then do
msgPieceMgr (PutbackBlocks [e])
s <- get
put $! s { blockQueue = S.delete e q }
else do
infoP "Peer rejected piece/block we never requested, stopping"
stopP
-- | Handle a cancel message from the peer
cancelMsg :: PieceNum -> Block -> Process CF ST ()
cancelMsg n blk = outChan $ SenderQ.SenderQCancel n blk
-- | 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 CF ST ()
considerInterest = do
c <- asks interestTV
pcs <- gets peerPieces
msgPieceMgr (AskInterested pcs c)
interested <- liftIO $ do atomically $ takeTMVar c
if interested
then do oldI <- gets weInterested
when (interested /= oldI)
(do modify (\s -> s { weInterested = True })
debugP "We are interested"
outChan $ SenderQ.SenderQM Interested)
else do oldI <- gets weInterested
when (interested /= oldI)
(do modify (\s -> s { weInterested = False})
debugP "We are not interested"
outChan $ SenderQ.SenderQM NotInterested)
-- | Try to fill up the block queue at the peer. The reason we pipeline a
-- number of blocks is to get around the line delay present on the internet.
fillBlocks :: Process CF ST ()
fillBlocks = do
choked <- gets peerChoke
interested <- gets weInterested
when (not choked && interested) checkWatermark
-- | check the current Watermark level. If we are below the lower one, then
-- fill till the upper one. This in turn keeps the pipeline of pieces full as
-- long as the peer is interested in talking to us.
-- TODO: Decide on a queue size based on the current download rate.
checkWatermark :: Process CF ST ()
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)
queuePieces toQueue)
-- These three values are chosen rather arbitrarily at the moment.
loMark :: Int
loMark = 5
hiMark :: Int
hiMark = 25
-- Low mark when running in endgame mode
endgameLoMark :: Int
endgameLoMark = 1
-- | Queue up pieces for retrieval at the Peer
queuePieces :: [(PieceNum, Block)] -> Process CF ST ()
queuePieces toQueue = do
s <- get
let bq = blockQueue s
q <- forM toQueue
(\(p, b) -> do
if S.member (p, b) bq
then return Nothing -- Ignore pieces which are already in queue
else do outChan $ SenderQ.SenderQM $ Request p b
return $ Just (p, b))
put $! s { blockQueue = S.union bq (S.fromList $ catMaybes q) }
-- | Tell the PieceManager to store the given block
storeBlock :: PieceNum -> Block -> B.ByteString -> Process CF ST ()
storeBlock n blk bs = msgPieceMgr (StoreBlock n blk bs)
-- | The call @grabBlocks n@ will attempt to grab (up to) @n@ blocks from the
-- piece Manager for request at the peer.
grabBlocks :: Int -> Process CF ST [(PieceNum, Block)]
grabBlocks n = do
c <- asks grabBlockTV
ps <- gets peerPieces
msgPieceMgr (GrabBlocks n ps c)
blks <- liftIO $ do atomically $ takeTMVar c
case blks of
Leech bs -> return bs
Endgame bs ->
modify (\s -> s { runningEndgame = True }) >> return bs
createAllPieces :: MonadIO m => Int -> Bool -> m PS.PieceSet
createAllPieces n False = PS.fromList n []
createAllPieces n True = PS.fromList n [0..(n-1)]
createPeerPieces :: MonadIO m => Int -> B.ByteString -> m PS.PieceSet
createPeerPieces nPieces =
PS.fromList nPieces . map fromIntegral . concat . decodeBytes 0 . B.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 CF ST ()
outChan qm = do
modify (\st -> st { lastMessage = 0 })
ch <- asks outCh
liftIO . atomically $ writeTChan ch qm
msgPieceMgr :: PieceMgrMsg -> Process CF ST ()
msgPieceMgr m = do
pmc <- asks pieceMgrCh
{-# SCC "Channel_Write" #-} liftIO . atomically $ writeTChan pmc m
-- IP address is given in host byte-order
allowedFast :: Word32 -> InfoHash -> Int -> Int -> IO [Word32]
allowedFast ip ihash sz n = generate n [] x []
where
-- Take pieces from the generated ones and refill when it is exhausted.
-- While taking pieces, kill duplicates
generate 0 pcs _ _ = return $ reverse pcs
generate k pcs hsh (p : rest)
| p `elem` pcs = generate k pcs hsh rest
| otherwise = generate (k-1) (p : pcs) hsh rest
generate k pcs hsh [] = do
nhsh <- Digest.digestBS hsh
generate k pcs nhsh (genPieces nhsh)
genPieces hash | B.null hash = []
| otherwise =
let (h, rest) = B.splitAt 4 hash
bytes :: [Word32]
bytes = [fromIntegral z `shiftL` s |
(z, s) <- zip (B.unpack h) [24,16,8,0]]
ntohl = fromIntegral . sum
in ((ntohl bytes) `mod` fromIntegral sz) : genPieces rest
-- To prevent a Peer to reconnect, obtain a new IP and thus new FAST-set pieces, we mask out
-- the lower bits
ipBytes = B.pack $ map fromIntegral
[ (ip .&. 0xFF000000) `shiftR` 24
, (ip .&. 0x00FF0000) `shiftR` 16
, (ip .&. 0x0000FF00) `shiftR` 8
, 0
]
x = B.concat [ipBytes, ihash]
testSuite :: Test
testSuite = testGroup "Process/Peer"
[ testCase "AllowedFast" testAllowedFast ]
testAllowedFast :: Assertion
testAllowedFast = do
pcs <- allowedFast (ip32 [80,4,4,200]) tHash 1313 7
assertEqual "Test1" [1059,431,808,1217,287,376,1188] pcs
pcs' <- allowedFast (ip32 [80,4,4,200]) tHash 1313 9
assertEqual "Test2" [1059,431,808,1217,287,376,1188,353,508] pcs'
where ip32 :: [Int] -> Word32
ip32 bytes = fromIntegral $ sum [b `shiftL` s | (b, s) <- zip bytes [24,16,8,0]]
tHash :: B.ByteString
tHash = B.pack $ take 20 (repeat 0xaa)