Skip to content

Commit

Permalink
Merge pull request #86 from GoogleCloudPlatform:fixSizeLimitHandler
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 591117937
Change-Id: I903eb5ac1f2d28ac4c093331d367d94e8907dad5
  • Loading branch information
@gae-java-bot
gae-java-bot committed Dec 15, 2023
2 parents 6765b66 + 37e4578 commit f410df1
Showing 1 changed file with 72 additions and 62 deletions.
134 changes: 72 additions & 62 deletions ..._impl_jetty12/src/main/java/com/google/apphosting/runtime/jetty/CoreSizeLimitHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,21 +16,20 @@

package com.google.apphosting.runtime.jetty;

import java.nio.ByteBuffer;
import org.eclipse.jetty.http.BadMessageException;
import org.eclipse.jetty.http.HttpException;
import org.eclipse.jetty.http.HttpField;
import org.eclipse.jetty.http.HttpFields;
import org.eclipse.jetty.http.HttpHeader;
import org.eclipse.jetty.http.MetaData;
import org.eclipse.jetty.http.HttpStatus;
import org.eclipse.jetty.io.Content;
import org.eclipse.jetty.server.Handler;
import org.eclipse.jetty.server.HttpStream;
import org.eclipse.jetty.server.Request;
import org.eclipse.jetty.server.Response;
import org.eclipse.jetty.util.BufferUtil;
import org.eclipse.jetty.util.Callback;

import java.nio.ByteBuffer;

/**
* A handler that can limit the size of message bodies in requests and responses.
*
Expand All @@ -46,8 +45,6 @@ public class CoreSizeLimitHandler extends Handler.Wrapper
{
private final long _requestLimit;
private final long _responseLimit;
private long _read = 0;
private long _written = 0;

/**
* @param requestLimit The request body size limit in bytes or -1 for no limit
Expand All @@ -66,85 +63,98 @@ public boolean handle(Request request, Response response, Callback callback) thr
if (contentLengthField != null)
{
long contentLength = contentLengthField.getLongValue();
if (_requestLimit > 0 && contentLength > _requestLimit) {
response.setStatus(413);
if (_requestLimit >= 0 && contentLength > _requestLimit)
{
String s = "Request body is too large: " + contentLength + ">" + _requestLimit;
response.write(true, BufferUtil.toBuffer(s), callback);
Response.writeError(request, response, callback, HttpStatus.PAYLOAD_TOO_LARGE_413, s);
return true;
}
}

HttpFields.Mutable.Wrapper httpFields = new HttpFields.Mutable.Wrapper(response.getHeaders())
SizeLimitRequestWrapper wrappedRequest = new SizeLimitRequestWrapper(request);
SizeLimitResponseWrapper wrappedResponse = new SizeLimitResponseWrapper(wrappedRequest, response);
return super.handle(wrappedRequest, wrappedResponse, callback);
}

private class SizeLimitRequestWrapper extends Request.Wrapper
{
private long _read = 0;

public SizeLimitRequestWrapper(Request wrapped)
{
@Override
public HttpField onAddField(HttpField field)
super(wrapped);
}

@Override
public Content.Chunk read()
{
Content.Chunk chunk = super.read();
if (chunk == null)
return null;
if (chunk.getFailure() != null)
return chunk;

// Check request content limit.
ByteBuffer content = chunk.getByteBuffer();
if (content != null && content.remaining() > 0)
{
if ((field.getHeader()!=null) && (field.getHeader().is(HttpHeader.CONTENT_LENGTH.asString())))
_read += content.remaining();
if (_requestLimit >= 0 && _read > _requestLimit)
{
long contentLength = field.getLongValue();
if (_responseLimit > 0 && contentLength > _responseLimit)
throw new HttpException.RuntimeException(500, "Response body is too large: " + contentLength + ">" + _responseLimit);
BadMessageException e = new BadMessageException(HttpStatus.PAYLOAD_TOO_LARGE_413, "Request body is too large: " + _read + ">" + _requestLimit);
getWrapped().fail(e);
return null;
}
return super.onAddField(field);
}
};

response = new Response.Wrapper(request, response)
{
@Override
public HttpFields.Mutable getHeaders()
{
return httpFields;
}
};
return chunk;
}
}

request.addHttpStreamWrapper(httpStream -> new HttpStream.Wrapper(httpStream)
{
@Override
public Content.Chunk read()
{
Content.Chunk chunk = super.read();
if (chunk == null)
return null;
if (chunk.getFailure() != null)
return chunk;
private class SizeLimitResponseWrapper extends Response.Wrapper
{
private final HttpFields.Mutable _httpFields;
private long _written = 0;

public SizeLimitResponseWrapper(Request request, Response wrapped) {
super(request, wrapped);

// Check request content limit.
ByteBuffer content = chunk.getByteBuffer();
if (content != null && content.remaining() > 0)
_httpFields = new HttpFields.Mutable.Wrapper(wrapped.getHeaders())
{
@Override
public HttpField onAddField(HttpField field)
{
_read += content.remaining();
if (_requestLimit >= 0 && _read > _requestLimit)
if (field.getHeader().is(HttpHeader.CONTENT_LENGTH.asString()))
{
System.err.println("we actually read too much content in coreSizeLimitHandler for request size limit");
BadMessageException e = new BadMessageException(413, "Request body is too large: " + _read + ">" + _requestLimit);
request.fail(e);
return null;
long contentLength = field.getLongValue();
if (_responseLimit >= 0 && contentLength > _responseLimit)
throw new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, "Response body is too large: " + contentLength + ">" + _responseLimit);
}
return super.onAddField(field);
}
};
}

return chunk;
}
@Override
public HttpFields.Mutable getHeaders() {
return _httpFields;
}

@Override
public void send(MetaData.Request request, MetaData.Response response, boolean last, ByteBuffer content, Callback callback)
@Override
public void write(boolean last, ByteBuffer content, Callback callback)
{
if (content != null && content.remaining() > 0)
{
// Check response content limit.
if (content != null && content.remaining() > 0)
if (_responseLimit >= 0 && (_written + content.remaining()) > _responseLimit)
{
if (_responseLimit >= 0 && (_written + content.remaining()) > _responseLimit)
{
callback.failed(new HttpException.RuntimeException(500, "Response body is too large: " +
_written + content.remaining() + ">" + _responseLimit));
return;
}
_written += content.remaining();
callback.failed(new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, "Response body is too large: " +
_written + content.remaining() + ">" + _responseLimit));
return;
}

super.send(request, response, last, content, callback);
_written += content.remaining();
}
});

return super.handle(request, response, callback);
super.write(last, content, callback);
}
}
}

0 comments on commit f410df1

Please sign in to comment.