2 files changed, 37 insertions, 11 deletions

diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 15d1099c..49629755 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -214,8 +214,15 @@ instance (Typeable ip, Serialize ip) =>
 -- | Announce that the peer, controlling the querying node, is
 -- downloading a torrent on a port.
 data Announce = Announce
-  { -- | infohash of the torrent;
-    topic    :: InfoHash
+  { -- | If set, the 'port' field should be ignored and the source
+    -- port of the UDP packet should be used as the peer's port
+    -- instead. This is useful for peers behind a NAT that may not
+    -- know their external port, and supporting uTP, they accept
+    -- incoming connections on the same port as the DHT port.
+    impliedPort :: Bool
+
+    -- | infohash of the torrent;
+  , topic    :: InfoHash
 
     -- | the port /this/ peer is listening;
   , port     :: PortNumber
@@ -227,16 +234,26 @@ data Announce = Announce
 port_key :: BKey
 port_key = "port"
 
+implied_port_key :: BKey
+implied_port_key = "implied_port"
+
 instance BEncode Announce where
   toBEncode Announce {..} = toDict $
-       info_hash_key .=! topic
-    .: port_key      .=! port
-    .: token_key     .=! sessionToken
+       implied_port_key .=? flagField impliedPort
+    .: info_hash_key    .=! topic
+    .: port_key         .=! port
+    .: token_key        .=! sessionToken
     .: endDict
+    where
+      flagField flag = if flag then Just (1 :: Int) else Nothing
+
   fromBEncode = fromDict $ do
-    Announce <$>! info_hash_key
+    Announce <$> (boolField <$> optional (field (req implied_port_key)))
+             <*>! info_hash_key
              <*>! port_key
              <*>! token_key
+    where
+      boolField = maybe False (/= (0 :: Int))
 
 -- | The queried node must verify that the token was previously sent
 -- to the same IP address as the querying node. Then the queried node
diff --git a/tests/Network/BitTorrent/DHT/MessageSpec.hs b/tests/Network/BitTorrent/DHT/MessageSpec.hs
index 52bda932..11a8fb15 100644
--- a/tests/Network/BitTorrent/DHT/MessageSpec.hs
+++ b/tests/Network/BitTorrent/DHT/MessageSpec.hs
@@ -143,15 +143,24 @@ spec = do
                  \5:token8:aoeusnth\
                 \e" `shouldBe` Right
         (Query "abcdefghij0123456789"
-                   (Announce "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
+                   (Announce False "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
+
+      BE.decode "d2:id20:abcdefghij0123456789\
+                 \12:implied_porti1e\
+                 \9:info_hash20:mnopqrstuvwxyz123456\
+                 \4:porti6881e\
+                 \5:token8:aoeusnth\
+                \e" `shouldBe` Right
+        (Query "abcdefghij0123456789"
+                   (Announce True "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
 
 
       BE.decode "d2:id20:mnopqrstuvwxyz123456e"
         `shouldBe` Right
         (Response "mnopqrstuvwxyz123456" Announced)
 
-    it "properly bencoded (iso)" $ property $ \ nid topic port token -> do
-      prop_bencode (Query nid    (Announce topic port token))
+    it "properly bencoded (iso)" $ property $ \ nid flag topic port token -> do
+      prop_bencode (Query    nid (Announce flag topic port token))
       prop_bencode (Response nid (Announced))
 
 
@@ -160,7 +169,7 @@ spec = do
       Response _remoteId Announced <- rpc $ do
         Response _ GotPeers {..} <- query remoteAddr (Query nid (GetPeers def))
         let _ = peers :: Either [NodeInfo IPv4] [PeerAddr IPv4]
-        query remoteAddr (Query nid (Announce def thisPort grantedToken))
+        query remoteAddr (Query nid (Announce False def thisPort grantedToken))
       return ()
 
     it "does fail on invalid token" $ do
@@ -169,6 +178,6 @@ spec = do
         Response _ GotPeers {..} <- query remoteAddr (Query nid (GetPeers def))
         let _ = peers :: Either [NodeInfo IPv4] [PeerAddr IPv4]
         let invalidToken = ""
-        query remoteAddr (Query nid (Announce def thisPort invalidToken)))
+        query remoteAddr (Query nid (Announce False def thisPort invalidToken)))
           `shouldThrow` isProtocolError
       return ()