@ParameterizedTest
@ArgumentsSource(ClientProvider.class)
void testTooLargeContent(WebClient client) throws Exception {
    final int maxContentLength = 65536;
    serverMaxRequestLength = maxContentLength;

    final HttpRequestWriter req = HttpRequest.streaming(HttpMethod.POST, "/count");
    final CompletableFuture<AggregatedHttpResponse> f = client.execute(req).aggregate();

    stream(req, maxContentLength + 1, 1024);

    final AggregatedHttpResponse res = f.get();

    assertThat(res.status()).isEqualTo(HttpStatus.REQUEST_ENTITY_TOO_LARGE);
    assertThat(res.contentType()).isEqualTo(MediaType.PLAIN_TEXT_UTF_8);
    assertThat(res.contentUtf8()).isEqualTo("413 Request Entity Too Large");
}
 

@ParameterizedTest
@ArgumentsSource(ClientAndProtocolProvider.class)
void testStreamRequestLongerThanTimeout(WebClient client) throws Exception {
    // Disable timeouts and length limits so that test does not fail due to slow transfer.
    clientWriteTimeoutMillis = 0;
    clientResponseTimeoutMillis = 0;
    clientMaxResponseLength = 0;
    serverRequestTimeoutMillis = 0;

    final HttpRequestWriter request = HttpRequest.streaming(HttpMethod.POST, "/echo");
    final HttpResponse response = client.execute(request);
    request.write(HttpData.ofUtf8("a"));
    Thread.sleep(2000);
    request.write(HttpData.ofUtf8("b"));
    Thread.sleep(2000);
    request.write(HttpData.ofUtf8("c"));
    Thread.sleep(2000);
    request.write(HttpData.ofUtf8("d"));
    Thread.sleep(2000);
    request.write(HttpData.ofUtf8("e"));
    request.close();
    assertThat(response.aggregate().get().contentUtf8()).isEqualTo("abcde");
}
 

@Override
public HttpResponse authFailed(
    HttpService delegate, ServiceRequestContext ctx, HttpRequest req, @Nullable Throwable cause)
    throws Exception {
  if (cause != null) {
    logger.warn("Unexpected exception during authorization.", cause);
    return HttpResponse.of(HttpStatus.UNAUTHORIZED);
  }
  if (!config.getIncludedPaths().isEmpty()) {
    if (config.getIncludedPaths().contains(ctx.path())) {
      return HttpResponse.of(HttpStatus.UNAUTHORIZED);
    } else {
      return delegate.serve(ctx, req);
    }
  }
  if (config.getExcludedPaths().contains(ctx.path())) {
    return delegate.serve(ctx, req);
  }
  return HttpResponse.of(HttpStatus.UNAUTHORIZED);
}
 

@Override
public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
    return HttpResponse.from(
            first.findFile(ctx, req)
                 .readAttributes(ctx.blockingTaskExecutor())
                 .thenApply(firstAttrs -> {
                     try {
                         if (firstAttrs != null) {
                             return first.serve(ctx, req);
                         }

                         return second.serve(ctx, req);
                     } catch (Exception e) {
                         return Exceptions.throwUnsafely(e);
                     }
                 }));
}
 

/**
 * Serves the specified {@link HttpRequest} by delegating it to the matching {@code 'doMETHOD()'} method.
 */
@Override
public final HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
    switch (req.method()) {
        case OPTIONS:
            return doOptions(ctx, req);
        case GET:
            return doGet(ctx, req);
        case HEAD:
            return doHead(ctx, req);
        case POST:
            return doPost(ctx, req);
        case PUT:
            return doPut(ctx, req);
        case PATCH:
            return doPatch(ctx, req);
        case DELETE:
            return doDelete(ctx, req);
        case TRACE:
            return doTrace(ctx, req);
        default:
            return HttpResponse.of(HttpStatus.METHOD_NOT_ALLOWED);
    }
}
 

