quick and dirty initial extended metadata protocol implementation

[functorrent.git] / src / FuncTorrent / Peer.hs

diff --git a/src/FuncTorrent/Peer.hs b/src/FuncTorrent/Peer.hs
index 87e7d23bff683a448beb8f644b8820601db10886..4396c1c004bd5364c9a75deecb34b862d3c10dc7 100644 (file)
--- a/src/FuncTorrent/Peer.hs
+++ b/src/FuncTorrent/Peer.hs
@@ -25,14 +25,16 @@ module FuncTorrent.Peer
 
 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
 
+import Control.Concurrent.MVar (MVar, readMVar, putMVar, takeMVar)
 import Control.Monad.State
-import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
+import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty, singleton)
 import Data.Bits
 import Data.Word (Word8)
-import Data.Map ((!), adjust)
+import Data.Map (Map, (!), adjust, fromList, insert)
 import Network (connectTo, PortID(..))
 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
 
+import FuncTorrent.Bencode(BVal(..), encode, decode, decodeWithLeftOvers)
 import FuncTorrent.Metainfo (Metainfo(..))
 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
 import FuncTorrent.Utils (splitNum, verifyHash)
@@ -46,6 +48,11 @@ data PState = PState { handle :: Handle
                      , heChoking :: Bool
                      , heInterested :: Bool}
 
+data InfoPieceMap = InfoPieceMap { infoLength :: Integer
+                                 , infoMap :: Map Integer (Maybe ByteString) }
+
+newtype InfoState = InfoState (MVar InfoPieceMap)
+
 havePiece :: PieceMap -> Integer -> Bool
 havePiece pm index =
   dlstate (pm ! index) == Have
@@ -56,18 +63,25 @@ connectToPeer (Peer ip port) = do
   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))
+  hsMsg <- hGet h (length (unpack hs))
   putStrLn $ "<-- handshake from peer: " ++ show p
   return ()
+  -- if doesPeerSupportExtendedMsg hsMsg
+  --   then
+  --   return doExtendedHandshake h
+  --   else
+  --   return Nothing
 doHandshake False h p infohash peerid = do
   let hs = genHandshakeMsg infohash peerid
   putStrLn "waiting for a handshake"
-  hsMsg <- hGet h (length (unpack hs))
+  -- read 28 bytes. '19' ++ 'BitTorrent Protocol' ++ 8 reserved bytes
+  hsMsg <- hGet h 28
   putStrLn $ "<-- handshake from peer: " ++ show p
   let rxInfoHash = take 20 $ drop 28 hsMsg
   if rxInfoHash /= infohash
@@ -78,7 +92,12 @@ doHandshake False h p infohash peerid = do
     else do
     _ <- hPut h hs
     putStrLn $ "--> handhake to peer: " ++ show p
-    return ()
+    -- if doesPeerSupportExtendedMsg hsMsg
+    --   then do
+    --   doExtendedHandshake h
+    --   else
+    --   return Nothing
+
 
 bitfieldToList :: [Word8] -> [Integer]
 bitfieldToList bs = go bs 0
@@ -270,5 +289,64 @@ downloadPiece h index pieceLength = do
 
    At this point, we have the infodict.
 
--)
+-}
 
+{-
+data InfoPieceMap = { infoLength :: Integer
+                    , infoMap :: Map Integer (Maybe ByteString)
+                    }
+
+newtype InfoState = InfoState (MVar InfoPieceMap)
+
+-}
+
+
+metadataMsgLoop :: Handle -> InfoState -> IO ()
+metadataMsgLoop h (InfoState st) = do
+    infoState <- readMVar st
+    let metadataLen = infoLength infoState
+        -- send the handshake msg
+        metadata = encode (metadataMsg metadataLen)
+    sendMsg h (ExtendedMsg 0 metadata)
+    -- recv return msg from the peer. Will have 'metadata_size'
+    msg <- getMsg h
+    case msg of
+      ExtendedMsg 0 rBs -> do
+        -- decode rBs
+        let (Right (Bdict msgMap)) = decode rBs
+            (Bdict mVal) = msgMap ! "m" -- which is another dict
+            (Bint metadata_msgID) = mVal ! "ut_metadata"
+            (Bint metadata_size) = msgMap ! "metadata_size"
+            -- divide metadata_size into 16384 sized pieces, find number of pieces
+            (q, r) = metadata_size `divMod` 16384
+            -- pNumLengthPairs = zip [0..q-1] (take q (repeat 16384)) ++ (q, r)
+            -- TODO: corner case where infodict size is a multiple of 16384
+            -- and start sending request msg for each.
+        if metadataLen == 0
+          then -- We don't have any piece. Send request msg for all pieces.
+          mapM_ (\n -> do
+                    sendMsg h (ExtendedMsg metadata_msgID (encode (requestMsg n)))
+                    dataOrRejectMsg <- getMsg h
+                    case dataOrRejectMsg of
+                      ExtendedMsg 3 payload -> do
+                        -- bencoded dict followed by XXXXXX
+                        infoState <- takeMVar st
+                        let (Right (Bdict bval, pieceData)) = decodeWithLeftOvers payload
+                            (Bint pieceIndex) = bval ! "piece"
+                            payloadLen = length (unpack pieceData)
+                            infoMapVal = infoMap infoState
+                        putMVar st infoState {
+                          infoMap = insert pieceIndex (Just payload) infoMapVal }
+                )
+          [0..q]
+          else
+          return () -- TODO: reject for now
+      where
+        metadataMsg 0 = Bdict (fromList [("m", Bdict (fromList [("ut_metadata", (Bint 3))]))])
+        metadataMsg l = Bdict (fromList [("m", Bdict (fromList [("ut_metadata", (Bint 3))])),
+                                         ("metadata_size", (Bint l))])
+        requestMsg i = Bdict (fromList [("msg_type", (Bint 0)), ("piece", (Bint i))])
+        rejectmsg i = Bdict (fromList [("msg_type", (Bint 2)), ("piece", (Bint i))])
+
+doesPeerSupportExtendedMsg :: ByteString -> Bool
+doesPeerSupportExtendedMsg bs = take 1 (drop 5 bs) == singleton 0x10