Add getMetadata function

2 files changed, 53 insertions, 5 deletions

diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index e93f8bbe..d0d2bb03 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -83,6 +83,8 @@ module Network.BitTorrent.Exchange.Message
        , metadataPieceSize
        , defaultMetadataFactor
        , defaultMaxInfoDictSize
+       , isLastPiece
+       , isValidPiece
        ) where
 
 import Control.Applicative
@@ -859,13 +861,18 @@ instance PeerMessage ExtendedMetadata where
 metadataPieceSize :: Int
 metadataPieceSize = 16 * 1024
 
+isLastPiece :: P.Piece a -> Int -> Bool
+isLastPiece P.Piece {..} total = succ pieceIndex == pcnt
+  where
+    pcnt = q + if r > 0 then 1 else 0
+    (q, r) = quotRem total metadataPieceSize
+
 -- TODO we can check if the piece payload bytestring have appropriate
 -- length; otherwise serialization MUST fail.
-isLastMetadata :: ExtendedMetadata -> Bool
-isLastMetadata = undefined -- FIXME
-
-checkPiece :: ExtendedMetadata -> Bool
-checkPiece = undefined -- FIXME
+isValidPiece :: P.Piece BL.ByteString -> Int -> Bool
+isValidPiece p @ P.Piece {..} total
+  | isLastPiece p total = P.pieceSize p <= metadataPieceSize
+  |       otherwise     = P.pieceSize p == metadataPieceSize
 
 setMetadataPayload :: BS.ByteString -> ExtendedMetadata -> ExtendedMetadata
 setMetadataPayload bs (MetadataData (P.Piece pix _) t) =
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 369f0746..427d7676 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -55,6 +55,7 @@ module Network.BitTorrent.Exchange.Wire
        , getConnection
        , getExtCaps
        , getStats
+       , getMetadata
        ) where
 
 import Control.Applicative
@@ -689,3 +690,43 @@ acceptWire :: Socket -> PeerAddr IP -> Wire () -> IO ()
 acceptWire sock peerAddr wire = do
   bracket (return sock) close $ \ _ -> do
     error "acceptWire: not implemented"
+
+{-----------------------------------------------------------------------
+--  Metadata exchange
+-----------------------------------------------------------------------}
+-- TODO introduce new metadata exchange specific exceptions
+
+fetchMetadata :: Wire [BS.ByteString]
+fetchMetadata = loop 0
+  where
+    recvData = recvMessage >>= inspect
+      where
+        inspect (Extended (EMetadata _ meta)) =
+          case meta of
+            MetadataRequest pix -> do
+              sendMessage (MetadataReject pix)
+              recvData
+            MetadataData   {..} -> return (piece, totalSize)
+            MetadataReject   _  -> disconnectPeer
+            MetadataUnknown   _ -> recvData
+        inspect _ = recvData
+
+    loop i = do
+      sendMessage (MetadataRequest i)
+      (piece, totalSize) <- recvData
+      unless (pieceIndex piece == i) $ do
+        disconnectPeer
+
+      if piece `isLastPiece` totalSize
+        then pure [pieceData piece]
+        else (pieceData piece :) <$> loop (succ i)
+
+getMetadata :: Wire InfoDict
+getMetadata = do
+  chunks <- fetchMetadata
+  Connection {..} <- getConnection
+  case BE.decode (BS.concat chunks) of
+     Right (infodict @ InfoDict {..})
+       | connTopic == idInfoHash -> return infodict
+       |         otherwise       -> error "broken infodict"
+     Left err -> error $ "unable to parse infodict" ++ err