Store external port number, as well as IP.

1 files changed, 11 insertions, 11 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index e067ab52..44083d81 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -232,7 +232,7 @@ refreshNodes nid = do
 
 -- | This operation do not block but acquire exclusive access to
 --   routing table.
-insertNode :: Address ip => NodeInfo ip -> Maybe ReflectedIP -> DHT ip ThreadId
+insertNode :: forall ip. Address ip => NodeInfo ip -> Maybe ReflectedIP -> DHT ip ThreadId
 insertNode info witnessed_ip0 = fork $ do
   var <- asks routingInfo
   tm <- getTimestamp
@@ -247,23 +247,24 @@ insertNode info witnessed_ip0 = fork $ do
   fallbackid <- asks tentativeNodeId
   let atomicInsert arrival witnessed_ip = do
           minfo <- readTVar var
-          let change ip = fromMaybe fallbackid
-                            $ listToMaybe
-                            $ rank id (nodeId $ foreignNode arrival)
-                            $ bep42s ip fallbackid
+          let change ip0 = fromMaybe fallbackid $ do
+                            ip <- fromSockAddr ip0 :: Maybe ip
+                            listToMaybe
+                                $ rank id (nodeId $ foreignNode arrival)
+                                $ bep42s ip fallbackid
           case minfo of
             Just inf -> do
               (ps,t') <- R.insert tm arrival $ myBuckets inf
               writeTVar var $ Just $ inf { myBuckets = t' }
               return $ do
                 case witnessed_ip of
-                    Just (ReflectedIP ip0)
-                      | fromSockAddr ip0 /= Just (myAddress inf)
+                    Just (ReflectedIP ip)
+                      | ip /= myAddress inf
                         -> $(logInfo) ( T.pack $ L.unwords
                               $ [ "Possible NAT?"
                                 , show (toSockAddr $ nodeAddr $ foreignNode arrival)
                                 , "reports my address:"
-                                , show ip0 ] )
+                                , show ip ] )
                             -- TODO: Let routing table vote on my IP/NodeId.
                     _   -> return ()
                 return ps
@@ -275,8 +276,7 @@ insertNode info witnessed_ip0 = fork $ do
                           <> T.pack (show (toSockAddr $ nodeAddr $ foreignNode arrival)))
                       return []
                 in fromMaybe dropped $ do
-                      ReflectedIP ip0 <- witnessed_ip
-                      ip <- fromSockAddr ip0
+                      ReflectedIP ip <- witnessed_ip
                       let nil = nullTable (change ip) maxbuckets
                       return $ do
                         (ps,t') <- R.insert tm arrival nil
@@ -285,7 +285,7 @@ insertNode info witnessed_ip0 = fork $ do
                         return $ do
                            $(logInfo) ( T.pack $ L.unwords
                                 [ "External IP address:"
-                                , show ip0
+                                , show ip
                                 , "(reported by"
                                 , show (toSockAddr $ nodeAddr $ foreignNode arrival)
                                     <> ")"