@Test
void setResponseTimeoutAfter() throws InterruptedException {
    final HttpRequest req = HttpRequest.of(HttpMethod.GET, "/");
    final DefaultClientRequestContext ctx = (DefaultClientRequestContext) ClientRequestContext.of(req);
    final long tolerance = 500;

    ctx.eventLoop().execute(() -> {
        ctx.setResponseTimeoutMillis(TimeoutMode.SET_FROM_NOW, 1000);
        final long oldResponseTimeoutMillis = ctx.responseTimeoutMillis();
        ctx.setResponseTimeoutMillis(TimeoutMode.SET_FROM_NOW, 2000);
        assertThat(ctx.responseTimeoutMillis()).isBetween(oldResponseTimeoutMillis + 1000 - tolerance,
                                                          oldResponseTimeoutMillis + 1000 + tolerance);
        finished.set(true);
    });

    await().untilTrue(finished);
}
 

@Override
protected void configure(ServerBuilder sb) throws Exception {
    sb.service("/hello", new AbstractHttpService() {
        final AtomicInteger reqCount = new AtomicInteger();

        @Override
        protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req) {
            ctx.mutateAdditionalResponseTrailers(
                    mutator -> mutator.add(HttpHeaderNames.of("foo"), "bar"));
            if (reqCount.getAndIncrement() < 1) {
                return HttpResponse.of(HttpStatus.INTERNAL_SERVER_ERROR);
            } else {
                return HttpResponse.of(HttpStatus.OK, MediaType.PLAIN_TEXT_UTF_8, "hello");
            }
        }
    });
}
 

private static ClientRequestContext newClientContext(
        String path, @Nullable String query) throws Exception {

    final InetSocketAddress remoteAddress = new InetSocketAddress(
            InetAddress.getByAddress("server.com", new byte[] { 1, 2, 3, 4 }), 8080);
    final InetSocketAddress localAddress = new InetSocketAddress(
            InetAddress.getByAddress("client.com", new byte[] { 5, 6, 7, 8 }), 5678);

    final String pathAndQuery = path + (query != null ? '?' + query : "");
    final HttpRequest req = HttpRequest.of(RequestHeaders.of(HttpMethod.GET, pathAndQuery,
                                                             HttpHeaderNames.AUTHORITY, "server.com:8080",
                                                             HttpHeaderNames.USER_AGENT, "some-client"));

    final ClientRequestContext ctx =
            ClientRequestContext.builder(req)
                                .remoteAddress(remoteAddress)
                                .localAddress(localAddress)
                                .endpoint(Endpoint.of("server.com", 8080))
                                .sslSession(newSslSession())
                                .build();

    ctx.setAttr(MY_ATTR, new CustomObject("some-name", "some-value"));
    return ctx;
}
 

@Test
void doesNotContain() {
    assertThat(client.execute(HttpRequest.of(HttpMethod.GET, "/doesNotContain?your-param=your-value"))
                     .aggregate().join().contentUtf8()).isEqualTo("!my-param");
    assertThat(client.execute(HttpRequest.of(HttpMethod.GET, "/doesNotContain?my-param=my-value"))
                     .aggregate().join().contentUtf8()).isEqualTo("fallback");
}
 

@Test
void returnToThePoolAfterRequestIsComplete() {
    final CompleteInterceptableFuture<Void> future = new CompleteInterceptableFuture<>();
    final HttpRequest httpRequest = httpRequest(future);

    final WebClient webClient = WebClient.of(server.uri(SessionProtocol.H1C));
    final AggregatedHttpResponse res = webClient.execute(httpRequest).aggregate().join();
    assertThat(res.status()).isSameAs(HttpStatus.OK);
    assertThat(webClient.get("/").aggregate().join().status()).isSameAs(HttpStatus.OK);
    future.completeValue(null); // This will make the first connection return to the pool.
    assertThat(webClient.get("/").aggregate().join().status()).isSameAs(HttpStatus.OK);
    assertThat(remoteAddresses.get(0)).isNotSameAs(remoteAddresses.get(1));
    assertThat(remoteAddresses.get(0)).isSameAs(remoteAddresses.get(2));
}
 

