module FuncTorrent.Peer
(Peer(..),
PieceMap,
- handlePeerMsgs,
- bytesDownloaded,
- initPieceMap,
- pieceMapFromFile
+ handlePeerMsgs
) where
-import Prelude hiding (lookup, concat, replicate, splitAt, take, drop, filter)
+import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
-import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
-import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
-import qualified Data.ByteString.Char8 as BC (length)
-import Network (connectTo, PortID(..))
+import Control.Concurrent.Chan (writeChan)
import Control.Monad.State
+import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
import Data.Bits
import Data.Word (Word8)
-import Data.Map (Map, fromList, toList, (!), mapWithKey, traverseWithKey, adjust, filter)
-import Safe (headMay)
+import Data.Map ((!), adjust)
+import Network (connectTo, PortID(..))
+import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
-import FuncTorrent.Metainfo (Info(..), Metainfo(..))
-import FuncTorrent.Utils (splitN, splitNum, writeFileAtOffset, readFileAtOffset, verifyHash)
+import FuncTorrent.Metainfo (Metainfo(..))
import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
+import FuncTorrent.Utils (splitNum, verifyHash)
+import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
+import qualified FuncTorrent.FileSystem as FS (MsgChannel, Msg(..), Piece(..))
data PState = PState { handle :: Handle
, peer :: Peer
, heChoking :: Bool
, heInterested :: Bool}
-data PieceDlState = Pending
- | Downloading
- | Have
- deriving (Show, Eq)
-
--- todo - map with index to a new data structure (peers who have that piece and state)
-data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
- , dlstate :: PieceDlState -- ^ state of the piece from download perspective.
- , hash :: ByteString -- ^ piece hash
- , len :: Integer } -- ^ piece length
-
--- which piece is with which peers
-type PieceMap = Map Integer PieceData
-
-
--- Make the initial Piece map, with the assumption that no peer has the
--- piece and that every piece is pending download.
-initPieceMap :: ByteString -> Integer -> Integer -> PieceMap
-initPieceMap pieceHash fileLen pieceLen = fromList kvs
- where
- numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
- kvs = [(i, PieceData { peers = []
- , dlstate = Pending
- , hash = h
- , len = pLen })
- | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
- hashes = splitN 20 pieceHash
- pLengths = splitNum fileLen pieceLen
-
-pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
-pieceMapFromFile filePath pieceMap =
- traverseWithKey f pieceMap
- where
- f k v = do
- let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
- isHashValid <- flip verifyHash (hash v) <$> readFileAtOffset filePath offset (len v)
- if isHashValid
- then return $ v { dlstate = Have }
- else return v
-
havePiece :: PieceMap -> Integer -> Bool
havePiece pm index =
dlstate (pm ! index) == Have
, meChoking = meCh
, meInterested = meIn }
--- simple algorithm to pick piece.
--- pick the first piece from 0 that is not downloaded yet.
-pickPiece :: PieceMap -> Maybe Integer
-pickPiece =
- (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
-
-bytesDownloaded :: PieceMap -> Integer
-bytesDownloaded =
- sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
-
-updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
-updatePieceAvailability pieceStatus p pieceList =
- mapWithKey (\k pd -> if k `elem` pieceList
- then (pd { peers = p : peers pd })
- else pd) pieceStatus
-
-handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> IO ()
-handlePeerMsgs p peerId m pieceMap isClient = do
+handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
+handlePeerMsgs p peerId m pieceMap isClient c = do
h <- connectToPeer p
doHandshake isClient h p (infoHash m) peerId
let pstate = toPeerState h p False False True True
- filePath = name (info m)
- _ <- runStateT (msgLoop pieceMap filePath) pstate
+ _ <- runStateT (msgLoop pieceMap c) pstate
return ()
-msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
-msgLoop pieceStatus file = do
+msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
+msgLoop pieceStatus msgchannel = do
h <- gets handle
st <- get
case st of
liftIO $ sendMsg h InterestedMsg
gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
modify (\st' -> st' { meInterested = True })
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
PState { meInterested = True, heChoking = False } ->
case pickPiece pieceStatus of
Nothing -> liftIO $ putStrLn "Nothing to download"
then
liftIO $ putStrLn "Hash mismatch"
else do
- let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
- liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
- liftIO $ writeFileAtOffset file fileOffset pBS
- msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
+ liftIO $ putStrLn $ "Write piece: " ++ show workPiece
+ liftIO $ writeChan msgchannel $ FS.WritePiece (FS.Piece workPiece pBS)
+ msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
_ -> do
msg <- liftIO $ getMsg h
gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
KeepAliveMsg -> do
liftIO $ sendMsg h KeepAliveMsg
gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
BitFieldMsg bss -> do
p <- gets peer
let pieceList = bitfieldToList (unpack bss)
-- for each pieceIndex in pieceList, make an entry in the pieceStatus
-- map with pieceIndex as the key and modify the value to add the peer.
-- download each of the piece in order
- msgLoop pieceStatus' file
+ msgLoop pieceStatus' msgchannel
UnChokeMsg -> do
modify (\st' -> st' {heChoking = False })
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
ChokeMsg -> do
modify (\st' -> st' {heChoking = True })
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
InterestedMsg -> do
modify (\st' -> st' {heInterested = True})
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
NotInterestedMsg -> do
modify (\st' -> st' {heInterested = False})
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
CancelMsg _ _ _ -> -- check if valid index, begin, length
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
PortMsg _ ->
- msgLoop pieceStatus file
+ msgLoop pieceStatus msgchannel
-- handle RequestMsg, HaveMsg. No need to handle PieceMsg here.
-- also BitFieldMsg
++ show begin
return block
_ -> do
- putStrLn "ignoring irrelevant msg"
+ putStrLn $ "ignoring irrelevant msg: " ++ show msg
return empty)
projects / functorrent.git / blobdiff