{-# LANGUAGE OverloadedStrings #-}
module FuncTorrent.Peer
(Peer(..),
- handlePeerMsgs
+ PieceMap,
+ handlePeerMsgs,
+ bytesDownloaded,
+ initPieceMap,
+ pieceMapFromFile
) where
-import Prelude hiding (lookup, concat, replicate, splitAt, take)
+import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
import System.IO (Handle, BufferMode(..), hSetBuffering)
import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
import qualified Data.ByteString.Char8 as BC (length)
import Network (connectTo, PortID(..))
-import Control.Monad (liftM, forM)
+import Control.Monad.State
import Data.Bits
import Data.Word (Word8)
-import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
-import qualified Crypto.Hash.SHA1 as SHA1 (hash)
+import Data.Map (Map, fromList, toList, (!), mapWithKey, traverseWithKey, adjust, filter)
+import Safe (headMay)
import FuncTorrent.Metainfo (Info(..), Metainfo(..))
-import FuncTorrent.Utils (splitN, splitNum)
-import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
+import FuncTorrent.Utils (splitN, splitNum, writeFileAtOffset, readFileAtOffset, verifyHash)
import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
--- PeerState is a misnomer
-data PeerState = PeerState { handle :: Handle
- , peer :: Peer
- , meChoking :: Bool
- , meInterested :: Bool
- , heChoking :: Bool
- , heInterested :: Bool}
+data PState = PState { handle :: Handle
+ , peer :: Peer
+ , meChoking :: Bool
+ , meInterested :: Bool
+ , heChoking :: Bool
+ , heInterested :: Bool}
+
+type PeerState = State PState
data PieceDlState = Pending
- | InProgress
+ | Downloading
| Have
deriving (Show, Eq)
--- todo - map with index to a new data structure (peers who have that piece amd state)
+-- 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
- , state :: PieceDlState -- ^ state of the piece from download perspective.
+ , dlstate :: PieceDlState -- ^ state of the piece from download perspective.
, hash :: ByteString -- ^ piece hash
, len :: Integer } -- ^ piece length
-- Make the initial Piece map, with the assumption that no peer has the
-- piece and that every piece is pending download.
-mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
-mkPieceMap numPieces pieceHash pLengths = fromList kvs
- where kvs = [(i, PieceData { peers = []
- , state = Pending
- , hash = h
- , len = pLen })
- | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
- hashes = splitN 20 pieceHash
+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 = do
+ 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 =
- state (pm ! index) == Have
+ dlstate (pm ! index) == Have
connectToPeer :: Peer -> IO Handle
connectToPeer (Peer _ ip port) = do
in
setBits ++ go bs' (pos + 1)
--- recvMsg :: Peer -> Handle -> Msg
-msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
-msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file =
- do
- -- if me NOT Interested and she is Choking, tell her that
- -- I am interested.
- let h = handle pState
- sendMsg h InterestedMsg
- putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
- msgLoop (pState { meInterested = True }) pieceStatus file
-msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
- -- if me Interested and she not Choking, send her a request
- -- for a piece.
- case pickPiece pieceStatus of
- Nothing -> putStrLn "Nothing to download"
- Just workPiece -> do
- let pLen = len (pieceStatus ! workPiece)
- putStrLn $ "piece length = " ++ show pLen
- pBS <- downloadPiece (handle pState) workPiece pLen
- if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
- then
- putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
- else do
- let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
- putStrLn $ "Write into file at offset: " ++ show fileOffset
- writeFileAtOffset file fileOffset pBS
- msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
-msgLoop pState pieceStatus file = do
- msg <- getMsg (handle pState)
- putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
- case msg of
- KeepAliveMsg -> do
- sendMsg (handle pState) KeepAliveMsg
- putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
- msgLoop pState pieceStatus file
- BitFieldMsg bss -> do
- let pieceList = bitfieldToList (unpack bss)
- pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
- putStrLn $ show (length pieceList) ++ " Pieces"
- -- 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 pState pieceStatus' file
- UnChokeMsg ->
- msgLoop (pState { heChoking = False }) pieceStatus file
- _ ->
- msgLoop pState pieceStatus file
+-- helper functions to manipulate PeerState
+toPeerState :: Handle
+ -> Peer
+ -> Bool -- ^ meChoking
+ -> Bool -- ^ meInterested
+ -> Bool -- ^ heChoking
+ -> Bool -- ^ heInterested
+ -> PState
+toPeerState h p meCh meIn heCh heIn =
+ PState { handle = h
+ , peer = p
+ , heChoking = heCh
+ , heInterested = heIn
+ , 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 m =
- let pieceList = toList m
- allPending = filter (\(_, v) -> state v == Pending) pieceList
- in
- case allPending of
- [] -> Nothing
- ((i, _):_) -> Just i
+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 =
then (pd { peers = p : peers pd })
else pd) pieceStatus
-handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
-handlePeerMsgs p m peerId = do
+handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> IO ()
+handlePeerMsgs p peerId m pieceMap = do
h <- connectToPeer p
doHandshake h p (infoHash m) peerId
- let state = PeerState { handle = h
- , peer = p
- , heInterested = False
- , heChoking = True
- , meInterested = False
- , meChoking = True }
- pieceHash = pieces (info m)
- numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
- pLen = pieceLength (info m)
- fileLen = lengthInBytes (info m)
- fileName = name (info m)
- pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
- createDummyFile fileName (fromIntegral fileLen)
- msgLoop state pieceStatus fileName
-
+ let pstate = toPeerState h p False False True True
+ filePath = name (info m)
+ _ <- runStateT (msgLoop pieceMap filePath) pstate
+ return ()
+
+msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
+msgLoop pieceStatus file = do
+ h <- gets handle
+ st <- get
+ case st of
+ PState { meInterested = False, heChoking = True } -> do
+ liftIO $ sendMsg h InterestedMsg
+ gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
+ modify (\st -> st { meInterested = True })
+ msgLoop pieceStatus file
+ PState { meInterested = True, heChoking = False } ->
+ case pickPiece pieceStatus of
+ Nothing -> liftIO $ putStrLn "Nothing to download"
+ Just workPiece -> do
+ let pLen = len (pieceStatus ! workPiece)
+ liftIO $ putStrLn $ "piece length = " ++ show pLen
+ pBS <- liftIO $ downloadPiece h workPiece pLen
+ if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
+ 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
+ _ -> do
+ msg <- liftIO $ getMsg h
+ gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
+ case msg of
+ KeepAliveMsg -> do
+ liftIO $ sendMsg h KeepAliveMsg
+ gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
+ msgLoop pieceStatus file
+ BitFieldMsg bss -> do
+ p <- gets peer
+ let pieceList = bitfieldToList (unpack bss)
+ pieceStatus' = updatePieceAvailability pieceStatus p pieceList
+ liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
+ -- 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
+ UnChokeMsg -> do
+ modify (\st -> st {heChoking = False })
+ msgLoop pieceStatus file
+
+
downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
downloadPiece h index pieceLength = do
let chunks = splitNum pieceLength 16384
- liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
- sendMsg h (RequestMsg index (i*pLen) pLen)
- putStrLn $ "--> " ++ "RequestMsg for Piece "
- ++ show index ++ ", part: " ++ show i ++ " of length: "
- ++ show pLen
- msg <- getMsg h
- case msg of
- PieceMsg index begin block -> do
- putStrLn $ " <-- PieceMsg for Piece: "
- ++ show index
- ++ ", offset: "
- ++ show begin
- return block
- _ -> do
- putStrLn "ignoring irrelevant msg"
- return empty)
-
-verifyHash :: ByteString -> ByteString -> Bool
-verifyHash bs pieceHash =
- take 20 (SHA1.hash bs) == pieceHash
+ concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
+ sendMsg h (RequestMsg index (i*pLen) pLen)
+ putStrLn $ "--> " ++ "RequestMsg for Piece "
+ ++ show index ++ ", part: " ++ show i ++ " of length: "
+ ++ show pLen
+ msg <- getMsg h
+ case msg of
+ PieceMsg index begin block -> do
+ putStrLn $ " <-- PieceMsg for Piece: "
+ ++ show index
+ ++ ", offset: "
+ ++ show begin
+ return block
+ _ -> do
+ putStrLn "ignoring irrelevant msg"
+ return empty)
+
projects / functorrent.git / blobdiff