X-Git-Url: https://git.rkrishnan.org/?a=blobdiff_plain;f=src%2FFuncTorrent%2FPeer.hs;h=c833c061bd4285be88f3754e62e5b83774c1d806;hb=08f3238a8162039a1b77f4fb8825a56a34203c02;hp=798f60acf0712eef0c0644b320c8fe97750cb212;hpb=de4c6fbd9b420c7832ad5df25d70527dfad139d8;p=functorrent.git

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 798f60a..c833c06 100644
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -1,59 +1,202 @@
+{-
+ - 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(..),
-     handShake
+    (PieceMap,
+     handlePeerMsgs
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt)
+import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
 
-import System.IO
-import Data.ByteString (ByteString, unpack, concat, hGet, hPut, singleton)
-import Data.ByteString.Char8 (replicate, pack)
+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 ()
+
+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.
+
+
+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)
 
-type ID = String
-type IP = String
-type Port = Integer
-
-data PeerState = PeerState { am_choking :: Bool
-                           , am_interested :: Bool
-                           , peer_choking :: Bool
-                           , peer_interested :: Bool }
-
--- | Peer is a PeerID, IP address, port tuple
-data Peer = Peer ID IP Port
-          deriving (Show, Eq)
-
-data Msg = HandShakeMsg ByteString ID
-         | KeepAliveMsg
-         | ChokeMsg
-         | UnChokeMsg
-         | InterestedMsg
-         | NotInterestedMsg
-         | HaveMsg Integer
-         | BitFieldMsg Integer
-         | RequestMsg Integer Integer Integer
-         | PieceMsg Integer Integer Integer
-         | CancelMsg Integer Integer Integer
-         | PortMsg Port
-         deriving (Show)
-
-genHandShakeMsg :: ByteString -> String -> ByteString
-genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
-  where pstrlen = singleton 19
-        pstr = pack "BitTorrent protocol"
-        reserved = replicate 8 '\0'
-        peerID = pack peer_id
-
-handShake :: Peer -> ByteString -> String -> IO ByteString
-handShake (Peer _ ip port) infoHash peerid = do
-  let hs = genHandShakeMsg infoHash peerid
-  handle <- connectTo ip (PortNumber (fromIntegral port))
-  hSetBuffering handle LineBuffering
-  hPut handle hs
-  rlenBS <- hGet handle 1
-  let rlen = fromIntegral $ (unpack rlenBS) !! 0
-  hGet handle rlen
-
--- sendMsg :: Peer -> Handle -> PeerMsg -> IO ()
--- recvMsg :: Peer -> Handle -> Msg