use token bucket as the random salt generation rate limiter

This avoids a 1s delay in requests, except when an attacker is flooding
them.

1 files changed, 9 insertions, 11 deletions

diff --git a/HTTP/RateLimit.hs b/HTTP/RateLimit.hs
index 737f7dc..81e3ed2 100644
--- a/HTTP/RateLimit.hs
+++ b/HTTP/RateLimit.hs
@@ -9,7 +9,6 @@ import Types.Cost
 import HTTP
 import HTTP.ProofOfWork
 import Servant
-import Control.Concurrent
 import Control.Concurrent.STM
 import Control.Concurrent.TokenBucket
 import qualified Data.BloomFilter.Mutable as BloomFilter
@@ -34,7 +33,7 @@ data RateLimiter = RateLimiter
 	, usedRandomSalts :: BloomFilter
 	, usedRandomSaltsOld :: BloomFilter
 	, numRandomSalts :: TMVar Int
-	, accessLock :: TMVar ()
+	, randomSaltGenerationLimiter :: TokenBucket
 	}
 
 type BloomFilter = TMVar (BloomFilter.MBloom RealWorld RandomSalt)
@@ -54,7 +53,7 @@ newRateLimiter = RateLimiter
 	<*> mkBloomFilter
 	<*> mkBloomFilter
 	<*> newTMVarIO 0
-	<*> newTMVarIO ()
+	<*> newTokenBucket
   where
 	-- The last bucket takes half of maxProofOfWork to access, and
 	-- each earlier bucket halves that time, down to the first bucket,
@@ -152,18 +151,17 @@ assignWork :: RateLimiter -> [Bucket] -> Handler (POWGuarded a)
 assignWork ratelimiter bs = case mapMaybe mkProofReq bs of
 	[] -> throwError err404
 	(mkreq:_) -> liftIO $ do
-		-- 10 second pause here, with a lock held so that only
-		-- this thread can run. This prevents an attacker
-		-- flooding requests that cause new random salts to be
-		-- assigned, in order to fill up the bloom table and cause
-		-- salts assigned to other clients to be rejected.
+		-- This prevents an attacker flooding requests that 
+		-- cause new random salts to be assigned, in order 
+		-- to fill up the bloom table and cause salts assigned 
+		-- to other clients to be rejected.
 		-- Since the bloom filters hold 1 million salts,
 		-- the attacker would need to send requests for over 10
 		-- hours to force a bloom filter rotation, so would not
 		-- impact many users.
-		atomically $ takeTMVar (accessLock ratelimiter)
-		threadDelay 10000000
-		atomically $ putTMVar (accessLock ratelimiter) ()
+		tokenBucketWait (randomSaltGenerationLimiter ratelimiter)
+			100 -- burst
+			100000 -- refill 1 token per second
 
 		salt <- liftIO mkRandomSalt
 		withBloomFilter ratelimiter assignedRandomSalts