Return Bool instead of raising exception in checkToken

2 files changed, 6 insertions, 4 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