quick and dirty initial extended metadata protocol implementation

diff --git a/src/FuncTorrent/Bencode.hs b/src/FuncTorrent/Bencode.hs
index d4eb58c5ac12d0e4f811eb050f95266733b9f262..ac6221119663077f83aeb17757d26a7afc057926 100644 (file)
--- a/src/FuncTorrent/Bencode.hs
+++ b/src/FuncTorrent/Bencode.hs
@@ -23,6 +23,7 @@ module FuncTorrent.Bencode
     , bValToInteger
     , bstrToString
     , decode
+    , decodeWithLeftOvers
     , encode
     ) where
 
@@ -165,6 +166,10 @@ bencVal = Bstr <$> bencStr <|>
 decode :: ByteString -> Either ParseError BVal
 decode = parse bencVal "BVal"
 
+decodeWithLeftOvers :: ByteString -> Either ParseError (BVal, ByteString)
+decodeWithLeftOvers = parse ((,) <$> bencVal <*> (fmap pack leftOvers)) "BVal with LeftOvers"
+  where leftOvers = manyTill anyToken eof
+
 -- Encode BVal into a bencoded ByteString. Inverse of decode
 
 -- TODO: Use builders and lazy byte string to get O(1) concatenation over O(n)

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

diff --git a/src/FuncTorrent/PeerMsgs.hs b/src/FuncTorrent/PeerMsgs.hs
index d6bbdcfadefac6d3f7672f090a952492d7299fa2..467ac8d6f72e94a41bc9233643e731ded0a4e802 100644 (file)
--- a/src/FuncTorrent/PeerMsgs.hs
+++ b/src/FuncTorrent/PeerMsgs.hs
@@ -64,7 +64,12 @@ data PeerMsg = KeepAliveMsg
              | CancelMsg Integer Integer Integer
              | PortMsg Port
              | ExtendedMsg Integer ByteString
-             deriving (Show)
+  deriving (Show)
+
+data ExtMetadataMsg = Request Integer
+                    | Data Integer Integer
+                    | Reject Integer
+  deriving (Eq, Show)
 
 instance Binary PeerMsg where
   put msg = case msg of
@@ -101,14 +106,13 @@ instance Binary PeerMsg where
              PortMsg p -> do putWord32be 3
                              putWord8 9
                              putWord16be (fromIntegral p)
-             ExtendedHandshakeMsg t b-> do putWord32be msgLen
-                                           putWord8 20
-                                           putWord8 t -- 0 => handshake msg
-                                           -- actual extension msg follows
-                                           mapM_ putWord8 blockList
-                                             where blockList = unpack b
-                                                   blockLen  = length blockList
-
+             ExtendedMsg t b-> do putWord32be (fromIntegral blockLen)
+                                  putWord8 20
+                                  putWord8 (fromIntegral t) -- 0 => handshake msg
+                                  -- actual extension msg follows
+                                  mapM_ putWord8 blockList
+                                    where blockList = unpack b
+                                          blockLen  = length blockList
 
     where putIndexOffsetLength i o l = do
             putWord32be (fromIntegral i)
@@ -152,9 +156,9 @@ genHandshakeMsg :: ByteString -> String -> ByteString
 genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved1, reserved2, reserved3, infoHash, peerID]
   where pstrlen = singleton 19
         pstr = BC.pack "BitTorrent protocol"
-        reserved1 = BC.replicate 4 '\0'
+        reserved1 = BC.replicate 5 '\0'
         reserved2 = singleton 0x10 -- support extension protocol
-        reserved3 = BC.replicate 3 '\0'
+        reserved3 = BC.replicate 2 '\0'
         peerID = BC.pack peer_id
 
 bsToInt :: ByteString -> Int