@Test
void path() {
    final HttpRequest req = HttpRequest.of(HttpMethod.GET, "/foo");
    final HttpServerRequest braveReq = ServiceRequestContextAdapter.asHttpServerRequest(
            ServiceRequestContext.of(req));

    assertThat(braveReq.path()).isEqualTo("/foo");
}
 

private Function<? super HttpClient, ? extends HttpClient> loggingDecorator() {
    return delegate -> new SimpleDecoratingHttpClient(delegate) {
        @Override
        public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Exception {
            ctx.log().whenRequestComplete().thenAccept(log -> listener.accept(log.partial()));
            ctx.log().whenComplete().thenAccept(listener);
            return unwrap().execute(ctx, req);
        }
    };
}
 

@Override
public Function<? super HttpService, ? extends HttpService> newDecorator(MyDecorator3 parameter) {
    return delegate -> new SimpleDecoratingHttpService(delegate) {
        @Override
        public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
            appendAttribute(ctx, " (decorated-3)");
            return unwrap().serve(ctx, req);
        }
    };
}
 

@Override
protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req) throws Exception {
    final String value1 = ctx.pathParam("var1");
    final String value2 = ctx.pathParam("var2");
    final String value3 = ctx.pathParam("var3");
    return HttpResponse.of(HttpStatus.OK, MediaType.PLAIN_TEXT_UTF_8,
                           "SERVICE_NAMED_PARAM %s %s %s", value1, value2, value3);
}
 

@Override
public HttpResponse convertResponse(ServiceRequestContext ctx, ResponseHeaders headers,
                                    @Nullable Object resObj,
                                    HttpHeaders trailingHeaders) throws Exception {
    try {
        final HttpRequest request = RequestContext.current().request();
        if (resObj == null || HttpMethod.DELETE == request.method() ||
            (resObj instanceof Iterable && Iterables.size((Iterable<?>) resObj) == 0)) {
            return HttpResponse.of(HttpStatus.NO_CONTENT);
        }

        final ResponseHeaders resHeaders;
        if (headers.contentType() == null) {
            final ResponseHeadersBuilder builder = headers.toBuilder();
            builder.contentType(MediaType.JSON_UTF_8);
            resHeaders = builder.build();
        } else {
            resHeaders = headers;
        }

        final HttpData httpData = HttpData.wrap(Jackson.writeValueAsBytes(resObj));
        return HttpResponse.of(resHeaders, httpData, trailingHeaders);
    } catch (JsonProcessingException e) {
        logger.debug("Failed to convert a response:", e);
        return HttpApiUtil.newResponse(ctx, HttpStatus.INTERNAL_SERVER_ERROR, e);
    }
}
 

@Override
protected void configure(ServerBuilder sb) throws Exception {
    // Bind a service that returns the remote address of the connection to determine
    // if the same connection was used to handle more than one request.
    sb.service("/", new AbstractHttpService() {
        @Override
        protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req) throws Exception {
            // Consume the request completely so that the connection can be returned to the pool.
            return HttpResponse.from(req.aggregate().handle((unused1, unused2) -> {
                // Signal the main thread that the connection has been returned to the pool.
                // Note that this is true only when pipelining is enabled. The connection is returned
                // after response is fully sent if pipelining is disabled.
                lock.lock();
                try {
                    connectionReturnedToPool = true;
                    condition.signal();
                } finally {
                    lock.unlock();
                }

                semaphore.acquireUninterruptibly();
                try {
                    return HttpResponse.of(HttpStatus.OK, MediaType.PLAIN_TEXT_UTF_8,
                                           String.valueOf(ctx.remoteAddress()));
                } finally {
                    semaphore.release();
                }
            }));
        }
    });
}
 

@BeforeEach
void setUp() {
    when(eventLoop.inEventLoop()).thenReturn(true);
    ctx = ServiceRequestContext.builder(HttpRequest.of(HttpMethod.GET, "/"))
                               .eventLoop(eventLoop)
                               .build();
}
 

