/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * 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.
 */

package io.undertow.server.handlers;

import io.undertow.UndertowLogger;
import io.undertow.connector.PooledByteBuffer;
import io.undertow.server.Connectors;
import io.undertow.server.HandlerWrapper;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.handlers.builder.HandlerBuilder;
import io.undertow.server.protocol.http.HttpContinue;
import org.xnio.ChannelListener;
import org.xnio.IoUtils;
import org.xnio.channels.StreamSourceChannel;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.Map;
import java.util.Set;
Handler that will buffer all request data
Author:Stuart Douglas
/** * * Handler that will buffer all request data * * @author Stuart Douglas */
public class RequestBufferingHandler implements HttpHandler { private final HttpHandler next; private final int maxBuffers; public RequestBufferingHandler(HttpHandler next, int maxBuffers) { this.next = next; this.maxBuffers = maxBuffers; } @Override public void handleRequest(final HttpServerExchange exchange) throws Exception { if(!exchange.isRequestComplete() && !HttpContinue.requiresContinueResponse(exchange.getRequestHeaders())) { final StreamSourceChannel channel = exchange.getRequestChannel(); int readBuffers = 0; final PooledByteBuffer[] bufferedData = new PooledByteBuffer[maxBuffers]; PooledByteBuffer buffer = exchange.getConnection().getByteBufferPool().allocate(); try { do { int r; ByteBuffer b = buffer.getBuffer(); r = channel.read(b); if (r == -1) { if (b.position() == 0) { buffer.close(); } else { b.flip(); bufferedData[readBuffers] = buffer; } break; } else if (r == 0) { final PooledByteBuffer finalBuffer = buffer; final int finalReadBuffers = readBuffers; channel.getReadSetter().set(new ChannelListener<StreamSourceChannel>() { PooledByteBuffer buffer = finalBuffer; int readBuffers = finalReadBuffers; @Override public void handleEvent(StreamSourceChannel channel) { try { do { int r; ByteBuffer b = buffer.getBuffer(); r = channel.read(b); if (r == -1) { if (b.position() == 0) { buffer.close(); } else { b.flip(); bufferedData[readBuffers] = buffer; } Connectors.ungetRequestBytes(exchange, bufferedData); Connectors.resetRequestChannel(exchange); channel.getReadSetter().set(null); channel.suspendReads(); Connectors.executeRootHandler(next, exchange); return; } else if (r == 0) { return; } else if (!b.hasRemaining()) { b.flip(); bufferedData[readBuffers++] = buffer; if (readBuffers == maxBuffers) { Connectors.ungetRequestBytes(exchange, bufferedData); Connectors.resetRequestChannel(exchange); channel.getReadSetter().set(null); channel.suspendReads(); Connectors.executeRootHandler(next, exchange); return; } buffer = exchange.getConnection().getByteBufferPool().allocate(); } } while (true); } catch (Throwable t) { if (t instanceof IOException) { UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) t); } else { UndertowLogger.REQUEST_IO_LOGGER.handleUnexpectedFailure(t); } for (int i = 0; i < bufferedData.length; ++i) { IoUtils.safeClose(bufferedData[i]); } if (buffer != null && buffer.isOpen()) { IoUtils.safeClose(buffer); } exchange.endExchange(); } } }); channel.resumeReads(); return; } else if (!b.hasRemaining()) { b.flip(); bufferedData[readBuffers++] = buffer; if (readBuffers == maxBuffers) { break; } buffer = exchange.getConnection().getByteBufferPool().allocate(); } } while (true); Connectors.ungetRequestBytes(exchange, bufferedData); Connectors.resetRequestChannel(exchange); } catch (Exception | Error e) { for (int i = 0; i < bufferedData.length; ++i) { IoUtils.safeClose(bufferedData[i]); } if (buffer != null && buffer.isOpen()) { IoUtils.safeClose(buffer); } throw e; } } next.handleRequest(exchange); } public static final class Wrapper implements HandlerWrapper { private final int maxBuffers; public Wrapper(int maxBuffers) { this.maxBuffers = maxBuffers; } @Override public HttpHandler wrap(HttpHandler handler) { return new RequestBufferingHandler(handler, maxBuffers); } } @Override public String toString() { return "buffer-request( " + maxBuffers + " )"; } public static final class Builder implements HandlerBuilder { @Override public String name() { return "buffer-request"; } @Override public Map<String, Class<?>> parameters() { return Collections.<String, Class<?>>singletonMap("buffers", Integer.class); } @Override public Set<String> requiredParameters() { return Collections.singleton("buffers"); } @Override public String defaultParameter() { return "buffers"; } @Override public HandlerWrapper build(Map<String, Object> config) { return new Wrapper((Integer) config.get("buffers")); } } }