-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #34 from privacybydesign/security/ratelimiting
Ratelimiting
- Loading branch information
Showing
8 changed files
with
303 additions
and
7 deletions.
There are no files selected for viewing
43 changes: 43 additions & 0 deletions
43
src/main/java/foundation/privacybydesign/email/CleanupBackgroundJob.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,43 @@ | ||
package foundation.privacybydesign.email; | ||
|
||
import foundation.privacybydesign.email.ratelimit.MemoryRateLimit; | ||
import org.slf4j.Logger; | ||
import org.slf4j.LoggerFactory; | ||
|
||
import jakarta.servlet.ServletContextEvent; | ||
import jakarta.servlet.ServletContextListener; | ||
import jakarta.servlet.annotation.WebListener; | ||
import java.util.concurrent.Executors; | ||
import java.util.concurrent.ScheduledExecutorService; | ||
import java.util.concurrent.TimeUnit; | ||
|
||
/** | ||
* Clean up in-memory data structures once in a while. | ||
*/ | ||
@WebListener | ||
public class CleanupBackgroundJob implements ServletContextListener { | ||
private static Logger logger = LoggerFactory.getLogger(CleanupBackgroundJob.class); | ||
private ScheduledExecutorService scheduler; | ||
|
||
@Override | ||
public void contextInitialized(ServletContextEvent event) { | ||
logger.info("Setting up background cleanup task"); | ||
scheduler = Executors.newSingleThreadScheduledExecutor(); | ||
|
||
scheduler.scheduleAtFixedRate(new Runnable() { | ||
@Override public void run() { | ||
try { | ||
MemoryRateLimit.getInstance().periodicCleanup(); | ||
} catch (Exception e) { | ||
logger.error("Failed to run periodic cleanup:"); | ||
e.printStackTrace(); | ||
} | ||
} | ||
}, 5, 5, TimeUnit.MINUTES); | ||
} | ||
|
||
@Override | ||
public void contextDestroyed(ServletContextEvent event) { | ||
scheduler.shutdownNow(); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
122 changes: 122 additions & 0 deletions
122
src/main/java/foundation/privacybydesign/email/ratelimit/MemoryRateLimit.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,122 @@ | ||
package foundation.privacybydesign.email.ratelimit; | ||
|
||
import java.util.Map; | ||
import java.util.concurrent.ConcurrentHashMap; | ||
|
||
/** | ||
* Store rate limits in memory. Useful for debugging and rate limits that | ||
* aren't very long. | ||
* | ||
* How it works: | ||
* How much budget a user has, is expressed in a timestamp. The timestamp is | ||
* initially some period in the past, but with every usage (countEmail) | ||
* this timestamp is incremented. For e-mail addresses this | ||
* amount is exponential. | ||
* | ||
* An algorithm with a similar goal is the Token Bucket algorithm. This | ||
* algorithm probably works well, but seemed harder to implement. | ||
* https://en.wikipedia.org/wiki/Token_bucket | ||
*/ | ||
public class MemoryRateLimit extends RateLimit { | ||
private static final long SECOND = 1000; // 1000ms = 1s | ||
private static final long MINUTE = SECOND * 60; | ||
private static final long HOUR = MINUTE * 60; | ||
private static final long DAY = HOUR * 24; | ||
|
||
private static MemoryRateLimit instance; | ||
|
||
private final Map<String, Limit> emailLimits; | ||
|
||
public MemoryRateLimit() { | ||
emailLimits = new ConcurrentHashMap<>(); | ||
} | ||
|
||
public static MemoryRateLimit getInstance() { | ||
if (instance == null) { | ||
instance = new MemoryRateLimit(); | ||
} | ||
return instance; | ||
} | ||
|
||
// Is the user over the rate limit per e-mail address? | ||
@Override | ||
protected synchronized long nextTryEmail(String email, long now) { | ||
// Rate limiter durations (sort-of logarithmic): | ||
// 1 10 seconds | ||
// 2 1 minute | ||
// 3 5 minutes | ||
// 4 1 hour | ||
// 5 12 hours | ||
// 6+ 2 per day | ||
// Keep log 2 days for proper limiting. | ||
Limit limit = emailLimits.get(email); | ||
if (limit == null) { | ||
limit = new Limit(now); | ||
emailLimits.put(email, limit); | ||
} | ||
long nextTry; // timestamp when the next request is allowed | ||
switch (limit.tries) { | ||
case 0: // try 1: always succeeds | ||
nextTry = limit.timestamp; | ||
break; | ||
case 1: // try 2: allowed after 10 seconds | ||
nextTry = limit.timestamp + 10 * SECOND; | ||
break; | ||
case 2: // try 3: allowed after 1 minute | ||
nextTry = limit.timestamp + MINUTE; | ||
break; | ||
case 3: // try 4: allowed after 5 minutes | ||
nextTry = limit.timestamp + 5 * MINUTE; | ||
break; | ||
case 4: // try 5: allowed after 1 hour | ||
nextTry = limit.timestamp + HOUR; | ||
break; | ||
case 5: // try 6: allowed after 12 hours | ||
nextTry = limit.timestamp + 12 * HOUR; | ||
break; | ||
default: | ||
throw new IllegalStateException("invalid tries count"); | ||
} | ||
return nextTry; | ||
} | ||
|
||
// Count the usage of this rate limit - adding to the budget for this | ||
// e-mail address. | ||
@Override | ||
protected synchronized void countEmail(String email, long now) { | ||
long nextTry = nextTryEmail(email, now); | ||
Limit limit = emailLimits.get(email); | ||
if (nextTry > now) { | ||
throw new IllegalStateException("counting rate limit while over the limit"); | ||
} | ||
limit.tries = Math.min(limit.tries+1, 6); // add 1, max at 6 | ||
// If the last usage was e.g. ≥2 days ago, we should allow them 2 | ||
// extra tries this day. | ||
long lastTryDaysAgo = (now-limit.timestamp)/DAY; | ||
long bonusTries = limit.tries - lastTryDaysAgo; | ||
if (bonusTries >= 1) { | ||
limit.tries = (int)bonusTries; | ||
} | ||
limit.timestamp = now; | ||
} | ||
|
||
public void periodicCleanup() { | ||
long now = System.currentTimeMillis(); | ||
// Use enhanced for loop, because an iterator makes sure concurrency issues cannot occur. | ||
for (Map.Entry<String, Limit> entry : emailLimits.entrySet()) { | ||
if (entry.getValue().timestamp < now - 2*DAY) { | ||
emailLimits.remove(entry.getKey()); | ||
} | ||
} | ||
} | ||
} | ||
|
||
class Limit { | ||
long timestamp; | ||
int tries; | ||
|
||
Limit(long now) { | ||
tries = 0; | ||
timestamp = now; | ||
} | ||
} |
34 changes: 34 additions & 0 deletions
34
src/main/java/foundation/privacybydesign/email/ratelimit/RateLimit.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,34 @@ | ||
package foundation.privacybydesign.email.ratelimit; | ||
|
||
import org.slf4j.Logger; | ||
import org.slf4j.LoggerFactory; | ||
|
||
/** | ||
* Base class for rate limiting. Subclasses provide storage methods (memory | ||
* for easier debugging and database for production). | ||
*/ | ||
public abstract class RateLimit { | ||
private static Logger logger = LoggerFactory.getLogger(RateLimit.class); | ||
|
||
/** Take an e-mail address and rate limit it. | ||
* @param email e-mail address | ||
* @return the number of milliseconds that the client should wait - 0 if | ||
* it shouldn't wait. | ||
*/ | ||
public long rateLimited(String email) { | ||
long now = System.currentTimeMillis(); | ||
long retryAfter = nextTryEmail(email, now); | ||
|
||
if (retryAfter > now) { | ||
logger.warn("Denying request: email rate limit email exceeded"); | ||
// Don't count this request if it has been denied. | ||
return retryAfter - now; | ||
} | ||
|
||
countEmail(email, now); | ||
return 0; | ||
} | ||
|
||
protected abstract long nextTryEmail(String email, long now); | ||
protected abstract void countEmail(String email, long now); | ||
} |
25 changes: 25 additions & 0 deletions
25
src/test/java/foundation/privacybydesign/email/RateLimitTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
package foundation.privacybydesign.email; | ||
|
||
import org.junit.Test; | ||
|
||
import foundation.privacybydesign.email.ratelimit.MemoryRateLimit; | ||
|
||
import static org.junit.Assert.assertEquals; | ||
import static org.junit.Assert.assertNotEquals; | ||
|
||
/** | ||
* Test whether the rate limiter works as expected. | ||
*/ | ||
public class RateLimitTest { | ||
|
||
@Test | ||
public void testRateLimit() { | ||
|
||
var rateLimit = MemoryRateLimit.getInstance(); | ||
var email = "[email protected]"; | ||
|
||
assertEquals("not rate limited", 0, rateLimit.rateLimited(email)); | ||
|
||
assertNotEquals("rate limited", 0, rateLimit.rateLimited(email)); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters