1 files changed, 54 insertions, 4 deletions

diff --git a/src/Network/BitTorrent/DHT/Token.hs b/src/Network/BitTorrent/DHT/Token.hs
index e3a6b1f6..08079b75 100644
--- a/src/Network/BitTorrent/DHT/Token.hs
+++ b/src/Network/BitTorrent/DHT/Token.hs
@@ -24,6 +24,10 @@ module Network.BitTorrent.DHT.Token
 
          -- * Session tokens
        , TokenMap
+       , SessionTokens
+       , nullSessionTokens
+       , checkToken
+       , grantToken
 
          -- ** Construction
        , Network.BitTorrent.DHT.Token.tokens
@@ -52,7 +56,7 @@ import Data.Hashable
 import Data.String
 import Data.Time
 import System.Random
-
+import Control.Concurrent.STM
 import Network.Address
 
 -- TODO use ShortByteString
@@ -77,7 +81,7 @@ type Secret = Int
 
 -- The BitTorrent implementation uses the SHA1 hash of the IP address
 -- concatenated onto a secret, we use hashable instead.
-makeToken :: Hashable a => NodeAddr a -> Secret -> Token
+makeToken :: Hashable a => a -> Secret -> Token
 makeToken n s = Token $ toBS $ hashWithSalt s n
   where
     toBS = toStrict . toLazyByteString . int64BE . fromIntegral
@@ -104,7 +108,7 @@ tokens seed = (`evalState` mkStdGen seed) $
 -- 'update's.
 --
 -- Typically used to handle find_peers query.
-lookup :: Hashable a => NodeAddr a -> TokenMap -> Token
+lookup :: Hashable a => a -> TokenMap -> Token
 lookup addr TokenMap {..} = makeToken addr curSecret
 
 -- | Check if token is valid.
@@ -112,7 +116,7 @@ lookup addr TokenMap {..} = makeToken addr curSecret
 -- Typically used to handle 'Network.DHT.Mainline.Announce'
 -- query. If token is invalid the 'Network.KRPC.ProtocolError' should
 -- be sent back to the malicious node.
-member :: Hashable a => NodeAddr a -> Token -> TokenMap -> Bool
+member :: Hashable a => a -> Token -> TokenMap -> Bool
 member addr token TokenMap {..} = token `L.elem` valid
   where valid = makeToken addr <$> [curSecret, prevSecret]
 
@@ -130,3 +134,49 @@ update TokenMap {..} = TokenMap
     }
   where
     (newSecret, newGen) = next generator
+
+data SessionTokens = SessionTokens
+  { tokenMap    :: !TokenMap
+  , lastUpdate  :: !UTCTime
+  , maxInterval :: !NominalDiffTime
+  }
+
+nullSessionTokens :: IO SessionTokens
+nullSessionTokens = SessionTokens
+  <$> (tokens <$> randomIO)
+  <*> getCurrentTime
+  <*> pure defaultUpdateInterval
+
+-- TODO invalidate *twice* if needed
+invalidateTokens :: UTCTime -> SessionTokens -> SessionTokens
+invalidateTokens curTime ts @ SessionTokens {..}
+  | curTime `diffUTCTime` lastUpdate > maxInterval = SessionTokens
+    { tokenMap    = update tokenMap
+    , lastUpdate  = curTime
+    , maxInterval = maxInterval
+    }
+  |                  otherwise                     = ts
+
+{-----------------------------------------------------------------------
+-- Tokens
+-----------------------------------------------------------------------}
+
+tryUpdateSecret :: TVar SessionTokens -> IO ()
+tryUpdateSecret toks = do
+  curTime <- getCurrentTime
+  atomically $ modifyTVar' toks (invalidateTokens curTime)
+
+grantToken :: Hashable addr => TVar SessionTokens -> addr -> IO Token
+grantToken sessionTokens addr = do
+  tryUpdateSecret sessionTokens
+  toks <- readTVarIO sessionTokens
+  return $ Network.BitTorrent.DHT.Token.lookup addr $ tokenMap toks
+
+-- | Throws 'HandlerError' if the token is invalid or already
+-- expired. See 'TokenMap' for details.
+checkToken :: Hashable addr => TVar SessionTokens -> addr -> Token -> IO Bool
+checkToken sessionTokens addr questionableToken = do
+  tryUpdateSecret sessionTokens
+  toks <- readTVarIO sessionTokens
+  return $ member addr questionableToken (tokenMap toks)
+