X-Git-Url: https://git.rkrishnan.org/?a=blobdiff_plain;f=src%2FFuncTorrent%2FPeer.hs;h=63aaa5afe9ede340862383318370616a8c9813e8;hb=9beb0fb9814b33725f6adfa5adabb3225a54277b;hp=ae95f838b38215279ca3150a6002a1086caefd1a;hpb=be3d38d002fe56d86bf824afe8e6356f7c58c639;p=functorrent.git diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs index ae95f83..63aaa5a 100644 --- a/src/FuncTorrent/Peer.hs +++ b/src/FuncTorrent/Peer.hs @@ -1,27 +1,44 @@ +{- + - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan + - + - 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 + -} + {-# LANGUAGE OverloadedStrings #-} + module FuncTorrent.Peer - (Peer(..), - handlePeerMsgs, - bytesDownloaded + (PieceMap, + handlePeerMsgs ) where -import Prelude hiding (lookup, concat, replicate, splitAt, take, filter) +import Prelude hiding (lookup, concat, replicate, splitAt, take, drop) -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.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, adjust, filter) -import qualified Crypto.Hash.SHA1 as SHA1 (hash) -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) -import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset) +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 @@ -30,52 +47,39 @@ data PState = PState { handle :: Handle , 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. -mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap -mkPieceMap numPieces pieceHash pLengths = fromList kvs - where kvs = [(i, PieceData { peers = [] - , dlstate = Pending - , hash = h - , len = pLen }) - | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths] - hashes = splitN 20 pieceHash - 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 :: Handle -> Peer -> ByteString -> String -> IO () -doHandshake h peer infoHash peerid = do - let hs = genHandshakeMsg infoHash peerid +doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO () +doHandshake True h p infohash peerid = do + let hs = genHandshakeMsg infohash peerid hPut h hs - putStrLn $ "--> handhake to peer: " ++ show peer + putStrLn $ "--> handhake to peer: " ++ show p _ <- hGet h (length (unpack hs)) - putStrLn $ "<-- handshake from peer: " ++ show peer + 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 @@ -101,48 +105,25 @@ toPeerState h p meCh meIn heCh 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 = - (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 -> Metainfo -> String -> IO () -handlePeerMsgs p m peerId = do +handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO () +handlePeerMsgs p peerId m pieceMap isClient c = do h <- connectToPeer p - doHandshake h p (infoHash m) peerId + doHandshake isClient h p (infoHash m) peerId let pstate = toPeerState h p False False True 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) - (r, _) <- runStateT (msgLoop pieceStatus fileName) 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 - PState { meInterested = True, heChoking = False } -> do + modify (\st' -> st' { meInterested = True }) + msgLoop pieceStatus msgchannel + PState { meInterested = True, heChoking = False } -> case pickPiece pieceStatus of Nothing -> liftIO $ putStrLn "Nothing to download" Just workPiece -> do @@ -151,20 +132,19 @@ msgLoop pieceStatus file = do pBS <- liftIO $ downloadPiece h workPiece pLen if not $ verifyHash pBS (hash (pieceStatus ! workPiece)) then - liftIO $ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS)) + 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) @@ -173,32 +153,50 @@ msgLoop pieceStatus file = do -- 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 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: " ++ show msg + return empty) +