UNDERTOW-895 Improve MAX_PARAMETERS and MAX_HEADERS handling across a…

…ll protocols

core/src/main/java/io/undertow/UndertowMessages.java

@@ -33,6 +33,7 @@
 import io.undertow.security.api.AuthenticationMechanism;
 import io.undertow.server.handlers.builder.HandlerBuilder;
 import io.undertow.util.HttpString;
+import io.undertow.util.ParameterLimitException;
 
 /**
  * @author Stuart Douglas
@@ -172,7 +173,7 @@ public interface UndertowMessages {
     IllegalStateException tooManyCookies(int maxCookies);
 
     @Message(id = 47, value = "The number of parameters exceeded the maximum of %s")
-    IllegalStateException tooManyParameters(int maxValues);
+    ParameterLimitException tooManyParameters(int maxValues);
 
     @Message(id = 48, value = "No request is currently active")
     IllegalStateException noRequestActive();

core/src/main/java/io/undertow/client/ajp/AjpClientExchange.java

@@ -79,13 +79,19 @@ class AjpClientExchange extends AbstractAttachable implements ClientExchange {
 
     void terminateRequest() {
         state |= REQUEST_TERMINATED;
+        if(!clientConnection.isOpen()) {
+            state |= RESPONSE_TERMINATED;
+        }
         if (anyAreSet(state, RESPONSE_TERMINATED)) {
             clientConnection.requestDone();
         }
     }
 
     void terminateResponse() {
         state |= RESPONSE_TERMINATED;
+        if(!clientConnection.isOpen()) {
+            state |= REQUEST_TERMINATED;
+        }
         if (anyAreSet(state, REQUEST_TERMINATED)) {
             clientConnection.requestDone();
         }

core/src/main/java/io/undertow/protocols/http2/Http2HeaderBlockParser.java

@@ -74,9 +74,6 @@ abstract class Http2HeaderBlockParser extends Http2PushBackParser implements Hpa
 
     @Override
     protected void handleData(ByteBuffer resource, Http2FrameHeaderParser header) throws IOException {
-        if(maxHeaders > 0 && headerMap.size() >= maxHeaders) {
-            throw new IOException(UndertowMessages.MESSAGES.tooManyHeaders(maxHeaders));
-        }
         boolean continuationFramesComing = Bits.anyAreClear(header.flags, Http2Channel.HEADERS_FLAG_END_HEADERS);
         if (frameRemaining == -1) {
             frameRemaining = header.length;
@@ -110,6 +107,10 @@ protected void handleData(ByteBuffer resource, Http2FrameHeaderParser header) th
             } catch (HpackException e) {
                 throw new ConnectionErrorException(e.getCloseCode(), e);
             }
+
+            if(maxHeaders > 0 && headerMap.size() > maxHeaders) {
+                throw new StreamErrorException(Http2Channel.ERROR_FRAME_SIZE_ERROR);
+            }
             if(oldLimit != -1) {
                 if(resource.remaining() == 0) {
                     int paddingInBuffer = oldLimit - resource.limit();
@@ -135,7 +136,11 @@ HeaderMap getHeaderMap() {
 
     @Override
     public void emitHeader(HttpString name, String value, boolean neverIndex) throws HpackException {
+        if(maxHeaders > 0 && headerMap.size() > maxHeaders) {
+            return;
+        }
         headerMap.add(name, value);
+
         if(name.length() == 0) {
             throw UndertowMessages.MESSAGES.invalidHeader();
         }

core/src/main/java/io/undertow/server/Connectors.java

@@ -23,6 +23,7 @@
 import io.undertow.server.handlers.Cookie;
 import io.undertow.util.DateUtils;
 import io.undertow.util.Headers;
+import io.undertow.util.ParameterLimitException;
 import io.undertow.util.StatusCodes;
 import io.undertow.util.URLUtils;
 import io.undertow.connector.PooledByteBuffer;
@@ -253,7 +254,11 @@ public static void executeRootHandler(final HttpHandler handler, final HttpServe
      */
     @Deprecated
     public static void setExchangeRequestPath(final HttpServerExchange exchange, final String encodedPath, final String charset, boolean decode, final boolean allowEncodedSlash, StringBuilder decodeBuffer) {
-        setExchangeRequestPath(exchange, encodedPath, charset, decode, allowEncodedSlash, decodeBuffer, exchange.getConnection().getUndertowOptions().get(UndertowOptions.MAX_PARAMETERS, UndertowOptions.DEFAULT_MAX_PARAMETERS));
+        try {
+            setExchangeRequestPath(exchange, encodedPath, charset, decode, allowEncodedSlash, decodeBuffer, exchange.getConnection().getUndertowOptions().get(UndertowOptions.MAX_PARAMETERS, UndertowOptions.DEFAULT_MAX_PARAMETERS));
+        } catch (ParameterLimitException e) {
+            throw new RuntimeException(e);
+        }
     }
         /**
          * Sets the request path and query parameters, decoding to the requested charset.
@@ -262,7 +267,7 @@ public static void setExchangeRequestPath(final HttpServerExchange exchange, fin
          * @param encodedPath        The encoded path
          * @param charset     The charset
          */
