package com.linkedin.databus.core.test.netty;
/*
*
* Copyright 2013 LinkedIn Corp. All rights reserved
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*
*/
import java.util.List;
import java.util.Map.Entry;
import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;
import org.jboss.netty.channel.ChannelHandlerContext;
import org.jboss.netty.channel.Channels;
import org.jboss.netty.channel.MessageEvent;
import org.jboss.netty.channel.SimpleChannelUpstreamHandler;
import org.jboss.netty.handler.codec.frame.TooLongFrameException;
import org.jboss.netty.handler.codec.http.HttpChunk;
import org.jboss.netty.handler.codec.http.HttpChunkTrailer;
import org.jboss.netty.handler.codec.http.HttpHeaders;
import org.jboss.netty.handler.codec.http.HttpMessage;
/**
* Based on netty's HttpChunkAggregator but can handle footers
* @author cbotev
*
*/
public class FooterAwareHttpChunkAggregator extends SimpleChannelUpstreamHandler {
private final int maxContentLength;
private HttpMessage currentMessage;
/**
* Creates a new instance.
*
* @param maxContentLength
* the maximum length of the aggregated content.
* If the length of the aggregated content exceeds this value,
* a {@link TooLongFrameException} will be raised.
*/
public FooterAwareHttpChunkAggregator(int maxContentLength) {
if (maxContentLength <= 0) {
throw new IllegalArgumentException(
"maxContentLength must be a positive integer: " +
maxContentLength);
}
this.maxContentLength = maxContentLength;
}
@Override
public void messageReceived(ChannelHandlerContext ctx, MessageEvent e)
throws Exception {
Object msg = e.getMessage();
HttpMessage currentMessage = this.currentMessage;
if (msg instanceof HttpMessage) {
HttpMessage m = (HttpMessage) msg;
if (m.isChunked()) {
// A chunked message - remove 'Transfer-Encoding' header,
// initialize the cumulative buffer, and wait for incoming chunks.
List<String> encodings = m.getHeaders(HttpHeaders.Names.TRANSFER_ENCODING);
encodings.remove(HttpHeaders.Values.CHUNKED);
if (encodings.isEmpty()) {
m.removeHeader(HttpHeaders.Names.TRANSFER_ENCODING);
}
m.setContent(ChannelBuffers.dynamicBuffer(e.getChannel().getConfig().getBufferFactory()));
this.currentMessage = m;
} else {
// Not a chunked message - pass through.
this.currentMessage = null;
ctx.sendUpstream(e);
}
} else if (msg instanceof HttpChunk) {
// Sanity check
if (currentMessage == null) {
throw new IllegalStateException(
"received " + HttpChunk.class.getSimpleName() +
" without " + HttpMessage.class.getSimpleName());
}
// Merge the received chunk into the content of the current message.
HttpChunk chunk = (HttpChunk) msg;
ChannelBuffer content = currentMessage.getContent();
if (content.readableBytes() > maxContentLength - chunk.getContent().readableBytes()) {
throw new TooLongFrameException(
"HTTP content length exceeded " + maxContentLength +
" bytes.");
}
content.writeBytes(chunk.getContent());
if (chunk.isLast()) {
this.currentMessage = null;
currentMessage.setHeader(
HttpHeaders.Names.CONTENT_LENGTH,
String.valueOf(content.readableBytes()));
if (chunk instanceof HttpChunkTrailer)
{
HttpChunkTrailer chunkTrailer = (HttpChunkTrailer)chunk;
for(Entry<String, String> footer: chunkTrailer.getHeaders())
{
currentMessage.setHeader(footer.getKey(), footer.getValue());
}
}
Channels.fireMessageReceived(ctx, currentMessage, e.getRemoteAddress());
}
} else {
// Neither HttpMessage or HttpChunk
ctx.sendUpstream(e);
}
}
}