path: root/HTTP/RateLimit.hs

{- Copyright 2016 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

module HTTP.RateLimit where

import Types.Cost
import HTTP
import HTTP.ProofOfWork
import Tunables
import Servant
import Control.Concurrent.STM
import Control.Concurrent.TokenBucket
import qualified Data.BloomFilter.Mutable as BloomFilter
import qualified Data.BloomFilter.Hash as BloomFilter
import Data.BloomFilter.Easy (suggestSizing)
import Control.Monad.ST
import Control.Exception.Lifted (bracket)
import System.IO
import Data.Maybe
import Data.Word
import Control.Monad.IO.Class

-- | A rate limiter is a series of buckets. Each bucket has a
-- successively more difficult proof of work access requirement.
-- 
-- To guard against DOS attacks that reuse the same proof of work,
-- RandomSalt values are used, and bloom filters keep track of
-- the ones that have been assigned and used.
data RateLimiter = RateLimiter
	{ buckets :: TMVar [Bucket]
	, unusedBuckets :: TMVar [Bucket]
	, assignedRandomSalts :: BloomFilter
	, assignedRandomSaltsOld :: BloomFilter
	, usedRandomSalts :: BloomFilter
	, usedRandomSaltsOld :: BloomFilter
	, numRandomSalts :: TMVar Int
	, randomSaltGenerationLimiter :: TokenBucket
	, blockedRequestQueue :: TMVar [()]
	}

type BloomFilter = TMVar (BloomFilter.MBloom RealWorld RandomSalt)

-- | Buckets fill up at a fixed rate, and accessing a bucket
-- removes one unit from it.
data Bucket = Bucket
	{ tokenBucket :: TokenBucket
	, proofOfWorkRequired :: Seconds
	, fillInterval :: Word64
	}

minFillInterval :: Word64
minFillInterval = 2 * 60 * 1000000 -- 1 token every other minute

-- | Size of the bucket. This allows a burst of accesses after an idle
-- period, which is especially useful when retrieving keys that were
-- split into multiple chunks. However, setting this too high lets clients
-- cheaply store lots of data on a server that has been idle for a while,
-- which could be an attractive way to abuse keysafe servers.
burstSize :: Word64
burstSize = 4 -- 256 kb immediate storage

newRateLimiter :: IO RateLimiter
newRateLimiter = RateLimiter 
	<$> (newTMVarIO =<< mkbuckets (sdiv maxProofOfWork 2) [])
	<*> newTMVarIO []
	<*> mkBloomFilter
	<*> mkBloomFilter
	<*> mkBloomFilter
	<*> mkBloomFilter
	<*> newTMVarIO 0
	<*> newTokenBucket
	<*> newTMVarIO []
  where
	-- The last bucket takes half of maxProofOfWork to access, and
	-- each earlier bucket quarters that time, down to the first bucket,
	-- which needs no proof of work. This ensures that in the edge case
	-- where a client keeps getting bumped up to more and more expensive
	-- buckets, it doesn't need to do more than maxProofOfWork total work.
	mkbuckets s@(Seconds n) bs
		| n <= 0 = finalbucket bs
		| otherwise = do
			case mkProofOfWorkRequirement s of
				Nothing -> finalbucket bs
				Just _ -> do
					b <- Bucket
						<$> newTokenBucket
						<*> pure s
						<*> pure minFillInterval
					mkbuckets (sdiv s 4) (b:bs)
	finalbucket bs = do
		b <- Bucket
			<$> newTokenBucket
			<*> pure (Seconds 0)
			<*> pure minFillInterval
		return (b:bs)
		
	sdiv (Seconds n) d = Seconds (n `div` d)

mkBloomFilter :: IO BloomFilter
mkBloomFilter = do
	b <- stToIO $ BloomFilter.new (BloomFilter.cheapHashes bloomhashes) bloomsize
	newTMVarIO b
  where
	-- Size the bloom filter to hold 1 million items, with a false
	-- positive rate of 1 in 100 thousand. This will use around 32 mb
	-- of memory.
	(bloomhashes, bloomsize) = suggestSizing bloomMaxSize (1/100000)

-- | Maximum number of RandomSalts that can be stored in a bloom filter
-- without the false positive rate getting bad.
bloomMaxSize :: Int
bloomMaxSize = 1000000

-- A request is tried in each bucket in turn which its proof of work allows
-- access to, until one is found that accepts it.
rateLimit :: POWIdent p => RateLimiter -> Maybe ProofOfWork -> p -> Handler a -> Handler (POWGuarded a)
rateLimit ratelimiter mpow p a = do
	validsalt <- liftIO $ checkValidSalt ratelimiter mpow
	bs <- getBuckets ratelimiter
	if validsalt
		then go bs
		else assignWork ratelimiter bs
  where
	go [] = allBucketsEmpty ratelimiter a
	go (b:bs) = case mkProofOfWorkRequirement (proofOfWorkRequired b) of
		Nothing -> checkbucket b bs
		Just mkreq -> case mpow of
			Nothing -> assignWork ratelimiter (b:bs)
			Just pow@(ProofOfWork _ salt) -> 
				if isValidProofOfWork pow (mkreq salt) p
					then checkbucket b bs
					else assignWork ratelimiter (b:bs)
	checkbucket b bs = do
		allowed <- liftIO $ tokenBucketTryAlloc (tokenBucket b) 
			burstSize (fillInterval b) 1
		if allowed
			then Result <$> a
			else go bs

checkValidSalt :: RateLimiter -> Maybe ProofOfWork -> IO Bool
checkValidSalt _ Nothing = return True
checkValidSalt rl (Just (ProofOfWork _ salt)) = do
	assigned <- iselem assignedRandomSalts
	oldassigned <- iselem assignedRandomSaltsOld
	used <- iselem usedRandomSalts
	oldused <- iselem usedRandomSaltsOld
	if assigned && not oldassigned && not used && not oldused
		then do
			withBloomFilter rl usedRandomSalts
				(`BloomFilter.insert` salt)
			return True
		else return False
  where
	iselem f = withBloomFilter rl f (BloomFilter.elem salt)

assignWork :: RateLimiter -> [Bucket] -> Handler (POWGuarded a)
assignWork ratelimiter bs = case mapMaybe (mkProofOfWorkRequirement . proofOfWorkRequired) bs of
	[] -> throwError err404
	(mkreq:_) -> liftIO $ do
		-- 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.
		tokenBucketWait (randomSaltGenerationLimiter ratelimiter)
			100 -- burst
			1000000 -- refill 1 token per second

		salt <- liftIO mkRandomSalt
		withBloomFilter ratelimiter assignedRandomSalts
			(`BloomFilter.insert` salt)
		needrot <- atomically $ do
			n <- takeTMVar (numRandomSalts ratelimiter)
			if n > bloomMaxSize `div` 2
				then return Nothing
				else do
					putTMVar (numRandomSalts ratelimiter) (n+1)
					return (Just n)
		handlerotation needrot
		return $ NeedProofOfWork $ mkreq salt
  where
	handlerotation Nothing = return ()
	handlerotation (Just n) = do
		hPutStrLn stderr $ "rotating bloom filters after processing " ++ show n ++ " requests"
		newassigned <- mkBloomFilter
		newused <- mkBloomFilter
		atomically $ do
			oldassigned <- takeTMVar (assignedRandomSalts ratelimiter)
			oldused <- takeTMVar (usedRandomSalts ratelimiter)
			putTMVar (assignedRandomSaltsOld ratelimiter) oldassigned
			putTMVar (usedRandomSaltsOld ratelimiter) oldused
			putTMVar (assignedRandomSalts ratelimiter) =<< takeTMVar newassigned
			putTMVar (usedRandomSalts ratelimiter) =<< takeTMVar newused
			putTMVar (numRandomSalts ratelimiter) 0

withBloomFilter
	:: RateLimiter
	-> (RateLimiter -> BloomFilter)
	-> (BloomFilter.MBloom RealWorld RandomSalt -> ST RealWorld a)
	-> IO a
withBloomFilter rl field a = do
	b <- atomically $ readTMVar (field rl)
	stToIO (a b)

getBuckets :: MonadIO m => RateLimiter -> m [Bucket]
getBuckets = liftIO . atomically . readTMVar . buckets

putBuckets :: MonadIO m => RateLimiter -> [Bucket] -> m ()
putBuckets rl bs = liftIO $ atomically $ do
	_ <- takeTMVar (buckets rl)
	putTMVar (buckets rl) bs

-- When all buckets are empty, and the client has provided a good enough
-- proof of work to access the last bucket, the request is still processed, 
-- but blocked until the last token bucket refills. 
--
-- Only 100 requests are allowed to block this way at a time, since the
-- requests are taking up server memory while blocked, and because we don't
-- want to stall legitimate clients too long.
allBucketsEmpty :: RateLimiter -> Handler a -> Handler (POWGuarded a)
allBucketsEmpty ratelimiter a = bracket (liftIO addq) (liftIO . removeq) go
  where
	q = blockedRequestQueue ratelimiter

	addq = liftIO $ atomically $ do
		l <- takeTMVar q
		if length l >= 100
			then do
				putTMVar q l
				return False
			else do
				putTMVar q (():l)
				return True

	removeq False = return ()
	removeq True = liftIO $ atomically $ do
		l <- takeTMVar q
		putTMVar q (drop 1 l)
	
	waitlast = do
		bs <- getBuckets ratelimiter
		case reverse bs of
			(lastb:_) -> do
				hPutStrLn stderr "** warning: All token buckets are empty. Delaying request.."
				tokenBucketWait (tokenBucket lastb) burstSize (fillInterval lastb)
				return True
			[] -> return False

	go False = giveup
	go True = do
		ok <- liftIO waitlast
		if ok
			then Result <$> a
			else giveup
	
	giveup = do
		liftIO $ hPutStrLn stderr "** warning: All token buckets are empty and request queue is large; possible DOS attack? Rejected request.."
		assignWork ratelimiter =<< getBuckets ratelimiter

-- | How much data could be stored, in bytes per second, assuming all
-- buckets in the rate limiter being constantly drained by requests,
-- and all requests store objects.
maximumStorageRate :: RateLimiter -> IO Int
maximumStorageRate ratelimiter = do
	bs <- getBuckets ratelimiter
	return $ sum $ map calc bs
  where
	storesize = maximum knownObjectSizes
	calc b = (storesize * 1000000) `div` fromIntegral (fillInterval b)

describeRateLimiter :: RateLimiter -> IO String
describeRateLimiter ratelimiter = do
	storerate <- maximumStorageRate ratelimiter
	bs <- getBuckets ratelimiter
	return $ concat
		[ "rate limiter buckets: " ++ show bs
		, " ; maximum allowed storage rate: "
		, showrate (storerate * 60 * 60 * 24 * 31) ++ "/month"
		]
  where
	showrate n
		| n < 1024*1024 = show (n `div` 1024) ++ " KiB"
		| n < 1024*1024*1024 = show (n `div` (1024 * 1024)) ++ " MiB"
		| otherwise = show (n `div` (1024 * 1024 * 1024)) ++ " GiB"

instance Show Bucket where
	show b = show (fillInterval b `div` (60 * 1000000)) ++ " Second/Request"
		++ " (PoW=" ++ show (proofOfWorkRequired b) ++ ")"

increaseDifficulty :: RateLimiter -> IO ()
increaseDifficulty ratelimiter = do
	bs <- getBuckets ratelimiter
	case bs of
		[] -> unable
		(b:[]) -> do
			-- Make the remaining bucket take longer to fill.
			let b' = b { fillInterval = fillInterval b * 2 }
			putBuckets ratelimiter [b']
			done
		(b:rest) -> do
			-- Remove less expensive to access buckets,
			-- so that clients have to do some work.
			-- This is done first to cut off any freeloaders
			-- that may be abusing the keysafe server.
			atomically $ do
				unused <- takeTMVar (unusedBuckets ratelimiter)
				putTMVar (unusedBuckets ratelimiter) (b:unused)
			putBuckets ratelimiter rest
			done
  where
	unable = putStrLn "unable to increase difficulty; out of buckets"
	done = do
		desc <- describeRateLimiter ratelimiter
		putStrLn $ "increased difficulty -- " ++ desc

-- Should undo the effect of increaseDifficulty.
reduceDifficulty :: RateLimiter -> IO ()
reduceDifficulty ratelimiter = do
	bs <- getBuckets ratelimiter
	case bs of
		(b:[]) | fillInterval b > minFillInterval -> do
			let b' = b { fillInterval = fillInterval b `div` 2 }
			putBuckets ratelimiter [b']
			done
		_ -> do
			mb <- getunused
			case mb of
				Nothing -> unable
				Just b -> do
					putBuckets ratelimiter (b:bs)
					done
  where
	getunused = atomically $ do
		unused <- takeTMVar (unusedBuckets ratelimiter)
		case unused of
			(b:bs) -> do
				putTMVar (unusedBuckets ratelimiter) bs
				return (Just b)
			[] -> do
				putTMVar (unusedBuckets ratelimiter) []
				return Nothing
	unable = return ()
	done = do
		desc <- describeRateLimiter ratelimiter
		putStrLn $ "reduced difficulty -- " ++ desc