@Test
void clearResponseTimeoutWithPendingTask() {
    final HttpRequest req = HttpRequest.of(HttpMethod.GET, "/");
    final DefaultClientRequestContext ctx =
            (DefaultClientRequestContext) ClientRequestContext.builder(req)
                                                              .noTimeoutController()
                                                              .build();

    ctx.clearResponseTimeout();
    final TimeoutController timeoutController = mock(TimeoutController.class);
    ctx.setResponseTimeoutController(timeoutController);

    verify(timeoutController, timeout(1000)).cancelTimeout();
    assertThat(ctx.responseTimeoutMillis()).isEqualTo(0);
}
 

@Test
void sanitizeRequestHeaders() throws Exception {

    final HttpRequest req = HttpRequest.of(RequestHeaders.of(HttpMethod.POST, "/hello/trustin",
                                                             HttpHeaderNames.SCHEME, "http",
                                                             HttpHeaderNames.AUTHORITY, "test.com"));

    final ServiceRequestContext ctx = ServiceRequestContext.of(req);
    ctx.logBuilder().endResponse(new Exception("not sanitized"));
    final Logger logger = LoggingTestUtil.newMockLogger(ctx, capturedCause);
    when(logger.isInfoEnabled()).thenReturn(true);
    when(logger.isWarnEnabled()).thenReturn(true);

    final LoggingService service =
            LoggingService.builder()
                          .logger(logger)
                          .requestLogLevel(LogLevel.INFO)
                          .successfulResponseLogLevel(LogLevel.INFO)
                          .requestHeadersSanitizer(RegexBasedSanitizer.of(
                                  Pattern.compile("trustin"),
                                  Pattern.compile("com")))
                          .newDecorator().apply(delegate);

    assertThat(ctx.logBuilder().toString()).contains("trustin");
    assertThat(ctx.logBuilder().toString()).contains("test.com");
    service.serve(ctx, ctx.request());
    assertThat(ctx.logBuilder().toString()).doesNotContain("trustin");
    assertThat(ctx.logBuilder().toString()).doesNotContain("com");
}
 

@Override
public CompletionStage<Boolean> authorize(ServiceRequestContext ctx, HttpRequest req) {
    final String value = req.headers().get(HttpHeaderNames.COOKIE);
    if (value == null) {
        return CompletableFuture.completedFuture(false);
    }

    // Authentication will be succeeded only if both the specified cookie name and value are matched.
    final Set<Cookie> cookies = Cookie.fromCookieHeader(value);
    final boolean result = cookies.stream().anyMatch(
            cookie -> cookieName.equals(cookie.name()) && cookieValue.equals(cookie.value()));
    return CompletableFuture.completedFuture(result);
}
 

@Test
void mapRequestLogLevelMapper() throws Exception {
    final ServiceRequestContext ctx = ServiceRequestContext.of(HttpRequest.of(RequestHeaders.of(
            HttpMethod.GET, "/", "x-req", "test", "x-res", "test")));
    final Logger logger = LoggingTestUtil.newMockLogger(ctx, capturedCause);
    when(logger.isWarnEnabled()).thenReturn(true);

    final LoggingService service =
            LoggingService.builder()
                          .logger(logger)
                          .requestLogLevelMapper(log -> {
                              if (log.requestHeaders().contains("x-req")) {
                                  return LogLevel.WARN;
                              } else {
                                  return LogLevel.INFO;
                              }
                          })
                          .responseLogLevelMapper(log -> {
                              if (log.requestHeaders().contains("x-res")) {
                                  return LogLevel.WARN;
                              } else {
                                  return LogLevel.INFO;
                              }
                          })
                          .newDecorator().apply(delegate);

    // Check if logs at WARN level if there are headers we're looking for.
    service.serve(ctx, ctx.request());
    verify(logger, never()).isInfoEnabled();
    verify(logger, times(2)).isWarnEnabled();
    verify(logger).warn(eq(REQUEST_FORMAT), same(ctx),
                        matches(".*headers=\\[:method=GET, :path=/, x-req=test, x-res=test].*"));
    verify(logger).warn(eq(RESPONSE_FORMAT), same(ctx), anyString());
}
 

