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
     , bValToInteger
     , bstrToString
     , decode

+    , decodeWithLeftOvers

     , encode
     ) where
 
     , encode
     ) where
 


@@ -165,6 +166,10 @@ bencVal = Bstr <$> bencStr <|>
 decode :: ByteString -> Either ParseError BVal
 decode = parse bencVal "BVal"
 
 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)
 -- 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 Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
 

+import Control.Concurrent.MVar (MVar, readMVar, putMVar, takeMVar)

 import Control.Monad.State
 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.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 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)
 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}
 
                      , 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
 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
 
   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
 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 ()
   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"
 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
   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
     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
 
 bitfieldToList :: [Word8] -> [Integer]
 bitfieldToList bs = go bs 0


@@ -270,5 +289,64 @@ downloadPiece h index pieceLength = do
 
    At this point, we have the infodict.
 
 
    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
              | 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
 
 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)
              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)
 
     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"
 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
         reserved2 = singleton 0x10 -- support extension protocol

-        reserved3 = BC.replicate 3 '\0'
+        reserved3 = BC.replicate 2 '\0'

         peerID = BC.pack peer_id
 
 bsToInt :: ByteString -> Int
         peerID = BC.pack peer_id
 
 bsToInt :: ByteString -> Int