1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
|
{- 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 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]
, assignedRandomSalts :: BloomFilter
, assignedRandomSaltsOld :: BloomFilter
, usedRandomSalts :: BloomFilter
, usedRandomSaltsOld :: BloomFilter
, numRandomSalts :: TMVar Int
, randomSaltGenerationLimiter :: TokenBucket
}
type BloomFilter = TMVar (BloomFilter.MBloom RealWorld RandomSalt)
-- | Buckets fill up at a fixed rate (which can be tuned by the server
-- operator), and accessing a bucket removes one unit from it.
data Bucket = Bucket
{ tokenBucket :: TokenBucket
, mkProofReq :: Maybe (RandomSalt -> ProofOfWorkRequirement)
}
newRateLimiter :: IO RateLimiter
newRateLimiter = RateLimiter
<$> (newTMVarIO =<< mkbuckets maxProofOfWork [])
<*> mkBloomFilter
<*> mkBloomFilter
<*> mkBloomFilter
<*> mkBloomFilter
<*> newTMVarIO 0
<*> newTokenBucket
where
-- The last bucket takes half of maxProofOfWork to access, and
-- each earlier bucket halves 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 (Seconds n) bs
| n <= 0 = return bs
| otherwise = do
let s = Seconds (n `div` 2)
let mkreq = mkProofOfWorkRequirement s
b <- Bucket
<$> newTokenBucket
<*> pure mkreq
case mkreq of
Nothing -> return (b:bs)
Just _ -> mkbuckets s (b:bs)
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
-- | 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, so keep the objectSize in mind.
burstSize :: Word64
burstSize = 4 -- allow 4 objects to be stored/retrieved w/o proof of work
-- | Rate that the bucket is filled.
fillRate :: Word64
fillRate = 60000000 -- 1 token per minute
-- | How much data could be stored, in bytes per second, assuming all
-- buckets in the rate limiter are kept drained, and all requests are
-- stores.
maximumStorageRate :: RateLimiter -> IO Int
maximumStorageRate ratelimiter = do
let storesize = maximum knownObjectSizes
bs <- liftIO $ atomically $ readTMVar (buckets ratelimiter)
return $ (length bs * storesize * 1000000) `div` fromIntegral fillRate
-- A request is tried in each bucket in turn which its proof of work allows
-- access to. If all accessible token buckets are empty, generate a
-- new ProofOfWorkRequirement for the client.
--
-- If all buckets are tried and are empty, we must be very overloaded.
-- In this case, the request is still processed, since the client has done
-- quite a lot of work.
rateLimit :: POWIdent p => RateLimiter -> Maybe ProofOfWork -> p -> Handler a -> Handler (POWGuarded a)
rateLimit ratelimiter mpow p a = do
validsalt <- liftIO $ checkValidSalt ratelimiter mpow
bs <- liftIO $ atomically $ readTMVar (buckets ratelimiter)
if validsalt
then go bs
else assignWork ratelimiter bs
where
go [] = do
liftIO $ hPutStrLn stderr "** warning: all token buckets are empty; possible DOS attack?"
Result <$> a
go (b:bs) = case mkProofReq 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 fillRate 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 mkProofReq 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
100000 -- 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)
|