Skip to content

Commit

Permalink
feat: early reject large request
Browse files Browse the repository at this point in the history
  • Loading branch information
yzfeng2020 committed Dec 20, 2024
1 parent d3303e2 commit 990f4cb
Show file tree
Hide file tree
Showing 9 changed files with 160 additions and 75 deletions.
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2016 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand Down Expand Up @@ -50,6 +50,7 @@ public static ContentTooLargeExceptionBuilder builder() {
private final long maxContentLength;
private final long contentLength;
private final long transferred;
private final boolean earlyRejection;

private ContentTooLargeException() {
this(false);
Expand All @@ -62,16 +63,18 @@ private ContentTooLargeException(boolean neverSample) {
maxContentLength = -1;
transferred = -1;
contentLength = -1;
earlyRejection = false;
}

ContentTooLargeException(long maxContentLength, long contentLength, long transferred,
@Nullable Throwable cause) {
super(toString(maxContentLength, contentLength, transferred), cause);
boolean earlyRejection, @Nullable Throwable cause) {
super(toString(maxContentLength, contentLength, transferred, earlyRejection), cause);

neverSample = false;
this.transferred = transferred;
this.contentLength = contentLength;
this.maxContentLength = maxContentLength;
this.earlyRejection = earlyRejection;
}

/**
Expand All @@ -96,6 +99,13 @@ public long maxContentLength() {
return maxContentLength;
}

/**
* Returns whether the exception is raised when reading content-length header.
*/
public boolean earlyRejection() {
return earlyRejection;
}

@Override
public Throwable fillInStackTrace() {
if (!neverSample && Flags.verboseExceptionSampler().isSampled(getClass())) {
Expand All @@ -105,7 +115,8 @@ public Throwable fillInStackTrace() {
}

@Nullable
private static String toString(long maxContentLength, long contentLength, long transferred) {
private static String toString(long maxContentLength, long contentLength, long transferred,
boolean earlyRejection) {
try (TemporaryThreadLocals ttl = TemporaryThreadLocals.acquire()) {
final StringBuilder buf = ttl.stringBuilder();
if (maxContentLength >= 0) {
Expand All @@ -117,6 +128,9 @@ private static String toString(long maxContentLength, long contentLength, long t
if (transferred >= 0) {
buf.append(", transferred: ").append(transferred);
}
if (earlyRejection) {
buf.append(", earlyRejection: ").append("true");
}
return buf.length() != 0 ? buf.substring(2) : null;
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2021 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand Down Expand Up @@ -30,6 +30,8 @@ public final class ContentTooLargeExceptionBuilder {
private long maxContentLength = -1;
private long contentLength = -1;
private long transferred = -1;
private boolean earlyRejection;

@Nullable
private Throwable cause;

Expand Down Expand Up @@ -83,13 +85,24 @@ public ContentTooLargeExceptionBuilder cause(Throwable cause) {
return this;
}

/**
* Sets the exception as early rejection.
*/
@UnstableApi
public ContentTooLargeExceptionBuilder earlyRejection(boolean isEarlyRejection) {
this.earlyRejection = isEarlyRejection;
return this;
}

/**
* Returns a new instance of {@link ContentTooLargeException}.
*/
public ContentTooLargeException build() {
if (maxContentLength < 0 && contentLength < 0 && transferred < 0 && cause == null) {
if (maxContentLength < 0 && contentLength < 0 &&
transferred < 0 && !earlyRejection && cause == null) {
return ContentTooLargeException.get();
}
return new ContentTooLargeException(maxContentLength, contentLength, transferred, cause);
return new ContentTooLargeException(maxContentLength, contentLength,
transferred, earlyRejection, cause);
}
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2021 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand Down Expand Up @@ -143,6 +143,16 @@ default CompletableFuture<Void> whenAggregated() {
*/
long requestStartTimeMicros();

/**
* Returns the maximum allowed length of the content of the request.
*/
long maxRequestLength();

/**
* Returns the transferred bytes of the request.
*/
long transferredBytes();

/**
* Returns whether the request is an HTTP/1.1 webSocket request.
*/
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2021 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand All @@ -19,10 +19,5 @@
import com.linecorp.armeria.common.HttpRequestWriter;

interface DecodedHttpRequestWriter extends DecodedHttpRequest, HttpRequestWriter {

long maxRequestLength();

long transferredBytes();

void increaseTransferredBytes(long delta);
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2021 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand Down Expand Up @@ -260,4 +260,14 @@ public long requestStartTimeNanos() {
public long requestStartTimeMicros() {
return requestStartTimeMicros;
}

@Override
public long maxRequestLength() {
return 0;
}

@Override
public long transferredBytes() {
return 0;
}
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2016 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand Down Expand Up @@ -208,8 +208,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
// Validate the 'content-length' header.
final String contentLengthStr = headers.get(HttpHeaderNames.CONTENT_LENGTH);
final boolean contentEmpty;
long contentLength = 0;
if (contentLengthStr != null) {
long contentLength;
try {
contentLength = Long.parseLong(contentLengthStr);
} catch (NumberFormatException ignored) {
Expand Down Expand Up @@ -280,6 +280,10 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
final boolean endOfStream = contentEmpty && !transferEncodingChunked;
this.req = req = DecodedHttpRequest.of(endOfStream, eventLoop, id, 1, headers,
keepAlive, inboundTrafficController, routingCtx);
final long maxRequestLength = req.maxRequestLength();
if (maxRequestLength > 0 && contentLength > maxRequestLength) {
abortLargeRequest(ctx, req, id, endOfStream, keepAliveHandler, true);
}
cfg.serverMetrics().increasePendingHttp1Requests();
ctx.fireChannelRead(req);
} else {
Expand Down Expand Up @@ -313,33 +317,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
final long maxContentLength = decodedReq.maxRequestLength();
final long transferredLength = decodedReq.transferredBytes();
if (maxContentLength > 0 && transferredLength > maxContentLength) {
final ContentTooLargeException cause =
ContentTooLargeException.builder()
.maxContentLength(maxContentLength)
.contentLength(req.headers())
.transferred(transferredLength)
.build();
discarding = true;
req = null;
final boolean shouldReset;
if (encoder instanceof ServerHttp1ObjectEncoder) {
if (encoder.isResponseHeadersSent(id, 1)) {
ctx.channel().close();
} else {
keepAliveHandler.disconnectWhenFinished();
}
shouldReset = false;
} else {
// Upgraded to HTTP/2. Reset only if the remote peer is still open.
shouldReset = !endOfStream;
}

// Wrap the cause with the returned status to let LoggingService correctly log the
// status.
final HttpStatusException httpStatusException =
HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
decodedReq.abortResponse(httpStatusException, true);
abortLargeRequest(ctx, decodedReq, id, endOfStream, keepAliveHandler, false);
return;
}

Expand Down Expand Up @@ -377,6 +355,38 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
}
}

private void abortLargeRequest(ChannelHandlerContext ctx, DecodedHttpRequest decodedReq, int id,
boolean endOfStream, KeepAliveHandler keepAliveHandler, boolean isEarlyRejection) {
final ContentTooLargeException cause =
ContentTooLargeException.builder()
.maxContentLength(decodedReq.maxRequestLength())
.transferred(decodedReq.transferredBytes())
.contentLength(decodedReq.headers())
.earlyRejection(isEarlyRejection)
.build();
discarding = true;
req = null;
final boolean shouldReset;
if (encoder instanceof ServerHttp1ObjectEncoder) {
if (encoder.isResponseHeadersSent(id, 1)) {
ctx.channel().close();
} else {
keepAliveHandler.disconnectWhenFinished();
}
shouldReset = false;
} else {
// Upgraded to HTTP/2. Reset only if the remote peer is still open.
shouldReset = !endOfStream;
}

// Wrap the cause with the returned status to let LoggingService correctly log the
// status.
final HttpStatusException httpStatusException =
HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
decodedReq.abortResponse(httpStatusException, true);
}

private void removeFromPipelineIfUpgraded(ChannelHandlerContext ctx, boolean endOfStream) {
if (endOfStream && encoder instanceof ServerHttp2ObjectEncoder) {
// An HTTP/1 connection has been upgraded to HTTP/2.
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2016 LINE Corporation
* Copyright 2024 LINE Corporation
*
* LINE Corporation licenses this file to you under the Apache License,
* version 2.0 (the "License"); you may not use this file except in compliance
Expand Down Expand Up @@ -165,8 +165,8 @@ public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers

// Validate the 'content-length' header if exists.
final String contentLengthStr = headers.get(HttpHeaderNames.CONTENT_LENGTH);
long contentLength = 0;
if (contentLengthStr != null) {
long contentLength;
try {
contentLength = Long.parseLong(contentLengthStr);
} catch (NumberFormatException ignored) {
Expand Down Expand Up @@ -206,6 +206,10 @@ public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers
final EventLoop eventLoop = ctx.channel().eventLoop();
req = DecodedHttpRequest.of(endOfStream, eventLoop, id, streamId, headers, true,
inboundTrafficController, routingCtx);
final long maxRequestLength = req.maxRequestLength();
if (maxRequestLength > 0 && contentLength > maxRequestLength) {
abortLargeRequest(req, endOfStream, true);
}
requests.put(streamId, req);
cfg.serverMetrics().increasePendingHttp2Requests();
ctx.fireChannelRead(req);
Expand Down Expand Up @@ -321,20 +325,7 @@ public int onDataRead(
final long maxContentLength = decodedReq.maxRequestLength();
final long transferredLength = decodedReq.transferredBytes();
if (maxContentLength > 0 && transferredLength > maxContentLength) {
assert encoder != null;
final ContentTooLargeException cause =
ContentTooLargeException.builder()
.maxContentLength(maxContentLength)
.contentLength(decodedReq.headers())
.transferred(transferredLength)
.build();

final boolean shouldReset = !endOfStream;

final HttpStatusException httpStatusException =
HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
decodedReq.abortResponse(httpStatusException, true);
abortLargeRequest(decodedReq, endOfStream, false);
} else if (decodedReq.isOpen()) {
try {
// The decodedReq will be automatically closed if endOfStream is true.
Expand All @@ -350,6 +341,25 @@ public int onDataRead(
return dataLength + padding;
}

private void abortLargeRequest(DecodedHttpRequest decodedReq, boolean endOfStream,
boolean isEarlyRejection) {
assert encoder != null;
final ContentTooLargeException cause =
ContentTooLargeException.builder()
.maxContentLength(decodedReq.maxRequestLength())
.contentLength(decodedReq.headers())
.transferred(decodedReq.transferredBytes())
.earlyRejection(isEarlyRejection)
.build();

final boolean shouldReset = !endOfStream;

final HttpStatusException httpStatusException =
HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
decodedReq.abortResponse(httpStatusException, true);
}

private void writeInvalidRequestPathResponse(int streamId, @Nullable RequestHeaders headers) {
writeErrorResponse(streamId, headers, HttpStatus.BAD_REQUEST,
"Invalid request path", null);
Expand Down
Loading

0 comments on commit 990f4cb

Please sign in to comment.