-    public static void setExchangeRequestPath(final HttpServerExchange exchange, final String encodedPath, final String charset, boolean decode, final boolean allowEncodedSlash, StringBuilder decodeBuffer, int maxParameters) {
+    public static void setExchangeRequestPath(final HttpServerExchange exchange, final String encodedPath, final String charset, boolean decode, final boolean allowEncodedSlash, StringBuilder decodeBuffer, int maxParameters) throws ParameterLimitException {
         boolean requiresDecode = false;
         for (int i = 0; i < encodedPath.length(); ++i) {
             char c = encodedPath.charAt(i);

core/src/main/java/io/undertow/server/handlers/form/FormData.java

@@ -75,7 +75,7 @@ public void add(String name, String value, final HeaderMap headers) {
         }
         values.add(new FormValueImpl(value, headers));
         if (++valueCount > maxValues) {
-            throw UndertowMessages.MESSAGES.tooManyParameters(maxValues);
+            throw new RuntimeException(UndertowMessages.MESSAGES.tooManyParameters(maxValues));
         }
     }
 
@@ -86,10 +86,10 @@ public void add(String name, Path value, String fileName, final HeaderMap header
         }
         values.add(new FormValueImpl(value, fileName, headers));
         if (values.size() > maxValues) {
-            throw UndertowMessages.MESSAGES.tooManyParameters(maxValues);
+            throw new RuntimeException(UndertowMessages.MESSAGES.tooManyParameters(maxValues));
         }
         if (++valueCount > maxValues) {
-            throw UndertowMessages.MESSAGES.tooManyParameters(maxValues);
+            throw new RuntimeException(UndertowMessages.MESSAGES.tooManyParameters(maxValues));
         }
     }
 
@@ -102,7 +102,7 @@ public void put(String name, String value, final HeaderMap headers) {
         values.add(new FormValueImpl(value, headers));
 
         if (++valueCount > maxValues) {
-            throw UndertowMessages.MESSAGES.tooManyParameters(maxValues);
+            throw new RuntimeException(UndertowMessages.MESSAGES.tooManyParameters(maxValues));
         }
     }
 

core/src/main/java/io/undertow/server/handlers/proxy/LoadBalancingProxyClient.java

@@ -365,6 +365,19 @@ protected Host findStickyHost(HttpServerExchange exchange) {
         return null;
     }
 
+    /**
+     * Should only be used for tests
+     *
+     * DO NOT CALL THIS METHOD WHEN REQUESTS ARE IN PROGRESS
+     *
+     * It is not thread safe so internal state can get messed up.
+     */
+    public void closeCurrentConnections() {
+        for(Host host : hosts) {
+            host.closeCurrentConnections();
+        }
+    }
+
     public final class Host extends ConnectionPoolErrorHandler.SimpleConnectionPoolErrorHandler implements ConnectionPoolManager {
         final ProxyConnectionPool connectionPool;
         final String jvmRoute;
@@ -411,6 +424,10 @@ public int getMaxQueueSize() {
         public URI getUri() {
             return uri;
         }
+
+        void closeCurrentConnections() {
+            connectionPool.closeCurrentConnections();
+        }
     }
 
     private static class ExclusiveConnectionHolder {

core/src/main/java/io/undertow/server/handlers/proxy/ProxyClient.java

@@ -60,7 +60,7 @@ public interface ProxyClient {
     /**
      * An opaque interface that may contain information about the proxy target
      */
-    public interface ProxyTarget {
+    interface ProxyTarget {
 
     }
 }

core/src/main/java/io/undertow/server/handlers/proxy/ProxyConnectionPool.java

@@ -42,6 +42,7 @@
 import java.net.URI;
 import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.Map;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -536,6 +537,21 @@ public void connect(ProxyClient.ProxyTarget proxyTarget, HttpServerExchange exch
         }
     }
 
+    /**
+     * Should only be used for tests.
+     *
+     */
+    void closeCurrentConnections() {
+        for(Map.Entry<XnioIoThread, HostThreadData> data : hostThreadData.entrySet()) {
+            ConnectionHolder d = data.getValue().availableConnections.poll();
+            while (d != null) {
+                IoUtils.safeClose(d.clientConnection);
+                d = data.getValue().availableConnections.poll();
+            }
+            data.getValue().connections = 0;
+        }
+    }
+
     private final class HostThreadData {
 
         int connections = 0;

core/src/main/java/io/undertow/server/handlers/proxy/ProxyHandler.java

@@ -585,7 +585,7 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception {
                                 result.getRequestChannel().getWriteSetter().set(ChannelListeners.flushingChannelListener(new ChannelListener<StreamSinkChannel>() {
                                     @Override
                                     public void handleEvent(StreamSinkChannel channel) {
-                                        Transfer.initiateTransfer(exchange.getRequestChannel(), result.getRequestChannel(), ChannelListeners.closingChannelListener(), new HTTPTrailerChannelListener(exchange, result), handler, handler, exchange.getConnection().getByteBufferPool());
+                                        Transfer.initiateTransfer(exchange.getRequestChannel(), result.getRequestChannel(), ChannelListeners.closingChannelListener(), new HTTPTrailerChannelListener(exchange, result, exchange, proxyClientHandler, idempotentPredicate), handler, handler, exchange.getConnection().getByteBufferPool());
 
                                     }
                                 }, handler));
@@ -596,7 +596,7 @@ public void handleEvent(StreamSinkChannel channel) {
                             handler.handleException(result.getRequestChannel(), e);
                         }
                     }
-                    HTTPTrailerChannelListener trailerListener = new HTTPTrailerChannelListener(exchange, result);
+                    HTTPTrailerChannelListener trailerListener = new HTTPTrailerChannelListener(exchange, result, exchange, proxyClientHandler, idempotentPredicate);
                     if(!exchange.isRequestComplete()) {
                         Transfer.initiateTransfer(exchange.getRequestChannel(), result.getRequestChannel(), ChannelListeners.closingChannelListener(), trailerListener, handler, handler, exchange.getConnection().getByteBufferPool());
                     } else {
@@ -607,24 +607,26 @@ public void handleEvent(StreamSinkChannel channel) {
 
                 @Override
                 public void failed(IOException e) {
-                    UndertowLogger.PROXY_REQUEST_LOGGER.proxyRequestFailed(exchange.getRequestURI(), e);
-                    if(idempotentPredicate.resolve(exchange) && proxyClientHandler != null) {
-                        proxyClientHandler.failed(exchange); //this will attempt a retry if configured to do so
-                    } else {
-                        if (!exchange.isResponseStarted()) {
-                            exchange.setStatusCode(StatusCodes.SERVICE_UNAVAILABLE);
-                            exchange.endExchange();
-                        } else {
-                            IoUtils.safeClose(exchange.getConnection());
-                        }
-                    }
+                    handleFailure(exchange, proxyClientHandler, idempotentPredicate, e);
                 }
             });
 
 
         }
     }
 
+    static void handleFailure(HttpServerExchange exchange, ProxyClientHandler proxyClientHandler, Predicate idempotentRequestPredicate, IOException e) {
+        UndertowLogger.PROXY_REQUEST_LOGGER.proxyRequestFailed(exchange.getRequestURI(), e);
+        if(exchange.isResponseStarted()) {
+            IoUtils.safeClose(exchange.getConnection());
+        } else if(idempotentRequestPredicate.resolve(exchange) && proxyClientHandler != null) {
+            proxyClientHandler.failed(exchange); //this will attempt a retry if configured to do so
+        } else {
+            exchange.setStatusCode(StatusCodes.SERVICE_UNAVAILABLE);
+            exchange.endExchange();
+        }
+    }
+
     private static final class ResponseCallback implements ClientCallback<ClientExchange> {
 
         private final HttpServerExchange exchange;
@@ -674,33 +676,29 @@ public void handleUpgrade(StreamConnection streamConnection, HttpServerExchange
                 });
             }
             final IoExceptionHandler handler = new IoExceptionHandler(exchange, result.getConnection());
-            Transfer.initiateTransfer(result.getResponseChannel(), exchange.getResponseChannel(), ChannelListeners.closingChannelListener(), new HTTPTrailerChannelListener(result, exchange), handler, handler, exchange.getConnection().getByteBufferPool());
+            Transfer.initiateTransfer(result.getResponseChannel(), exchange.getResponseChannel(), ChannelListeners.closingChannelListener(), new HTTPTrailerChannelListener(result, exchange, exchange, proxyClientHandler, idempotentPredicate), handler, handler, exchange.getConnection().getByteBufferPool());
         }
 
         @Override
         public void failed(IOException e) {
-            UndertowLogger.PROXY_REQUEST_LOGGER.proxyRequestFailed(exchange.getRequestURI(), e);
-            if (!exchange.isResponseStarted()) {
-                if(idempotentPredicate.resolve(exchange) && proxyClientHandler != null) {
-                    proxyClientHandler.failed(exchange);
-                } else {
-                    exchange.setStatusCode(StatusCodes.INTERNAL_SERVER_ERROR);
-                    exchange.endExchange();
-                }
-            } else {
-                IoUtils.safeClose(exchange.getConnection());
-            }
+            handleFailure(exchange, proxyClientHandler, idempotentPredicate, e);
         }
     }
 
     private static final class HTTPTrailerChannelListener implements ChannelListener<StreamSinkChannel> {
 
         private final Attachable source;
         private final Attachable target;
+        private final HttpServerExchange exchange;
+        private final ProxyClientHandler proxyClientHandler;
+        private final Predicate idempotentPredicate;
 
-        private HTTPTrailerChannelListener(final Attachable source, final Attachable target) {
+        private HTTPTrailerChannelListener(final Attachable source, final Attachable target, HttpServerExchange exchange, ProxyClientHandler proxyClientHandler, Predicate idempotentPredicate) {
             this.source = source;
             this.target = target;
+            this.exchange = exchange;
+            this.proxyClientHandler = proxyClientHandler;
+            this.idempotentPredicate = idempotentPredicate;
         }
 
         @Override
@@ -725,11 +723,9 @@ public void handleEvent(StreamSinkChannel channel) {
                     channel.shutdownWrites();
                 }
             } catch (IOException e) {
-                UndertowLogger.REQUEST_IO_LOGGER.ioException(e);
-                IoUtils.safeClose(channel);
+                handleFailure(exchange, proxyClientHandler, idempotentPredicate, e);
             } catch (Exception e) {
-                UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
-                IoUtils.safeClose(channel);
+                handleFailure(exchange, proxyClientHandler, idempotentPredicate, new IOException(e));
             }
 
         }

