1 files changed, 24 insertions, 17 deletions

diff --git a/Mainline.hs b/Mainline.hs
index e3a48976..19646aeb 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE DeriveTraversable          #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase                 #-}
 {-# LANGUAGE PatternSynonyms            #-}
 {-# LANGUAGE StandaloneDeriving         #-}
 module Mainline where
@@ -382,12 +383,13 @@ is4mapped ip
     | otherwise = False
 
 prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP
-prefer4or6 addr iptyp = fromMaybe unspecified iptyp
- where
-    unspecified = case nodeIP addr of
-        IPv4 _               -> Want_IP4
-        IPv6 a | is4mapped a -> Want_IP4
-               | otherwise   -> Want_IP6
+prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp
+
+ipFamily :: IP -> WantIP
+ipFamily ip = case ip of
+    IPv4 _               -> Want_IP4
+    IPv6 a | is4mapped a -> Want_IP4
+           | otherwise   -> Want_IP6
 
 findNodeH :: Routing -> NodeInfo -> FindNode -> IO NodeFound
 findNodeH routing addr (FindNode node iptyp) = do
@@ -438,8 +440,7 @@ instance BEncode GotPeers where
     .: 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.)
+    .: peers_key .=! map S.encode peers
     .: endDict
 
   fromBEncode = fromDict $ do
@@ -453,15 +454,21 @@ instance BEncode GotPeers where
 
 getPeersH :: Routing -> SwarmsDatabase -> NodeInfo -> GetPeers -> IO GotPeers
 getPeersH routing (SwarmsDatabase peers toks _) naddr (GetPeers ih iptyp) = do
-  ps <- do
-      tm <- getTimestamp
-      atomically $ do
-        (ps,store') <- Peers.freshPeers ih tm <$> readTVar peers
-        writeTVar peers store'
-        return ps
-  tok <- grantToken toks naddr
-  ns <- findNodeH routing naddr (FindNode (coerce ih) iptyp)
-  return $ GotPeers ps ns tok
+    ps <- do
+        tm <- getTimestamp
+        atomically $ do
+          (ps,store') <- Peers.freshPeers ih tm <$> readTVar peers
+          writeTVar peers store'
+          return ps
+    -- Filter peer results to only a single address family, IPv4 or IPv6, as
+    -- per BEP 32.
+    let notboth  = iptyp >>= \case Want_Both -> Nothing
+                                   specific  -> Just specific
+        selected = prefer4or6 naddr notboth
+        ps'      = filter ( (== selected) . ipFamily . peerHost ) ps
+    tok <- grantToken toks naddr
+    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.