3 files changed, 16 insertions, 10 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index dcb217d6..07fb80c2 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -71,7 +71,10 @@ getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
 
 announceH :: Address ip => NodeHandler ip
 announceH = nodeHandler $ \ naddr @ NodeAddr {..} (Announce {..}) -> do
-  checkToken naddr sessionToken
+  valid <- checkToken naddr sessionToken
+  unless valid $ do
+    throwIO $ InvalidParameter "token"
+
   let annPort  = if impliedPort then nodePort else port
   let peerAddr = PeerAddr Nothing nodeHost annPort
   insertPeer topic peerAddr
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 6f13cf78..10f3e9e4 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -349,12 +349,11 @@ grantToken addr = do
 
 -- | Throws 'HandlerError' if the token is invalid or already
 -- expired. See 'TokenMap' for details.
-checkToken :: Hashable a => NodeAddr a -> Token -> DHT ip ()
+checkToken :: Hashable a => NodeAddr a -> Token -> DHT ip Bool
 checkToken addr questionableToken = do
   tryUpdateSecret
   toks <- asks sessionTokens >>= liftIO . readTVarIO
-  unless (T.member addr questionableToken (tokenMap toks)) $
-    throwIO $ InvalidParameter "token"
+  return $ T.member addr questionableToken (tokenMap toks)
 
 {-----------------------------------------------------------------------
 -- Routing table
diff --git a/tests/Network/BitTorrent/DHT/SessionSpec.hs b/tests/Network/BitTorrent/DHT/SessionSpec.hs
index 75c95462..3e197875 100644
--- a/tests/Network/BitTorrent/DHT/SessionSpec.hs
+++ b/tests/Network/BitTorrent/DHT/SessionSpec.hs
@@ -2,6 +2,7 @@
 module Network.BitTorrent.DHT.SessionSpec (spec) where
 import Control.Monad.Reader
 import Data.Default
+import Data.List as L
 import Test.Hspec
 import Test.QuickCheck
 
@@ -33,14 +34,17 @@ spec = do
   describe "tokens" $ do
     it "should not complain about valid token" $
       property $ \ (addrs :: [NodeAddr IPv4]) -> do
-        simpleDHT $ do
-          forM_ addrs $ \ addr -> do
+        isOks <- simpleDHT $ do
+          forM addrs $ \ addr -> do
             token <- grantToken addr
             checkToken addr token
-{-
-    it "" $ property $ \ (addr :: NodeAddr IPv4) token -> do
-      simpleDHT (checkToken addr token) `shouldThrow` (== undefined)
--}
+        L.and isOks `shouldBe` True
+
+    it "should complain about invalid token" $
+      property $ \ (addr :: NodeAddr IPv4) token -> do
+        isOk <- simpleDHT (checkToken addr token)
+        isOk `shouldBe` False
+
   describe "routing table" $ do
     return ()