Implemented BEP 32 (IPv6 extensions) for get-peers.

1 files changed, 21 insertions, 34 deletions

diff --git a/Mainline.hs b/Mainline.hs
index 12d1540b..e3a48976 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -23,6 +23,7 @@ import qualified Data.ByteArray           as BA
 import qualified Data.ByteString          as B
          ;import Data.ByteString          (ByteString)
 import Data.ByteString.Lazy               (toStrict)
+import Data.Coerce
 import Data.Data
 import Data.Default
 import Data.Hashable
@@ -413,13 +414,13 @@ instance BEncode GetPeers where
   fromBEncode = fromDict $ GetPeers <$>! info_hash_key <*>? want_key
 
 
-type PeerList = Either [NodeInfo] [PeerAddr]
-
 data GotPeers = GotPeers
   { -- | If the queried node has no peers for the infohash, returned
     -- the K nodes in the queried nodes routing table closest to the
     -- infohash supplied in the query.
-    peers        :: PeerList
+    peers        :: [PeerAddr]
+
+  , nodes        :: NodeFound
 
     -- | The token value is a required argument for a future
     -- announce_peer query.
@@ -431,34 +432,22 @@ nodeIsIPv6 (NodeInfo _ (IPv6 _) _) = True
 nodeIsIPv6 _                       = False
 
 instance BEncode GotPeers where
-  toBEncode GotPeers {..} = toDict $
-    case peers of
-      Left  ns
-        | let (ns6,ns4) = partition nodeIsIPv6 ns
-        ->
-           nodes_key .=? (if null ns4 then Nothing
-                                      else Just $ S.runPut (mapM_ putNodeInfo4 ns4))
-        .: nodes6_key .=? (if null ns6 then Nothing
-                                       else Just $ S.runPut (mapM_ putNodeInfo4 ns6))
-        .: token_key .=! grantedToken
-        .: endDict
-      Right ps ->
-           token_key .=! grantedToken
-        .: peers_key .=! map S.encode ps -- TODO: Spec says we shouldn't mix ip4/ip6 here.
-                                         -- (We could filter in MethodHandler.)
-        .: endDict
+  toBEncode GotPeers { nodes = NodeFound ns4 ns6, ..} = toDict $
+       nodes_key  .=? (if null ns4 then Nothing
+                                   else Just $ S.runPut (mapM_ putNodeInfo4 ns4))
+    .: nodes6_key .=? (if null ns6 then Nothing
+                                   else Just $ S.runPut (mapM_ putNodeInfo4 ns6))
+    .: token_key .=! grantedToken
+    .: peers_key .=! map S.encode peers -- TODO: Spec says we shouldn't mix ip4/ip6 here.
+                                        -- (We could filter in MethodHandler.)
+    .: endDict
 
   fromBEncode = fromDict $ do
-     mns4 <- optional (binary getNodeInfo4 nodes_key)        -- "nodes"
-     mns6 <- optional (binary getNodeInfo6 nodes6_key)       -- "nodes6"
-     let mns = ((++) <$> mns4 <*> mns6)
-               <|> mns4
-               <|> mns6
-     tok <- field    (req    token_key)                      -- "token"
-     mps <- optional (field (req peers_key) >>= decodePeers) -- "values"
-     case (Right <$> mps) <|> (Left <$> mns) of
-       Nothing -> fail "get_peers: neihter peers nor nodes key is valid"
-       Just xs -> pure $ GotPeers xs tok
+    ns4 <- fromMaybe [] <$> optional (binary getNodeInfo4 nodes_key)        -- "nodes"
+    ns6 <- fromMaybe [] <$> optional (binary getNodeInfo6 nodes6_key)       -- "nodes6"
+    tok <- field    (req    token_key)                                      -- "token"
+    ps <- fromMaybe [] <$> optional (field (req peers_key) >>= decodePeers) -- "values"
+    pure $ GotPeers ps (NodeFound ns4 ns6) tok
      where
        decodePeers = either fail pure . mapM S.decode
 
@@ -466,15 +455,13 @@ getPeersH :: Routing -> SwarmsDatabase -> NodeInfo -> GetPeers -> IO GotPeers
 getPeersH routing (SwarmsDatabase peers toks _) naddr (GetPeers ih iptyp) = do
   ps <- do
       tm <- getTimestamp
-      ps <- atomically $ do
+      atomically $ do
         (ps,store') <- Peers.freshPeers ih tm <$> readTVar peers
         writeTVar peers store'
         return ps
-      if null ps
-        then Left <$> error "TODO: getClosest ih"
-        else return (Right ps)
   tok <- grantToken toks naddr
-  return $ GotPeers ps tok
+  ns <- findNodeH routing naddr (FindNode (coerce ih) iptyp)
+  return $ GotPeers ps ns tok
 
 -- | Announce that the peer, controlling the querying node, is
 -- downloading a torrent on a port.