+{-
+ - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
+ -
+ - This file is part of FuncTorrent.
+ -
+ - FuncTorrent is free software; you can redistribute it and/or modify
+ - it under the terms of the GNU General Public License as published by
+ - the Free Software Foundation; either version 3 of the License, or
+ - (at your option) any later version.
+ -
+ - FuncTorrent is distributed in the hope that it will be useful,
+ - but WITHOUT ANY WARRANTY; without even the implied warranty of
+ - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ - GNU General Public License for more details.
+ -
+ - You should have received a copy of the GNU General Public License
+ - along with FuncTorrent; if not, see <http://www.gnu.org/licenses/>
+ -}
+
{-# LANGUAGE OverloadedStrings #-}
+
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.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, writePieceToDisk)
data PState = PState { handle :: Handle
, peer :: Peer
, heChoking :: Bool
, heInterested :: Bool}
-type PeerState = State PState
-
-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 = 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 =
dlstate (pm ! index) == Have
connectToPeer :: Peer -> IO Handle
-connectToPeer (Peer _ ip port) = do
+connectToPeer (Peer ip port) = do
h <- connectTo ip (PortNumber (fromIntegral port))
hSetBuffering h LineBuffering
return h
doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
-doHandshake isClient h peer infoHash peerid =
- let hs = genHandshakeMsg infoHash peerid
- in
- if isClient
- then do
- hPut h hs
- putStrLn $ "--> handhake to peer: " ++ show peer
- _ <- hGet h (length (unpack hs))
- putStrLn $ "<-- handshake from peer: " ++ show peer
- return ()
- else do
- putStrLn $ "waiting for a handshake"
- hsMsg <- hGet h (length (unpack hs))
- putStrLn $ "<-- handshake from peer: " ++ show peer
- let rxInfoHash = take 20 $ drop 28 hsMsg
- if rxInfoHash /= infoHash
- then do
- putStrLn $ "infoHashes does not match"
- hClose h
- return ()
- else do
- _ <- hPut h hs
- putStrLn $ "--> handhake to peer: " ++ show peer
- return ()
+doHandshake True h p infohash peerid = do
+ let hs = genHandshakeMsg infohash peerid
+ hPut h hs
+ putStrLn $ "--> handhake to peer: " ++ show p
+ _ <- hGet h (length (unpack hs))
+ putStrLn $ "<-- handshake from peer: " ++ show p
+ return ()
+doHandshake False h p infohash peerid = do
+ let hs = genHandshakeMsg infohash peerid
+ putStrLn "waiting for a handshake"
+ hsMsg <- hGet h (length (unpack hs))
+ putStrLn $ "<-- handshake from peer: " ++ show p
+ let rxInfoHash = take 20 $ drop 28 hsMsg
+ if rxInfoHash /= infohash
+ then do
+ putStrLn "infoHashes does not match"
+ hClose h
+ return ()
+ else do
+ _ <- hPut h hs
+ putStrLn $ "--> handhake to peer: " ++ show p
+ return ()
bitfieldToList :: [Word8] -> [Integer]
bitfieldToList bs = go bs 0
, 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
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
+ modify (\st' -> st' { meInterested = True })
+ msgLoop pieceStatus msgchannel
PState { meInterested = True, heChoking = False } ->
case pickPiece pieceStatus of
Nothing -> liftIO $ putStrLn "Nothing to download"
pBS <- liftIO $ downloadPiece h workPiece pLen
if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
then
- liftIO $ putStrLn $ "Hash mismatch"
+ 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 $ FS.writePieceToDisk msgchannel 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)
+ 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
+ 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
+ modify (\st' -> st' {heChoking = False })
+ msgLoop pieceStatus msgchannel
+ ChokeMsg -> do
+ modify (\st' -> st' {heChoking = True })
+ msgLoop pieceStatus msgchannel
+ InterestedMsg -> do
+ modify (\st' -> st' {heInterested = True})
+ msgLoop pieceStatus msgchannel
+ NotInterestedMsg -> do
+ modify (\st' -> st' {heInterested = False})
+ msgLoop pieceStatus msgchannel
+ CancelMsg {} -> -- check if valid index, begin, length
+ msgLoop pieceStatus msgchannel
+ PortMsg _ ->
+ msgLoop pieceStatus msgchannel
+ HaveMsg idx -> do
+ p <- gets peer
+ let pieceStatus' = updatePieceAvailability pieceStatus p [idx]
+ msgLoop pieceStatus' msgchannel
+ _ -> do
+ liftIO $ putStrLn ".. not doing anything with the msg"
+ msgLoop pieceStatus msgchannel
+ -- No need to handle PieceMsg and RequestMsg here.
downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
++ show begin
return block
_ -> do
- putStrLn "ignoring irrelevant msg"
+ putStrLn $ "ignoring irrelevant msg: " ++ show msg
return empty)