Harden WebSocket transport and permessage-deflate for RFC 6455/7692 conformance

httpclient5-websocket/src/main/java/org/apache/hc/client5/http/websocket/client/impl/protocol/Http1UpgradeProtocol.java

@@ -158,7 +158,9 @@ public void completed(final WebSocketEndpointConnector.ProtoEndpoint endpoint) {
                                     ext.append("; client_no_context_takeover");
                                 }
                                 if (cfg.getOfferClientMaxWindowBits() != null) {
-                                    ext.append("; client_max_window_bits=").append(cfg.getOfferClientMaxWindowBits());
+                                    // The JDK Deflater always compresses with a 15-bit window, so a
+                                    // smaller configured value must not be promised to the server.
+                                    ext.append("; client_max_window_bits=15");
                                 }
                                 if (cfg.getOfferServerMaxWindowBits() != null) {
                                     ext.append("; server_max_window_bits=").append(cfg.getOfferServerMaxWindowBits());
@@ -359,9 +361,28 @@ public void cancelled() {
         }
     }
 
-    private static String headerValue(final HttpResponse r, final String name) {
-        final Header h = r.getFirstHeader(name);
-        return h != null ? h.getValue() : null;
+    static String headerValue(final HttpResponse r, final String name) {
+        // RFC 6455 section 9.1 permits Sec-WebSocket-Extensions and Sec-WebSocket-Protocol
+        // to be split across multiple header fields; combine them as a comma-separated list.
+        final Header[] headers = r.getHeaders(name);
+        if (headers.length == 0) {
+            return null;
+        }
+        if (headers.length == 1) {
+            return headers[0].getValue();
+        }
+        final StringBuilder buf = new StringBuilder();
+        for (final Header h : headers) {
+            final String value = h.getValue();
+            if (value == null || value.isEmpty()) {
+                continue;
+            }
+            if (buf.length() > 0) {
+                buf.append(", ");
+            }
+            buf.append(value);
+        }
+        return buf.length() > 0 ? buf.toString() : null;
     }
 
     private static boolean containsToken(final HttpResponse r, final String header, final String token) {
@@ -394,9 +415,10 @@ static ExtensionChain buildExtensionChain(final WebSocketClientConfig cfg, final
         }
         boolean pmceSeen = false, serverNoCtx = false, clientNoCtx = false;
         Integer clientBits = null, serverBits = null;
-        final boolean offerServerNoCtx = cfg.isOfferServerNoContextTakeover();
-        final boolean offerClientNoCtx = cfg.isOfferClientNoContextTakeover();
-        final Integer offerClientBits = cfg.getOfferClientMaxWindowBits();
+        // The offer always advertises client_max_window_bits=15 because the JDK Deflater
+        // cannot compress with a smaller window, so the response is validated against 15
+        // regardless of the configured value.
+        final Integer offerClientBits = cfg.getOfferClientMaxWindowBits() != null ? Integer.valueOf(15) : null;
         final Integer offerServerBits = cfg.getOfferServerMaxWindowBits();
 
         final String[] tokens = ext.split(",");
@@ -418,15 +440,12 @@ static ExtensionChain buildExtensionChain(final WebSocketClientConfig cfg, final
                 final String p = parts[i].trim();
                 final int eq = p.indexOf('=');
                 if (eq < 0) {
+                    // RFC 7692 sections 7.1.1.1 and 7.1.1.2 permit the server to include either
+                    // no-context-takeover parameter in the response even when the offer did not;
+                    // both are always safe to honour.
                     if ("server_no_context_takeover".equalsIgnoreCase(p)) {
-                        if (!offerServerNoCtx) {
-                            throw new IllegalStateException("Server selected server_no_context_takeover not offered");
-                        }
                         serverNoCtx = true;
                     } else if ("client_no_context_takeover".equalsIgnoreCase(p)) {
-                        if (!offerClientNoCtx) {
-                            throw new IllegalStateException("Server selected client_no_context_takeover not offered");
-                        }
                         clientNoCtx = true;
                     } else {
                         throw new IllegalStateException("Unsupported permessage-deflate parameter: " + p);
@@ -453,9 +472,8 @@ static ExtensionChain buildExtensionChain(final WebSocketClientConfig cfg, final
                             throw new IllegalStateException("Invalid client_max_window_bits: " + v, nfe);
                         }
                     } else if ("server_max_window_bits".equalsIgnoreCase(k)) {
-                        if (offerServerBits == null) {
-                            throw new IllegalStateException("Server selected server_max_window_bits not offered");
-                        }
+                        // RFC 7692 section 7.1.2.1 permits the server to include this parameter
+                        // uninvited; a server constraining its own window is always acceptable.
                         try {
                             if (v.isEmpty()) {
                                 throw new IllegalStateException("server_max_window_bits must have a value");

httpclient5-websocket/src/main/java/org/apache/hc/client5/http/websocket/client/impl/protocol/Http2ExtendedConnectProtocol.java

@@ -211,8 +211,9 @@ public void produceRequest(final RequestChannel channel, final HttpContext conte
         public void consumeResponse(final HttpResponse response, final EntityDetails entityDetails,
                                     final HttpContext context) throws HttpException, IOException {
 
-            if (response.getCode() != HttpStatus.SC_OK) {
-                failFuture(new IllegalStateException("Unexpected status: " + response.getCode()));
+            final int code = response.getCode();
+            if (code < HttpStatus.SC_OK || code >= HttpStatus.SC_REDIRECTION) {
+                failFuture(new IllegalStateException("Unexpected status: " + code));
                 return;
             }
 
@@ -321,8 +322,7 @@ public void cancel() {
         }
 
         private static String headerValue(final HttpResponse r, final String name) {
-            final Header h = r.getFirstHeader(name);
-            return h != null ? h.getValue() : null;
+            return Http1UpgradeProtocol.headerValue(r, name);
         }
 
         private void failFuture(final Exception ex) {

httpclient5-websocket/src/main/java/org/apache/hc/client5/http/websocket/transport/DataStreamChannelTransport.java

@@ -28,6 +28,7 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.channels.CancelledKeyException;
 
 import org.apache.hc.core5.annotation.Internal;
 import org.apache.hc.core5.http.nio.DataStreamChannel;
@@ -51,14 +52,24 @@ public int write(final ByteBuffer src) throws IOException {
         if (ch == null) {
             return 0;
         }
-        return ch.write(src);
+        try {
+            return ch.write(src);
+        } catch (final CancelledKeyException ignore) {
+            // The selection key was cancelled by a concurrent shutdown; the channel is gone.
+            return 0;
+        }
     }
 
     @Override
     public void requestOutput() {
         final DataStreamChannel ch = channel;
         if (ch != null) {
-            ch.requestOutput();
+            try {
+                ch.requestOutput();
+            } catch (final CancelledKeyException ignore) {
+                // requestOutput is a best-effort nudge; if the channel's selection key was
+                // cancelled by a concurrent shutdown there is nothing left to flush.
+            }
         }
     }
 

httpclient5-websocket/src/main/java/org/apache/hc/client5/http/websocket/transport/WebSocketSessionEngine.java

@@ -109,6 +109,7 @@ public final class WebSocketSessionEngine {
     final AtomicBoolean open = new AtomicBoolean(true);
     final AtomicBoolean closeSent = new AtomicBoolean(false);
     private final AtomicBoolean closeReceived = new AtomicBoolean(false);
+    private final AtomicBoolean released = new AtomicBoolean(false);
     volatile boolean closeAfterFlush;
     private volatile ScheduledFuture<?> closeTimeoutFuture;
 
@@ -332,6 +333,12 @@ private void handleFrame() {
             return;
         }
         if (FrameOpcode.isControl(op)) {
+            if (r1) {
+                // Control frames are never compressed, so an extension never defines RSV1 for them.
+                initiateClose(1002, "RSV1 set on control frame");
+                inbuf.clear();
+                return;
+            }
             if (!fin) {
                 initiateClose(1002, "fragmented control frame");
                 inbuf.clear();
@@ -379,8 +386,7 @@ private void handleFrame() {
                     inbuf.clear();
                     return;
                 }
-                appendToMessage(payload);
-                if (fin) {
+                if (appendToMessage(payload) && fin) {
                     deliverAssembledMessage();
                 }
                 break;
@@ -484,15 +490,16 @@ private void startMessage(final int opcode, final ByteBuffer payload, final bool
         appendToMessage(payload);
     }
 
-    private void appendToMessage(final ByteBuffer payload) {
+    private boolean appendToMessage(final ByteBuffer payload) {
         final int n = payload.remaining();
         assemblingSize += n;
         if (cfg.getMaxMessageSize() > 0 && assemblingSize > cfg.getMaxMessageSize()) {
             initiateClose(1009, "Message too big");
-            return;
+            return false;
         }
         assemblingBuf = WebSocketBufferOps.ensureCapacity(assemblingBuf, n);
         assemblingBuf.put(payload.asReadOnlyBuffer());
+        return true;
     }
 
     private void deliverAssembledMessage() {
@@ -673,6 +680,9 @@ boolean enqueueData(final OutFrame frame) {
     }
 
     private void drainAndRelease() {
+        if (!released.compareAndSet(false, true)) {
+            return;
+        }
         if (activeWrite != null) {
             if (activeWrite.dataFrame) {
                 dataQueuedBytes.addAndGet(-activeWrite.size);
@@ -690,6 +700,12 @@ private void drainAndRelease() {
                 dataQueuedBytes.addAndGet(-f.size);
             }
         }
+        if (encChain != null) {
+            encChain.close();
+        }
+        if (decChain != null) {
+            decChain.close();
+        }
         cancelCloseTimeout();
     }
 

httpclient5-websocket/src/main/java/org/apache/hc/core5/websocket/PerMessageDeflateExtension.java

@@ -144,17 +144,21 @@ public ByteBuffer encode(final WebSocketFrameType type, final boolean fin, final
         deflater.setInput(input);
         final ByteArrayOutputStream out = new ByteArrayOutputStream(Math.max(128, input.length / 2));
         final byte[] buffer = new byte[Math.min(16384, Math.max(1024, input.length))];
-        while (!deflater.needsInput()) {
-            final int count = deflater.deflate(buffer, 0, buffer.length, Deflater.SYNC_FLUSH);
+        // Drain until a SYNC_FLUSH leaves the output buffer partly filled; testing needsInput()
+        // would stop once the input is consumed and could drop the trailing 00 00 FF FF flush bytes.
+        int count;
+        do {
+            count = deflater.deflate(buffer, 0, buffer.length, Deflater.SYNC_FLUSH);
             if (count > 0) {
                 out.write(buffer, 0, count);
-            } else {
-                break;
             }
-        }
+        } while (count == buffer.length);
         final byte[] data = out.toByteArray();
         final ByteBuffer encoded;
-        if (data.length >= 4) {
+        // Strip the 00 00 FF FF flush trailer only on the final fragment; non-final fragments
+        // must keep it so each intermediate empty stored block stays valid and the reassembled
+        // DEFLATE stream decodes (RFC 7692 section 7.2.1).
+        if (fin && data.length >= 4) {
             encoded = ByteBuffer.wrap(data, 0, data.length - 4);
         } else {
             encoded = ByteBuffer.wrap(data);
@@ -183,6 +187,12 @@ public WebSocketExtensionData getResponseData() {
         return new WebSocketExtensionData(getName(), params);
     }
 
+    @Override
+    public void close() {
+        deflater.end();
+        inflater.end();
+    }
+
     private static boolean isDataFrame(final WebSocketFrameType type) {
         return type == WebSocketFrameType.TEXT || type == WebSocketFrameType.BINARY;
     }

httpclient5-websocket/src/main/java/org/apache/hc/core5/websocket/WebSocketExtension.java

@@ -77,4 +77,13 @@ default ByteBuffer encode(
     default WebSocketExtensionData getResponseData() {
         return new WebSocketExtensionData(getName(), null);
     }
+
+    /**
+     * Releases any native resources held by this extension (e.g. a {@code Deflater} or
+     * {@code Inflater}). Called once when the owning session terminates.
+     *
+     * @since 5.7
+     */
+    default void close() {
+    }
 }

httpclient5-websocket/src/main/java/org/apache/hc/core5/websocket/WebSocketExtensionRegistry.java

@@ -27,9 +27,11 @@
 package org.apache.hc.core5.websocket;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.apache.hc.core5.util.Args;
 
@@ -52,14 +54,22 @@ public WebSocketExtensionNegotiation negotiate(
             final boolean server) throws WebSocketException {
         final List<WebSocketExtension> extensions = new ArrayList<>();
         final List<WebSocketExtensionData> responseData = new ArrayList<>();
+        final Set<String> accepted = new HashSet<>();
         if (requested != null) {
             for (final WebSocketExtensionData request : requested) {
+                // At most one offer per extension name is accepted; a second accepted offer of the
+                // same extension would claim the same RSV bit and break both the response header and
+                // the RSV bookkeeping in the frame reader.
+                if (accepted.contains(request.getName())) {
+                    continue;
+                }
                 final WebSocketExtensionFactory factory = factories.get(request.getName());
                 if (factory != null) {
                     final WebSocketExtension extension = factory.create(request, server);
                     if (extension != null) {
                         extensions.add(extension);
                         responseData.add(extension.getResponseData());
+                        accepted.add(request.getName());
                     }
                 }
             }

httpclient5-websocket/src/main/java/org/apache/hc/core5/websocket/WebSocketFrameReader.java

@@ -98,6 +98,9 @@ WebSocketFrame readFrame() throws IOException {
                 len = (len << 8) | (readByte() & 0xFF);
             }
         }
+        if (len < 0) {
+            throw new WebSocketException("64-bit frame length must have the most significant bit set to 0");
+        }
         if (len > Integer.MAX_VALUE) {
             throw new WebSocketException("Frame payload too large: " + len);
         }

httpclient5-websocket/src/main/java/org/apache/hc/core5/websocket/WebSocketFrameWriter.java

@@ -34,6 +34,7 @@
 import java.util.List;
 
 import org.apache.hc.core5.util.Args;
+import org.apache.hc.core5.websocket.message.CloseCodec;
 
 class WebSocketFrameWriter {
 
@@ -63,13 +64,9 @@ void writePong(final ByteBuffer payload) throws IOException {
     }
 
     void writeClose(final int statusCode, final String reason) throws IOException {
-        final byte[] reasonBytes = reason != null ? reason.getBytes(StandardCharsets.UTF_8) : new byte[0];
-        final int len = 2 + reasonBytes.length;
-        final ByteBuffer buffer = ByteBuffer.allocate(len);
-        buffer.put((byte) ((statusCode >> 8) & 0xFF));
-        buffer.put((byte) (statusCode & 0xFF));
-        buffer.put(reasonBytes);
-        buffer.flip();
+        // CloseCodec truncates the reason to 123 UTF-8 bytes so the payload never
+        // exceeds the 125-byte control frame limit (RFC 6455 section 5.5).
+        final ByteBuffer buffer = ByteBuffer.wrap(CloseCodec.encode(statusCode, reason));
         writeFrame(WebSocketFrameType.CLOSE, buffer, false, false, false);
     }
 
@@ -102,6 +99,9 @@ private void writeFrame(
         Args.notNull(type, "Frame type");
         final ByteBuffer buffer = payload != null ? payload.asReadOnlyBuffer() : ByteBuffer.allocate(0);
         final int payloadLen = buffer.remaining();
+        if (type.isControl() && payloadLen > 125) {
+            throw new IllegalArgumentException("Control frame payload > 125 bytes: " + payloadLen);
+        }
         int firstByte = 0x80 | (type.getOpcode() & 0x0F);
         if (rsv1) {
             firstByte |= 0x40;

httpclient5-websocket/src/main/java/org/apache/hc/core5/websocket/WebSocketSession.java

@@ -52,6 +52,7 @@ public final class WebSocketSession {
     private final SocketAddress localAddress;
     private final WebSocketFrameReader reader;
     private final WebSocketFrameWriter writer;
+    private final List<WebSocketExtension> extensions;
     private final ReentrantLock writeLock = new ReentrantLock();
     private volatile boolean closeSent;
 
@@ -68,11 +69,24 @@ public WebSocketSession(
         this.remoteAddress = remoteAddress;
         this.localAddress = localAddress;
         final List<WebSocketExtension> negotiated = extensions != null ? extensions : Collections.emptyList();
+        this.extensions = negotiated;
         this.reader = new WebSocketFrameReader(this.config, this.inputStream, negotiated);
         this.writer = new WebSocketFrameWriter(this.outputStream, negotiated);
         this.closeSent = false;
     }
 
+    /**
+     * Releases native resources held by the negotiated extensions. Invoked once when the
+     * session processing loop terminates.
+     *
+     * @since 5.7
+     */
+    public void releaseExtensions() {
+        for (final WebSocketExtension extension : extensions) {
+            extension.close();
+        }
+    }
+
     public SocketAddress getRemoteAddress() {
         return remoteAddress;
     }