ThrottlingMiddleware extends AbstractWebMiddleware<HttpRequest, NRES> { @Override public HttpResponse handle(HttpRequest request, MiddlewareChain<HttpRequest, NRES, ?, ?> chain) { if (throttles.stream() .anyMatch(throttle -> throttle.apply(request))) { return builder(HttpResponseUtils.response("Too Many Request")) .set(HttpResponse::setStatus, 429) .build(); } return castToHttpResponse(chain.next(request)); }  @Override HttpResponse handle(HttpRequest request, MiddlewareChain<HttpRequest, NRES, ?, ?> chain); void setThrottles(List<Throttle> throttles);  }
@Test public void sameIP() throws InterruptedException { List<Throttle> throttles = Collections.singletonList( new Throttle("IP", new LimitRate(1, Duration.ofMillis(500L)), HttpRequest::getRemoteAddr) ); ThrottlingMiddleware<HttpResponse> middleware = builder(new ThrottlingMiddleware<HttpResponse>()) .set(ThrottlingMiddleware::setThrottles, throttles) .build(); DefaultHttpRequest req = new DefaultHttpRequest(); req.setRemoteAddr("127.0.0.1"); MiddlewareChain<HttpRequest, HttpResponse, ?, ?> chain = new DefaultMiddlewareChain<>(new AnyPredicate<>(), null, (Endpoint<HttpRequest, HttpResponse>) request -> HttpResponse.of("")); final AtomicInteger count200 = new AtomicInteger(0); ScheduledExecutorService service = Executors.newScheduledThreadPool(3); service.scheduleAtFixedRate(() -> { HttpResponse res = middleware.handle(req, chain); if (res.getStatus() == 200) { count200.addAndGet(1); } }, 0, 250, TimeUnit.MILLISECONDS); service.schedule(service::shutdown, 4, TimeUnit.SECONDS); service.awaitTermination(5, TimeUnit.SECONDS); assertThat(count200.get()).isGreaterThan(3); }
@Test public void randomIP() throws InterruptedException { List<Throttle> throttles = Collections.singletonList( new Throttle("IP", new LimitRate(1, Duration.ofMillis(500L)), HttpRequest::getRemoteAddr) ); ThrottlingMiddleware<HttpResponse> middleware = builder(new ThrottlingMiddleware<HttpResponse>()) .set(ThrottlingMiddleware::setThrottles, throttles) .build(); DefaultHttpRequest req = new DefaultHttpRequest(); MiddlewareChain<HttpRequest, HttpResponse, ?, ?> chain = new DefaultMiddlewareChain<>(new AnyPredicate<>(), null, (Endpoint<HttpRequest, HttpResponse>) request -> HttpResponse.of("")); final AtomicInteger count429 = new AtomicInteger(0); ScheduledExecutorService service = Executors.newScheduledThreadPool(3); service.scheduleAtFixedRate(() -> { req.setRemoteAddr(RandomUtils.nextInt(1, 255) + "." + RandomUtils.nextInt(1, 255) + "." + RandomUtils.nextInt(1, 255) + "." + RandomUtils.nextInt(1, 255) ); HttpResponse res = middleware.handle(req, chain); if (res.getStatus() == 429) { count429.addAndGet(1); } }, 0, 250, TimeUnit.MILLISECONDS); service.schedule(service::shutdown, 3, TimeUnit.SECONDS); service.awaitTermination(5, TimeUnit.SECONDS); assertThat(count429.get()).isEqualTo(0); }