diff --git a/core/src/main/java/com/github/games647/fastlogin/core/RateLimiter.java b/core/src/main/java/com/github/games647/fastlogin/core/RateLimiter.java index cb228892..6cd5dc4e 100644 --- a/core/src/main/java/com/github/games647/fastlogin/core/RateLimiter.java +++ b/core/src/main/java/com/github/games647/fastlogin/core/RateLimiter.java @@ -25,6 +25,8 @@ */ package com.github.games647.fastlogin.core; +import com.google.common.base.Ticker; + import java.util.Arrays; /** @@ -33,16 +35,22 @@ import java.util.Arrays; */ public class RateLimiter { + private final Ticker ticker; + private final long[] requests; private final long expireTime; private int position; - public RateLimiter(int maxLimit, long expireTime) { + public RateLimiter(Ticker ticker, int maxLimit, long expireTime) { + this.ticker = ticker; + this.requests = new long[maxLimit]; this.expireTime = expireTime; - // fill the array with the lowest values, so that the first uninitialized values will always expire - Arrays.fill(requests, Long.MIN_VALUE); + // fill the array with expired entry, because nanoTime could overflow and include negative numbers + long nowMilli = ticker.read() / 1_000_000; + long initialVal = nowMilli - expireTime; + Arrays.fill(requests, initialVal); } /** @@ -52,15 +60,12 @@ public class RateLimiter { */ public boolean tryAcquire() { // current time millis is not monotonic - it can jump back depending on user choice or NTP - long now = System.nanoTime() / 1_000_000; - - // after this the request should be expired - long toBeExpired = now - expireTime; + long nowMilli = ticker.read() / 1_000_000; synchronized (this) { // having synchronized will limit the amount of concurrency a lot long oldest = requests[position]; - if (oldest < toBeExpired) { - requests[position] = now; + if (nowMilli - oldest >= expireTime) { + requests[position] = nowMilli; position = (position + 1) % requests.length; return true; } diff --git a/core/src/main/java/com/github/games647/fastlogin/core/shared/FastLoginCore.java b/core/src/main/java/com/github/games647/fastlogin/core/shared/FastLoginCore.java index f08b47d8..630873f1 100644 --- a/core/src/main/java/com/github/games647/fastlogin/core/shared/FastLoginCore.java +++ b/core/src/main/java/com/github/games647/fastlogin/core/shared/FastLoginCore.java @@ -35,6 +35,7 @@ import com.github.games647.fastlogin.core.hooks.PasswordGenerator; import com.github.games647.fastlogin.core.storage.MySQLStorage; import com.github.games647.fastlogin.core.storage.SQLStorage; import com.github.games647.fastlogin.core.storage.SQLiteStorage; +import com.google.common.base.Ticker; import com.google.common.net.HostAndPort; import com.zaxxer.hikari.HikariConfig; @@ -122,7 +123,7 @@ public class FastLoginCore

> { expireTime = MAX_EXPIRE_RATE; } - rateLimiter = new RateLimiter(maxCon, expireTime); + rateLimiter = new RateLimiter(Ticker.systemTicker(), maxCon, expireTime); Set proxies = config.getStringList("proxies") .stream() .map(HostAndPort::fromString) diff --git a/core/src/test/java/com/github/games647/fastlogin/core/FakeTicker.java b/core/src/test/java/com/github/games647/fastlogin/core/FakeTicker.java new file mode 100644 index 00000000..7659e1f8 --- /dev/null +++ b/core/src/test/java/com/github/games647/fastlogin/core/FakeTicker.java @@ -0,0 +1,23 @@ +package com.github.games647.fastlogin.core; + +import com.google.common.base.Ticker; + +import java.time.Duration; + +public class FakeTicker extends Ticker { + + private long timestamp; + + public FakeTicker(long initial) { + timestamp = initial; + } + + @Override + public long read() { + return timestamp; + } + + public void add(Duration duration) { + timestamp += duration.toNanos(); + } +} diff --git a/core/src/test/java/com/github/games647/fastlogin/core/RateLimiterTest.java b/core/src/test/java/com/github/games647/fastlogin/core/RateLimiterTest.java index c93e3721..fd90a7eb 100644 --- a/core/src/test/java/com/github/games647/fastlogin/core/RateLimiterTest.java +++ b/core/src/test/java/com/github/games647/fastlogin/core/RateLimiterTest.java @@ -25,6 +25,7 @@ */ package com.github.games647.fastlogin.core; +import java.time.Duration; import java.util.concurrent.TimeUnit; import org.junit.Test; @@ -43,14 +44,34 @@ public class RateLimiterTest { public void allowExpire() throws InterruptedException { int size = 3; + FakeTicker ticker = new FakeTicker(5_000_000L); + // run twice the size to fill it first and then test it - RateLimiter rateLimiter = new RateLimiter(size, 0); + RateLimiter rateLimiter = new RateLimiter(ticker, size, 0); for (int i = 0; i < size; i++) { assertTrue("Filling up", rateLimiter.tryAcquire()); } for (int i = 0; i < size; i++) { - Thread.sleep(1); + ticker.add(Duration.ofSeconds(1)); + assertTrue("Should be expired", rateLimiter.tryAcquire()); + } + } + + @Test + public void allowExpireNegative() throws InterruptedException { + int size = 3; + + FakeTicker ticker = new FakeTicker(-5_000_000L); + + // run twice the size to fill it first and then test it + RateLimiter rateLimiter = new RateLimiter(ticker, size, 0); + for (int i = 0; i < size; i++) { + assertTrue("Filling up", rateLimiter.tryAcquire()); + } + + for (int i = 0; i < size; i++) { + ticker.add(Duration.ofSeconds(1)); assertTrue("Should be expired", rateLimiter.tryAcquire()); } } @@ -62,8 +83,28 @@ public class RateLimiterTest { public void shouldBlock() { int size = 3; + FakeTicker ticker = new FakeTicker(5_000_000L); + // fill the size - RateLimiter rateLimiter = new RateLimiter(size, TimeUnit.SECONDS.toMillis(30)); + RateLimiter rateLimiter = new RateLimiter(ticker, size, TimeUnit.SECONDS.toMillis(30)); + for (int i = 0; i < size; i++) { + assertTrue("Filling up", rateLimiter.tryAcquire()); + } + + assertFalse("Should be full and no entry should be expired", rateLimiter.tryAcquire()); + } + + /** + * Too many requests + */ + @Test + public void shouldBlockNegative() { + int size = 3; + + FakeTicker ticker = new FakeTicker(-5_000_000L); + + // fill the size + RateLimiter rateLimiter = new RateLimiter(ticker, size, TimeUnit.SECONDS.toMillis(30)); for (int i = 0; i < size; i++) { assertTrue("Filling up", rateLimiter.tryAcquire()); } @@ -76,17 +117,40 @@ public class RateLimiterTest { */ @Test public void blockedNotAdded() throws InterruptedException { + FakeTicker ticker = new FakeTicker(5_000_000L); + // fill the size - 100ms should be reasonable high - RateLimiter rateLimiter = new RateLimiter(1, 100); + RateLimiter rateLimiter = new RateLimiter(ticker, 1, 100); assertTrue("Filling up", rateLimiter.tryAcquire()); - Thread.sleep(50); + ticker.add(Duration.ofMillis(50)); // still is full - should fail assertFalse("Expired too early", rateLimiter.tryAcquire()); // wait the remaining time and add a threshold, because - Thread.sleep(50 + THRESHOLD_MILLI); + ticker.add(Duration.ofMillis(50)); + assertTrue("Request not released", rateLimiter.tryAcquire()); + } + + /** + * Blocked attempts shouldn't replace existing ones. + */ + @Test + public void blockedNotAddedNegative() throws InterruptedException { + FakeTicker ticker = new FakeTicker(-5_000_000L); + + // fill the size - 100ms should be reasonable high + RateLimiter rateLimiter = new RateLimiter(ticker, 1, 100); + assertTrue("Filling up", rateLimiter.tryAcquire()); + + ticker.add(Duration.ofMillis(50)); + + // still is full - should fail + assertFalse("Expired too early", rateLimiter.tryAcquire()); + + // wait the remaining time and add a threshold, because + ticker.add(Duration.ofMillis(50)); assertTrue("Request not released", rateLimiter.tryAcquire()); } }