X-Git-Url: https://git.rkrishnan.org/?a=blobdiff_plain;f=src%2FFuncTorrent%2FPeer.hs;h=63aaa5afe9ede340862383318370616a8c9813e8;hb=9beb0fb9814b33725f6adfa5adabb3225a54277b;hp=f9fdbc32c45fb6afae5378a6fd425aad68f6ec19;hpb=ca6948b18ecf506d2959216b98cec9e85a6c463d;p=functorrent.git

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index f9fdbc3..63aaa5a 100644
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -1,68 +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(..),
-     PeerResp(..),
-     getPeerResponse,
-     handShakeMsg
+    (PieceMap,
+     handlePeerMsgs
     ) where
 
-import Prelude hiding (lookup, concat, replicate, splitAt)
-import Data.ByteString.Char8 (ByteString, pack, unpack, concat, replicate, splitAt)
-import Data.ByteString.Lazy (toChunks)
-import Data.Int (Int8)
-import Data.List (intercalate)
-import Data.Map as M ((!), lookup)
-import qualified Data.Binary as Bin (encode)
-import qualified Data.ByteString.Base16 as B16 (encode)
-
-import FuncTorrent.Bencode (BVal(..), InfoDict, decode)
-import FuncTorrent.Tracker (infoHash)
-import FuncTorrent.Utils (splitN)
-
-
-type Address = String
-type Port = Integer
-
-data Peer = Peer Address Port
-            deriving (Show, Eq)
-
-data PeerResp = PeerResp { interval :: Maybe Integer
-                         , peers :: [Peer]
-                         , complete :: Maybe Integer
-                         , incomplete :: Maybe Integer
-                         } deriving (Show, Eq)
-
-toInt :: String -> Integer
-toInt = read
-
-getPeerResponse :: ByteString -> PeerResp
-getPeerResponse body = case decode body of
-                        Right (Bdict peerM) ->
-                          let (Just (Bint i)) = lookup "interval" peerM
-                              (Bstr peersBS) = peerM ! "peers"
-                              pl = map (\peer -> let (ip', port') = splitAt 4 peer
-                                                 in Peer (toIPNum ip') (toPortNum port'))
-                                   (splitN 6 peersBS)
-                          in PeerResp { interval = Just i
-                                      , peers = pl
-                                      , complete = Nothing
-                                      , incomplete = Nothing
-                                      }
-                          where toPortNum = read . ("0x" ++) . unpack . B16.encode
-                                toIPNum = intercalate "." .
-                                          map (show . toInt . ("0x" ++) . unpack) .
-                                          splitN 2 . B16.encode
-
-                        _ -> PeerResp { interval = Nothing
-                                      , peers = []
-                                      , complete = Nothing
-                                      , incomplete = Nothing
-                                      }
-
-
-handShakeMsg :: InfoDict -> String -> ByteString
-handShakeMsg m peer_id = let pstrlen = concat $ toChunks $ Bin.encode (19 :: Int8)
-                             pstr = pack "BitTorrent protocol"
-                             reserved = replicate 8 '\0'
-                             infoH = infoHash m
-                             peerID = pack peer_id
-                         in concat [pstrlen, pstr, reserved, infoH, peerID]
+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 ()
+
+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)
+