1 files changed, 40 insertions, 18 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 2ddd51ca..d1fa36e5 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -148,7 +148,8 @@ pingQ addr = do
 findNodeQ :: Address ip => TableKey key => key -> Iteration ip NodeInfo
 findNodeQ key NodeInfo {..} = do
   NodeFound closest <- FindNode (toNodeId key) <@> nodeAddr
-  $(logInfoS) "findNodeQ" $ "NodeFound " <> T.pack (show $ L.map pPrint closest)
+  $(logInfoS) "findNodeQ" $ "NodeFound\n"
+    <> T.pack (L.unlines $ L.map ((' ' :) . show . pPrint) closest)
   return $ Right closest
 
 getPeersQ :: Address ip => InfoHash -> Iteration ip PeerAddr
@@ -260,23 +261,44 @@ insertNode info witnessed_ip = fork $ do
           case minfo of
             Just info -> do
               (ps,t') <- R.insert tm arrival $ myBuckets info
-              -- TODO: Check witnessed_ip against myAddress.
-              -- If 3 nodes witness a different address, change the table.
-              -- Require these witnesses satisfy bep-42 and that their
-              -- first 3 bits are unique.
               writeTVar var $ Just $ info { myBuckets = t' }
-              return ps
-            -- Ignore non-witnessing nodes until somebody tells
-            -- us our ip address.
-            Nothing -> fromMaybe (return []) $ do
-              ReflectedIP ip0 <- witnessed_ip
-              ip <- fromSockAddr ip0
-              let nil = nullTable (change ip) maxbuckets
               return $ do
-                (ps,t') <- R.insert tm arrival nil
-                writeTVar var $ Just $ R.Info t' (change ip) ip
+                case witnessed_ip of
+                    Just (ReflectedIP ip0)
+                      | fromSockAddr ip0 /= Just (myAddress info)
+                        -> $(logInfo) ( T.pack $ L.unwords
+                              $ [ "Possible NAT?"
+                                , show (toSockAddr $ nodeAddr $ foreignNode arrival)
+                                , "reports my address:"
+                                , show ip0 ] )
+                            -- TODO: Let routing table vote on my IP/NodeId.
+                    _   -> return ()
                 return ps
-  ps <- liftIO $ atomically $ atomicInsert arrival witnessed_ip
+            Nothing ->
+                let dropped = return $ do
+                      -- Ignore non-witnessing nodes until somebody tells
+                      -- us our ip address.
+                      $(logWarnS) "insertNode" ("Dropped "
+                          <> T.pack (show (toSockAddr $ nodeAddr $ foreignNode arrival)))
+                      return []
+                in fromMaybe dropped $ do
+                      ReflectedIP ip0 <- witnessed_ip
+                      ip <- fromSockAddr ip0
+                      let nil = nullTable (change ip) maxbuckets
+                      return $ do
+                        (ps,t') <- R.insert tm arrival nil
+                        let new_info = R.Info t' (change ip) ip
+                        writeTVar var $ Just new_info
+                        return $ do
+                           $(logInfo) ( T.pack $ L.unwords
+                                [ "External IP address:"
+                                , show ip0
+                                , "(reported by"
+                                , show (toSockAddr $ nodeAddr $ foreignNode arrival)
+                                    <> ")"
+                                ] )
+                           return ps
+  ps <- join $ liftIO $ atomically $ atomicInsert arrival witnessed_ip
   showTable
   fork $ forM_ ps $ \(CheckPing ns)-> do
     forM_ ns $ \n -> do
@@ -284,7 +306,7 @@ insertNode info witnessed_ip = fork $ do
         let alive = PingResult n b
         $(logDebugS) "insertNode" $ T.pack ("PingResult "++show (nodeId n,b))
         tm <- getTimestamp
-        liftIO $ atomically $ atomicInsert alive mip
+        join $ liftIO $ atomically $ atomicInsert alive mip
         showTable
   return ()
 
@@ -299,8 +321,8 @@ queryNode' addr q = do
   nid <- myNodeIdAccordingTo addr
   let read_only = False -- TODO: check for NAT issues. (BEP 43)
   (Response remoteId r, witnessed_ip) <- query' (toSockAddr addr) (Query nid read_only q)
-  $(logInfoS) "queryNode" $ "Witnessed address: " <> T.pack (show witnessed_ip)
-                             <> " by " <> T.pack (show (toSockAddr addr))
+  -- $(logDebugS) "queryNode" $ "Witnessed address: " <> T.pack (show witnessed_ip)
+  --                             <> " by " <> T.pack (show (toSockAddr addr))
   insertNode (NodeInfo remoteId addr) witnessed_ip
   return (remoteId, r, witnessed_ip)