@Test
void requestLogAvailabilityException() {
    final String fullName = AccessLogFormatsTest.class.getSimpleName() + "/rpcMethod";
    final String expectedLogMessage = "\"GET /armeria/log#" + fullName + " h2c\" 200 1024";

    final ServiceRequestContext ctx = ServiceRequestContext.builder(
            HttpRequest.of(RequestHeaders.of(HttpMethod.GET, "/armeria/log",
                                             HttpHeaderNames.USER_AGENT, "armeria/x.y.z",
                                             HttpHeaderNames.REFERER, "http://log.example.com",
                                             HttpHeaderNames.COOKIE, "a=1;b=2"))).build();
    final RequestLog log = ctx.log().partial();
    final RequestLogBuilder logBuilder = ctx.logBuilder();

    // AccessLogger#format will be called after response is finished.
    final AtomicReference<RequestLog> logHolder = new AtomicReference<>();
    log.whenComplete().thenAccept(logHolder::set);

    // RequestLogAvailabilityException will be raised inside AccessLogger#format before injecting each
    // component to RequestLog. So we cannot get the expected log message here.
    assertThat(AccessLogger.format(AccessLogFormats.COMMON, log)).doesNotEndWith(expectedLogMessage);
    logBuilder.requestContent(RpcRequest.of(AccessLogFormatsTest.class, "rpcMethod"), null);
    assertThat(AccessLogger.format(AccessLogFormats.COMMON, log)).doesNotEndWith(expectedLogMessage);
    logBuilder.endRequest();
    assertThat(AccessLogger.format(AccessLogFormats.COMMON, log)).doesNotEndWith(expectedLogMessage);
    logBuilder.responseHeaders(ResponseHeaders.of(HttpStatus.OK));
    assertThat(AccessLogger.format(AccessLogFormats.COMMON, log)).doesNotEndWith(expectedLogMessage);
    logBuilder.responseLength(1024);
    assertThat(AccessLogger.format(AccessLogFormats.COMMON, log)).doesNotEndWith(expectedLogMessage);
    logBuilder.endResponse();

    assertThat(AccessLogger.format(AccessLogFormats.COMMON, logHolder.get()))
            .endsWith(expectedLogMessage);
}
 

@Test
void contains() {
    assertThat(client.execute(HttpRequest.of(HttpMethod.GET, "/contains?my-param=any-value"))
                     .aggregate().join().contentUtf8()).isEqualTo("my-param");
    assertThat(client.execute(HttpRequest.of(HttpMethod.GET, "/contains?your-param=your-value"))
                     .aggregate().join().contentUtf8()).isEqualTo("fallback");
}
 

@Override
public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
    final HttpServerRequest braveReq = ServiceRequestContextAdapter.asHttpServerRequest(ctx);
    final Span span = handler.handleReceive(braveReq);

    // For no-op spans, nothing special to do.
    if (span.isNoop()) {
        try (SpanInScope ignored = tracer.withSpanInScope(span)) {
            return unwrap().serve(ctx, req);
        }
    }

    ctx.log().whenComplete().thenAccept(log -> {
        span.start(log.requestStartTimeMicros());

        final Long wireReceiveTimeNanos = log.requestFirstBytesTransferredTimeNanos();
        assert wireReceiveTimeNanos != null;
        SpanTags.logWireReceive(span, wireReceiveTimeNanos, log);

        final Long wireSendTimeNanos = log.responseFirstBytesTransferredTimeNanos();
        if (wireSendTimeNanos != null) {
            SpanTags.logWireSend(span, wireSendTimeNanos, log);
        } else {
            // If the client timed-out the request, we will have never sent any response data at all.
        }

        final HttpServerResponse braveRes =
                ServiceRequestContextAdapter.asHttpServerResponse(log, braveReq);
        handler.handleSend(braveRes, span);
    });

    try (SpanInScope ignored = tracer.withSpanInScope(span)) {
        return unwrap().serve(ctx, req);
    }
}
 