core/src/main/java/io/undertow/server/protocol/ajp/AjpReadListener.java

@@ -34,6 +34,7 @@
 import io.undertow.util.Methods;
 import org.xnio.ChannelListener;
 import io.undertow.connector.PooledByteBuffer;
+import io.undertow.util.StatusCodes;
 import org.xnio.StreamConnection;
 import org.xnio.channels.StreamSinkChannel;
 import org.xnio.channels.StreamSourceChannel;
@@ -105,6 +106,7 @@ public void handleEvent(final StreamSourceChannel channel) {
             channel.suspendReads();
             return;
         }
+
         PooledByteBuffer existing = connection.getExtraBytes();
 
         final PooledByteBuffer pooled = existing == null ? connection.getByteBufferPool().allocate() : existing;
@@ -223,6 +225,7 @@ public void handleEvent(AjpServerResponseConduit channel) {
                 if(state.attributes != null) {
                     httpServerExchange.putAttachment(HttpServerExchange.REQUEST_ATTRIBUTES, state.attributes);
                 }
+                AjpRequestParseState oldState = state;
                 state = null;
                 this.httpServerExchange = null;
                 httpServerExchange.setPersistent(true);
@@ -234,7 +237,13 @@ public void handleEvent(AjpServerResponseConduit channel) {
                 if(connectorStatistics != null) {
                     connectorStatistics.setup(httpServerExchange);
                 }
-                Connectors.executeRootHandler(connection.getRootHandler(), httpServerExchange);
+
+                if(oldState.badRequest) {
+                    httpServerExchange.setStatusCode(StatusCodes.BAD_REQUEST);
+                    httpServerExchange.endExchange();
+                } else {
+                    Connectors.executeRootHandler(connection.getRootHandler(), httpServerExchange);
+                }
 
             } catch (Throwable t) {
                 //TODO: we should attempt to return a 500 status code in this situation

core/src/main/java/io/undertow/server/protocol/ajp/AjpRequestParseState.java

@@ -94,12 +94,14 @@ class AjpRequestParseState {
     public String sslCipher;
     public String sslCert;
     public String sslKeySize;
+    boolean badRequest;
 
     public void reset() {
         stringLength = -1;
         currentStringLength = 0;
         currentIntegerPart = -1;
         readHeaders = 0;
+        badRequest = false;
     }
     public boolean isComplete() {
         return state == 15;