X-Git-Url: https://git.rkrishnan.org/?a=blobdiff_plain;f=src%2FFuncTorrent%2FPeer.hs;h=c833c061bd4285be88f3754e62e5b83774c1d806;hb=08f3238a8162039a1b77f4fb8825a56a34203c02;hp=96c66b0c020f57a44d1153686f6134f3bda92aba;hpb=56904b7b9799590e5dc4899246157961ae02c543;p=functorrent.git diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs index 96c66b0..c833c06 100644 --- a/src/FuncTorrent/Peer.hs +++ b/src/FuncTorrent/Peer.hs @@ -1,27 +1,202 @@ +{- + - 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(..), - handShakeMsg + (PieceMap, + handlePeerMsgs ) where -import Prelude hiding (lookup, concat, replicate, splitAt) +import Prelude hiding (lookup, concat, replicate, splitAt, take, drop) + +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 ((!), adjust) +import Network (connectTo, PortID(..)) +import System.IO (Handle, BufferMode(..), hSetBuffering, hClose) + +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 + , meChoking :: Bool + , meInterested :: Bool + , heChoking :: Bool + , heInterested :: Bool} + +havePiece :: PieceMap -> Integer -> Bool +havePiece pm index = + dlstate (pm ! index) == Have + +connectToPeer :: Peer -> IO Handle +connectToPeer (Peer ip port) = do + h <- connectTo ip (PortNumber (fromIntegral port)) + hSetBuffering h LineBuffering + return h + +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 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 + where go [] _ = [] + go (b:bs') pos = + let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i] + in + setBits ++ go bs' (pos + 1) + +-- 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 } + +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 + _ <- runStateT (msgLoop pieceMap c) pstate + return () -import Data.ByteString.Char8 (ByteString, pack, concat, replicate) -import Data.ByteString.Lazy (toChunks) -import Data.Int (Int8) -import qualified Data.Binary as Bin (encode) +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 msgchannel + 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 + 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) + case msg of + KeepAliveMsg -> do + liftIO $ sendMsg h KeepAliveMsg + gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p) + msgLoop pieceStatus msgchannel + 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' msgchannel + UnChokeMsg -> do + 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. -import FuncTorrent.Bencode (InfoDict) -import FuncTorrent.Metainfo (infoHash) --- | Peer is a IP address, port tuple -data Peer = Peer String Integer - deriving (Show, Eq) +downloadPiece :: Handle -> Integer -> Integer -> IO ByteString +downloadPiece h index pieceLength = do + let chunks = splitNum pieceLength 16384 + 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) -handShakeMsg :: InfoDict -> String -> ByteString -handShakeMsg m peer_id = concat [pstrlen, pstr, reserved, infoH, peerID] - where pstrlen = concat $ toChunks $ Bin.encode (19 :: Int8) - pstr = pack "BitTorrent protocol" - reserved = replicate 8 '\0' - infoH = infoHash m - peerID = pack peer_id