@Test
void responseAbortWithException() throws InterruptedException {
    final WebClient client = WebClient.of(server.httpUri());
    final HttpRequest request = HttpRequest.streaming(HttpMethod.GET, "/client-aborted");
    final HttpResponse response = client.execute(request);

    await().until(() -> completed);
    final IllegalStateException badState = new IllegalStateException("bad state");
    response.abort(badState);
    assertThatThrownBy(() -> response.aggregate().join())
            .isInstanceOf(CompletionException.class)
            .hasCause(badState);
}
 

@Test
void requestLogWithEmptyCause() {
    final ServiceRequestContext ctx = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));

    final RequestLogBuilder logBuilder = ctx.logBuilder();

    final List<AccessLogComponent> format =
            AccessLogFormats.parseCustom("%{requestCause}L %{responseCause}L");

    logBuilder.endRequest();
    logBuilder.endResponse();

    assertThat(AccessLogger.format(format, ctx.log().ensureComplete())).isEqualTo("- -");
}
 

@Override
public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
  if (ctx.mappedPath().indexOf('.', ctx.mappedPath().lastIndexOf('/') + 1) != -1
      || ctx.mappedPath().charAt(ctx.mappedPath().length() - 1) == '/') {
    // A path that ends with '/' will be handled by HttpFileService correctly, and otherwise if
    // it has a '.' in the last path segment, assume it is a filename.
    return delegate().serve(ctx, req);
  }
  return delegate().serve(new ContextWrapper(ctx), req);
}
 

@Override public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) {
  // We aggregate the request body with pooled objects because signing implies reading it before
  // sending it to Elasticsearch.
  return HttpResponse.from(
      req.aggregateWithPooledObjects(ctx.contextAwareEventLoop(), ctx.alloc())
          .thenApply(aggReg -> {
            try {
              AggregatedHttpRequest signed = sign(ctx, aggReg);
              return delegate().execute(ctx, signed.toHttpRequest());
            } catch (Exception e) {
              return HttpResponse.ofFailure(e);
            }
          }));
}
 

@Override
public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
    if (sampler.isSampled(ctx)) {
        ctx.log().whenRequestComplete().thenAccept(requestLogger);
        ctx.log().whenComplete().thenAccept(responseLogger);
    }
    return unwrap().serve(ctx, req);
}
 

@Override
protected void configure(ServerBuilder sb) throws Exception {
    final TokenBucket tokenBucket = TokenBucket.builder()
                                               .limit(1L, Duration.ofSeconds(10L))
                                               .build();
    sb.service("/http-serve",
               SERVICE.decorate(
                       ThrottlingService.newDecorator(
                               TokenBucketThrottlingStrategy.<HttpRequest>builder(tokenBucket)
                                       .build())));
    sb.service("/http-throttle1",
               SERVICE.decorate(
                       ThrottlingService.newDecorator(
                               TokenBucketThrottlingStrategy.<HttpRequest>builder(tokenBucket)
                                       .withHeadersScheme(ThrottlingHeaders.X_RATELIMIT)
                                       .build())));
    sb.service("/http-throttle2",
               SERVICE.decorate(
                       ThrottlingService.newDecorator(
                               TokenBucketThrottlingStrategy.<HttpRequest>builder(tokenBucket)
                                       .withMinimumBackoff(Duration.ofSeconds(15L))
                                       .withHeadersScheme(ThrottlingHeaders.X_RATELIMIT, true)
                                       .build())));
    sb.service("/http-throttle3",
               SERVICE.decorate(
                       ThrottlingService.newDecorator(
                               TokenBucketThrottlingStrategy.<HttpRequest>builder(tokenBucket)
                                       .build())));
    sb.service("/http-throttle4",
               SERVICE.decorate(
                       ThrottlingService.newDecorator(
                               TokenBucketThrottlingStrategy.<HttpRequest>builder(tokenBucket)
                                       .withMinimumBackoff(Duration.ofSeconds(5L))
                                       .build(),
                               (delegate, ctx, req, cause) ->
                                       HttpResponse.of(HttpStatus.SERVICE_UNAVAILABLE))));
}