From cd8a0373475edb2ff4ee9de956995e1f6feac8a3 Mon Sep 17 00:00:00 2001 From: Simon Mavi Stewart Date: Thu, 7 May 2026 17:25:51 +0200 Subject: [PATCH] [grid] Apply the WebSocket frame fast path on the Node MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Router has had a direct frame-forwarding path between the Netty pipeline and the upstream JDK WebSocket since db9b07aa5a (2026-03-11, "[grid] Router WebSocket handle dropped close frames, idle disconnects, high-latency proxying", #17197). Once the client-side handshake completes, an inbound WebSocketFrameProxy forwards each Netty WebSocketFrame straight to the upstream WebSocket, and the outbound DirectForwardingListener writes upstream replies directly to the client channel. Together those removed the per-frame Message allocation and the executor hop in WebSocketMessageHandler on the Router side. The Node still did the full round-trip through MessageInboundConverter, WebSocketMessageHandler, the registered Consumer, and MessageOutboundConverter in both directions for every frame. Each frame allocated a TextMessage or BinaryMessage and hopped onto the channel executor on delivery. For a busy CDP or VNC session that is measurable allocation and executor-queue pressure on the Node. Apply the same PostUpgradeHook pattern on the Node side: the consumer returned from ProxyNodeWebsockets installs a WebSocketFrameProxy after the handshake so inbound frames forward straight to the browser-side WebSocket, and a DirectForwardingListener writes outbound frames directly to the client channel. Frames received before the handshake are buffered in arrival order and drained on handover, so a frame cannot land in a pipeline that has already had its Message-layer handlers removed. The hardening that the Router-side listener picked up in 8d8cf6463a (2026-05-14, "[grid] Close pre-handshake race in WebSocket proxy", #17435) is mirrored on the Node listener: the pre-handshake buffer is capped at 128 frames with a 1009 close recorded on overflow; the close code and reason are recorded on pre-handshake close or error so a late onUpgrade can write a clean close frame to the client and tear the channel down rather than leaving it open; and the buffer is released on close so ref-counted frames cannot leak if the handshake never completes. Close-frame reasons coming from the upstream are now truncated to the 123-byte UTF-8 cap that RFC 6455 §5.5.1 imposes. The truncation uses a CharsetEncoder writing into a 120-byte buffer so it stops at a clean character boundary on overflow — a naive byte-truncate-then- decode could split a multi-byte sequence, produce a U+FFFD replacement on decode, and re-encode back over 123 bytes, breaking the close frame. The helper lives as a public static on WebSocketFrameProxy because both DirectForwardingListener classes already depend on that class. The Router-side listener that landed in #17435 had the same unchecked path; apply the helper there too so both proxies share the same safe behaviour. The Node-specific behaviour is preserved: - Session-activity heartbeats (sessionConsumer.accept(sessionId)) fire per frame, both pre- and post-handshake. - The connectionReleased CAS still guards a single node.releaseConnection call across the close and error paths, including the overflow path introduced here. - VNC sessions still install a no-op heartbeat consumer so VNC traffic does not mark the session as recently active. The existing ProxyNodeWebsocketsTest continues to exercise the slot accounting, including the regression from #17197 where onError without a follow-on onClose used to leak the slot. New unit tests in NodeDirectForwardingListenerTest pin the per-frame heartbeat, the buffer-then-drain ordering, the surface-and-teardown behaviour on a pre-handshake close, the overflow path's clean release of the session slot, and the safe truncation of an overlong upstream close reason that contains multi-byte UTF-8 characters. The shared helper has a focused unit test alongside it in WebSocketFrameProxyTest. --- .../org/openqa/selenium/grid/node/BUILD.bazel | 4 + .../grid/node/ProxyNodeWebsockets.java | 374 +++++++++++++----- .../grid/router/ProxyWebsocketsIntoGrid.java | 7 +- .../netty/server/WebSocketFrameProxy.java | 32 ++ .../org/openqa/selenium/grid/node/BUILD.bazel | 2 + .../NodeDirectForwardingListenerTest.java | 298 ++++++++++++++ .../netty/server/WebSocketFrameProxyTest.java | 25 ++ 7 files changed, 647 insertions(+), 95 deletions(-) create mode 100644 java/test/org/openqa/selenium/grid/node/NodeDirectForwardingListenerTest.java diff --git a/java/src/org/openqa/selenium/grid/node/BUILD.bazel b/java/src/org/openqa/selenium/grid/node/BUILD.bazel index 5e63ebe73db7c..43b8307238f1e 100644 --- a/java/src/org/openqa/selenium/grid/node/BUILD.bazel +++ b/java/src/org/openqa/selenium/grid/node/BUILD.bazel @@ -21,9 +21,13 @@ java_library( "//java/src/org/openqa/selenium/grid/security", "//java/src/org/openqa/selenium/grid/web", "//java/src/org/openqa/selenium/json", + "//java/src/org/openqa/selenium/netty/server", "//java/src/org/openqa/selenium/remote", "//java/src/org/openqa/selenium/status", artifact("com.google.guava:guava"), + artifact("io.netty:netty-buffer"), + artifact("io.netty:netty-codec-http"), + artifact("io.netty:netty-transport"), artifact("org.jspecify:jspecify"), ], ) diff --git a/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java b/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java index a2569e67e1ce2..d6b9b34696fcc 100644 --- a/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java +++ b/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java @@ -20,8 +20,19 @@ import static org.openqa.selenium.internal.Debug.getDebugLogLevel; import static org.openqa.selenium.remote.http.HttpMethod.GET; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFutureListener; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelPipeline; +import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame; +import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame; +import io.netty.handler.codec.http.websocketx.TextWebSocketFrame; +import io.netty.handler.codec.http.websocketx.WebSocketFrame; import java.net.URI; import java.net.URISyntaxException; +import java.util.ArrayDeque; +import java.util.Deque; import java.util.Objects; import java.util.Optional; import java.util.Set; @@ -31,18 +42,19 @@ import java.util.logging.Level; import java.util.logging.Logger; import java.util.stream.Stream; +import org.jspecify.annotations.Nullable; import org.openqa.selenium.Capabilities; import org.openqa.selenium.devtools.CdpEndpointFinder; import org.openqa.selenium.grid.data.Session; import org.openqa.selenium.internal.Require; +import org.openqa.selenium.netty.server.PostUpgradeHook; +import org.openqa.selenium.netty.server.WebSocketFrameProxy; import org.openqa.selenium.remote.SessionId; -import org.openqa.selenium.remote.http.BinaryMessage; import org.openqa.selenium.remote.http.ClientConfig; import org.openqa.selenium.remote.http.CloseMessage; import org.openqa.selenium.remote.http.HttpClient; import org.openqa.selenium.remote.http.HttpRequest; import org.openqa.selenium.remote.http.Message; -import org.openqa.selenium.remote.http.TextMessage; import org.openqa.selenium.remote.http.UrlTemplate; import org.openqa.selenium.remote.http.WebSocket; @@ -256,73 +268,13 @@ private Consumer createWsEndPoint( AtomicBoolean upstreamClosing = new AtomicBoolean(false); HttpClient client = clientFactory.createClient(ClientConfig.defaultConfig().baseUri(uri)); + DirectForwardingListener listener = + new DirectForwardingListener( + node, sessionConsumer, sessionId, connectionReleased, client, upstreamClosing); try { - WebSocket upstream = - client.openSocket( - new HttpRequest(GET, uri.toString()), - new ForwardingListener( - node, - downstream, - sessionConsumer, - sessionId, - connectionReleased, - client, - upstreamClosing)); - - return (msg) -> { - // Fast path: once the browser has signalled close, there is no point sending further - // data frames — the JDK WebSocket output is already closing and the send would either - // be dropped or throw "Output closed". For the CloseMessage echo we skip the actual - // network write (the JDK stack handles the protocol-level echo internally when it fires - // onClose) and go straight to resource cleanup. - if (upstreamClosing.get()) { - if (msg instanceof CloseMessage) { - if (connectionReleased.compareAndSet(false, true)) { - node.releaseConnection(sessionId); - try { - client.close(); - } catch (Exception e) { - LOG.log(Level.FINE, "Failed to close client after upstream close for " + uri, e); - } - } - } else { - LOG.log(Level.FINE, "Dropping in-flight data frame for closing session " + sessionId); - } - return; - } - - // Slow path: upstream is (was) open — attempt the send and catch the narrow race where - // the browser closes between the upstreamClosing check above and the actual write. - try { - upstream.send(msg); - } catch (Exception e) { - LOG.log( - Level.FINE, - "Could not forward message to browser WebSocket for session " - + sessionId - + " (connection likely closed concurrently)", - e); - if (connectionReleased.compareAndSet(false, true)) { - node.releaseConnection(sessionId); - try { - client.close(); - } catch (Exception ce) { - LOG.log(Level.FINE, "Failed to close client after send error for " + uri, ce); - } - } - return; - } - if (msg instanceof CloseMessage) { - if (connectionReleased.compareAndSet(false, true)) { - node.releaseConnection(sessionId); - } - try { - client.close(); - } catch (Exception e) { - LOG.log(Level.WARNING, "Failed to shutdown the client of " + uri, e); - } - } - }; + WebSocket upstream = client.openSocket(new HttpRequest(GET, uri.toString()), listener); + return new FrameProxyConsumer( + upstream, client, listener, node, sessionId, connectionReleased, upstreamClosing); } catch (Exception e) { LOG.log(Level.WARNING, "Connecting to upstream websocket failed", e); client.close(); @@ -330,25 +282,147 @@ private Consumer createWsEndPoint( } } - private static class ForwardingListener implements WebSocket.Listener { + // --------------------------------------------------------------------------- + // Consumer returned to WebSocketUpgradeHandler — also implements PostUpgradeHook + // --------------------------------------------------------------------------- + + private static class FrameProxyConsumer implements Consumer, PostUpgradeHook { + + private final WebSocket upstream; + private final HttpClient client; + private final DirectForwardingListener listener; + private final Node node; + private final SessionId sessionId; + private final AtomicBoolean connectionReleased; + private final AtomicBoolean upstreamClosing; + + FrameProxyConsumer( + WebSocket upstream, + HttpClient client, + DirectForwardingListener listener, + Node node, + SessionId sessionId, + AtomicBoolean connectionReleased, + AtomicBoolean upstreamClosing) { + this.upstream = upstream; + this.client = client; + this.listener = listener; + this.node = node; + this.sessionId = sessionId; + this.connectionReleased = connectionReleased; + this.upstreamClosing = upstreamClosing; + } + + /** + * Called on the Netty IO thread once the client-side handshake has completed. Hand the channel + * to the listener (draining any pre-handshake buffer in arrival order), install {@link + * WebSocketFrameProxy} so subsequent inbound frames forward to the upstream {@link WebSocket} + * without going through the {@code Message} layer, and strip the now-redundant Message-layer + * handlers. + */ + @Override + public void onUpgradeComplete(ChannelHandlerContext ctx) { + Channel ch = ctx.channel(); + listener.onUpgrade(ch); + + WebSocketFrameProxy proxy = new WebSocketFrameProxy(upstream, upstreamClosing); + ChannelPipeline pipeline = ctx.pipeline(); + pipeline.addBefore("netty-to-se-messages", "frame-proxy", proxy); + removeIfPresent(pipeline, "netty-to-se-messages"); + removeIfPresent(pipeline, "se-to-netty-messages"); + removeIfPresent(pipeline, "se-websocket-handler"); + } + + private static void removeIfPresent(ChannelPipeline pipeline, String name) { + if (pipeline.get(name) != null) { + pipeline.remove(name); + } + } + + /** + * After pipeline rewiring this consumer is only invoked for {@link CloseMessage} (fired by + * {@code WebSocketUpgradeHandler} on the close handshake or channel-inactive event). Data + * frames are handled directly by {@link WebSocketFrameProxy}. + */ + @Override + public void accept(Message msg) { + if (upstreamClosing.get()) { + if (msg instanceof CloseMessage) { + releaseSlotAndCloseClient(); + } + return; + } + try { + upstream.send(msg); + } catch (Exception e) { + LOG.log( + Level.FINE, + "Could not forward message to browser WebSocket for session " + + sessionId + + " (connection likely closed concurrently)", + e); + releaseSlotAndCloseClient(); + return; + } + if (msg instanceof CloseMessage) { + releaseSlotAndCloseClient(); + } + } + + private void releaseSlotAndCloseClient() { + if (connectionReleased.compareAndSet(false, true)) { + node.releaseConnection(sessionId); + } + try { + client.close(); + } catch (Exception e) { + LOG.log(Level.FINE, "Failed to close client", e); + } + } + } + + // --------------------------------------------------------------------------- + // Listener for browser → client messages (fast path via direct frame writes) + // --------------------------------------------------------------------------- + + /** + * Writes browser-side messages directly to the client {@link Channel} as Netty WebSocket frames, + * bypassing the {@code MessageOutboundConverter}. Frames received before {@link + * #onUpgrade(Channel)} fires are buffered in arrival order and drained on handover, so a frame + * can never land in a pipeline that has already had its Message-layer handlers removed. + */ + static class DirectForwardingListener implements WebSocket.Listener { + + // Bound on the pre-handshake buffer. If the upstream produces more frames than this before + // the client-side handshake completes, the buffer is dropped and a 1009 close is recorded + // so the channel sees a clean close once the upgrade lands. + private static final int MAX_PENDING_FRAMES = 128; + + private final Object lock = new Object(); + private final Deque pending = new ArrayDeque<>(); + private volatile Channel clientChannel; + // Pre-handshake terminal state. When closed is true at onUpgrade time, the listener writes + // the recorded close frame to the channel and tears it down instead of publishing it for + // normal forwarding. Guarded by lock. + private boolean closed; + private int closeCode; + private @Nullable String closeReason; + private final Node node; - private final Consumer downstream; private final Consumer sessionConsumer; private final SessionId sessionId; private final AtomicBoolean connectionReleased; private final HttpClient client; private final AtomicBoolean upstreamClosing; - public ForwardingListener( + DirectForwardingListener( Node node, - Consumer downstream, Consumer sessionConsumer, SessionId sessionId, AtomicBoolean connectionReleased, HttpClient client, AtomicBoolean upstreamClosing) { - this.node = node; - this.downstream = Objects.requireNonNull(downstream); + this.node = Objects.requireNonNull(node); this.sessionConsumer = Objects.requireNonNull(sessionConsumer); this.sessionId = Objects.requireNonNull(sessionId); this.connectionReleased = Objects.requireNonNull(connectionReleased); @@ -356,46 +430,160 @@ public ForwardingListener( this.upstreamClosing = Objects.requireNonNull(upstreamClosing); } + /** + * Hand the client channel over after the WebSocket upgrade has completed. If the upstream + * already closed or errored, write the recorded close frame to the channel and tear it down; + * otherwise drain any frames received in the meantime and publish the channel for the fast + * path. + */ + void onUpgrade(Channel ch) { + boolean terminal; + int code; + String reason; + synchronized (lock) { + WebSocketFrame frame; + while ((frame = pending.pollFirst()) != null) { + ch.writeAndFlush(frame); + } + terminal = closed; + code = closeCode; + reason = closeReason == null ? "" : closeReason; + if (!terminal) { + clientChannel = ch; + } + } + if (terminal && ch.isActive()) { + ch.writeAndFlush(new CloseWebSocketFrame(code, reason)) + .addListener(ChannelFutureListener.CLOSE); + } + } + @Override - public void onBinary(byte[] data) { - downstream.accept(new BinaryMessage(data)); + public void onText(CharSequence data) { + Channel ch = clientChannel; + if (ch != null) { + WebSocketFrameProxy.writeTextFrame(ch, data); + } else { + enqueueOrWrite(new TextWebSocketFrame(data.toString())); + } sessionConsumer.accept(sessionId); } @Override - public void onClose(int code, String reason) { - // Signal the send lambda before forwarding the close downstream so that any data frames - // still queued in the Netty pipeline are discarded rather than attempted on a closing stream. - upstreamClosing.set(true); - downstream.accept(new CloseMessage(code, reason)); - if (connectionReleased.compareAndSet(false, true)) { - node.releaseConnection(sessionId); - // Close the HttpClient eagerly so the connection slot is freed even if the client-side - // Close echo never arrives (e.g. the client dropped the TCP connection). - try { - client.close(); - } catch (Exception e) { - LOG.log(Level.FINE, "Failed to close client on upstream WebSocket close", e); + public void onBinary(byte[] data) { + Channel ch = clientChannel; + if (ch != null) { + WebSocketFrameProxy.writeBinaryFrame(ch, data); + } else { + enqueueOrWrite(new BinaryWebSocketFrame(Unpooled.wrappedBuffer(data))); + } + sessionConsumer.accept(sessionId); + } + + private void enqueueOrWrite(WebSocketFrame frame) { + boolean overflow = false; + Channel ch; + synchronized (lock) { + ch = clientChannel; + if (ch == null) { + if (closed) { + frame.release(); + return; + } + if (pending.size() >= MAX_PENDING_FRAMES) { + // Stall protection: drop the buffer and record a 1009 close so the next onUpgrade + // closes the client cleanly instead of letting memory grow without bound. + frame.release(); + discardPendingLocked(); + closed = true; + closeCode = 1009; + closeReason = "websocket buffer overflow"; + overflow = true; + } else { + pending.addLast(frame); + return; + } } } + if (overflow) { + upstreamClosing.set(true); + LOG.log( + Level.WARNING, + "Dropping pre-handshake WebSocket buffer for session {0}: exceeded {1} pending" + + " frames", + new Object[] {sessionId, MAX_PENDING_FRAMES}); + releaseSlot(); + return; + } + ch.writeAndFlush(frame); } @Override - public void onText(CharSequence data) { - downstream.accept(new TextMessage(data)); - sessionConsumer.accept(sessionId); + public void onClose(int code, String reason) { + // Signal closing before forwarding the close so any data frames already queued for this + // listener are discarded rather than attempted on a closing stream. + upstreamClosing.set(true); + // The WebSocket spec caps close-frame reasons at 123 bytes UTF-8; truncate so an upstream + // sending a longer one cannot break the close frame's encoding. + String safeReason = WebSocketFrameProxy.truncateCloseReason(reason); + Channel ch = clientChannel; + if (ch == null) { + synchronized (lock) { + ch = clientChannel; + if (ch == null) { + // Pre-handshake close: drop the buffer so ref-counted frames are released even if + // the client-side handshake never lands, and record the code/reason so a late + // onUpgrade can still surface a clean close to the client. + discardPendingLocked(); + closed = true; + closeCode = code; + closeReason = safeReason; + } + } + } + if (ch != null && ch.isActive()) { + ch.writeAndFlush(new CloseWebSocketFrame(code, safeReason)); + } + releaseSlot(); } @Override public void onError(Throwable cause) { upstreamClosing.set(true); LOG.log(Level.WARNING, "Error proxying websocket command", cause); + Channel ch = clientChannel; + if (ch == null) { + synchronized (lock) { + ch = clientChannel; + if (ch == null) { + discardPendingLocked(); + closed = true; + closeCode = 1011; + closeReason = "upstream error"; + } + } + } + if (ch != null && ch.isActive()) { + ch.writeAndFlush(new CloseWebSocketFrame(1011, "upstream error")) + .addListener(ChannelFutureListener.CLOSE); + } + releaseSlot(); + } + + private void discardPendingLocked() { + WebSocketFrame frame; + while ((frame = pending.pollFirst()) != null) { + frame.release(); + } + } + + private void releaseSlot() { if (connectionReleased.compareAndSet(false, true)) { node.releaseConnection(sessionId); try { client.close(); } catch (Exception e) { - LOG.log(Level.FINE, "Failed to close client after WebSocket error", e); + LOG.log(Level.FINE, "Failed to close client", e); } } } diff --git a/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java b/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java index e63ea28b0b855..2f365d0491fe8 100644 --- a/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java +++ b/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java @@ -332,6 +332,9 @@ private void enqueueOrWrite(WebSocketFrame frame) { @Override public void onClose(int code, String reason) { upstreamClosing.set(true); + // The WebSocket spec caps close-frame reasons at 123 bytes UTF-8; truncate so an upstream + // sending a longer one cannot break the close frame's encoding. + String safeReason = WebSocketFrameProxy.truncateCloseReason(reason); Channel ch = clientChannel; if (ch == null) { synchronized (lock) { @@ -344,12 +347,12 @@ public void onClose(int code, String reason) { discardPendingLocked(); closed = true; closeCode = code; - closeReason = reason; + closeReason = safeReason; } } } if (ch != null && ch.isActive()) { - ch.writeAndFlush(new CloseWebSocketFrame(code, reason)); + ch.writeAndFlush(new CloseWebSocketFrame(code, safeReason)); } closeClient(); } diff --git a/java/src/org/openqa/selenium/netty/server/WebSocketFrameProxy.java b/java/src/org/openqa/selenium/netty/server/WebSocketFrameProxy.java index 40f3d9c622531..908888056cb3e 100644 --- a/java/src/org/openqa/selenium/netty/server/WebSocketFrameProxy.java +++ b/java/src/org/openqa/selenium/netty/server/WebSocketFrameProxy.java @@ -17,6 +17,8 @@ package org.openqa.selenium.netty.server; +import static java.nio.charset.StandardCharsets.UTF_8; + import io.netty.buffer.Unpooled; import io.netty.channel.Channel; import io.netty.channel.ChannelHandlerContext; @@ -28,6 +30,9 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.CharBuffer; +import java.nio.charset.CharsetEncoder; import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.Level; import java.util.logging.Logger; @@ -167,4 +172,31 @@ public static void writeTextFrame(Channel clientChannel, CharSequence text) { public static void writeBinaryFrame(Channel clientChannel, byte[] data) { clientChannel.writeAndFlush(new BinaryWebSocketFrame(Unpooled.wrappedBuffer(data))); } + + /** + * Shrink a WebSocket close-frame reason to fit RFC 6455 §5.5.1's 123-byte UTF-8 cap. + * + *

A naïve approach — encode to bytes, truncate, decode back — can split a multi-byte UTF-8 + * sequence at the boundary, which Java then decodes as {@code U+FFFD} (three bytes when + * re-encoded). That pushes the final encoded length back over the limit and breaks {@link + * io.netty.handler.codec.http.websocketx.CloseWebSocketFrame} encoding. Use {@link + * CharsetEncoder} into a 120-byte buffer instead — {@code encode()} stops at a clean character + * boundary on overflow, so no partial sequence is ever left behind. A three-byte ASCII ellipsis + * marks the truncation, keeping the encoded total at most 123 bytes regardless of the input. + */ + public static String truncateCloseReason(String reason) { + if (reason == null) { + return ""; + } + // Fast path: short reasons skip the encoder allocation. + if (reason.getBytes(UTF_8).length <= 123) { + return reason; + } + ByteBuffer out = ByteBuffer.allocate(120); + CharsetEncoder encoder = UTF_8.newEncoder(); + encoder.encode(CharBuffer.wrap(reason), out, true); + encoder.flush(out); + out.flip(); + return UTF_8.decode(out) + "..."; + } } diff --git a/java/test/org/openqa/selenium/grid/node/BUILD.bazel b/java/test/org/openqa/selenium/grid/node/BUILD.bazel index 1d3ab581d282c..c1250f8948add 100644 --- a/java/test/org/openqa/selenium/grid/node/BUILD.bazel +++ b/java/test/org/openqa/selenium/grid/node/BUILD.bazel @@ -25,6 +25,8 @@ java_test_suite( "//java/test/org/openqa/selenium/grid/testing", "//java/test/org/openqa/selenium/remote/tracing:tracing-support", artifact("com.google.guava:guava"), + artifact("io.netty:netty-codec-http"), + artifact("io.netty:netty-transport"), artifact("io.opentelemetry:opentelemetry-api"), artifact("org.junit.jupiter:junit-jupiter-api"), artifact("org.assertj:assertj-core"), diff --git a/java/test/org/openqa/selenium/grid/node/NodeDirectForwardingListenerTest.java b/java/test/org/openqa/selenium/grid/node/NodeDirectForwardingListenerTest.java new file mode 100644 index 0000000000000..c5967c470d02f --- /dev/null +++ b/java/test/org/openqa/selenium/grid/node/NodeDirectForwardingListenerTest.java @@ -0,0 +1,298 @@ +// Licensed to the Software Freedom Conservancy (SFC) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The SFC licenses this file +// to you 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 org.openqa.selenium.grid.node; + +import static org.assertj.core.api.Assertions.assertThat; + +import io.netty.channel.embedded.EmbeddedChannel; +import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame; +import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame; +import io.netty.handler.codec.http.websocketx.TextWebSocketFrame; +import java.net.URI; +import java.time.Duration; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import org.openqa.selenium.Capabilities; +import org.openqa.selenium.NoSuchSessionException; +import org.openqa.selenium.WebDriverException; +import org.openqa.selenium.grid.data.CreateSessionRequest; +import org.openqa.selenium.grid.data.CreateSessionResponse; +import org.openqa.selenium.grid.data.NodeId; +import org.openqa.selenium.grid.data.NodeStatus; +import org.openqa.selenium.grid.data.Session; +import org.openqa.selenium.grid.security.Secret; +import org.openqa.selenium.internal.Either; +import org.openqa.selenium.remote.SessionId; +import org.openqa.selenium.remote.http.HttpClient; +import org.openqa.selenium.remote.http.HttpRequest; +import org.openqa.selenium.remote.http.HttpResponse; +import org.openqa.selenium.remote.http.WebSocket; +import org.openqa.selenium.remote.tracing.DefaultTestTracer; + +class NodeDirectForwardingListenerTest { + + private static final Secret SECRET = new Secret("test"); + + @Test + void framesArrivingBeforeHandshakeAreBufferedAndDrainInOrder() { + SessionId sessionId = new SessionId(UUID.randomUUID()); + AtomicInteger heartbeats = new AtomicInteger(); + + ProxyNodeWebsockets.DirectForwardingListener listener = + new ProxyNodeWebsockets.DirectForwardingListener( + new NoopNode(), + id -> heartbeats.incrementAndGet(), + sessionId, + new AtomicBoolean(false), + new NoopHttpClient(), + new AtomicBoolean(false)); + + listener.onText("first"); + listener.onBinary(new byte[] {7, 8, 9}); + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + Object first = channel.readOutbound(); + Object second = channel.readOutbound(); + + assertThat(first).isInstanceOf(TextWebSocketFrame.class); + assertThat(((TextWebSocketFrame) first).text()).isEqualTo("first"); + assertThat(second).isInstanceOf(BinaryWebSocketFrame.class); + assertThat(((BinaryWebSocketFrame) second).content().readableBytes()).isEqualTo(3); + + // Heartbeats fire per frame regardless of which side of the handshake we're on. + assertThat(heartbeats.get()).isEqualTo(2); + + ((TextWebSocketFrame) first).release(); + ((BinaryWebSocketFrame) second).release(); + } + + @Test + void preHandshakeCloseIsSurfacedAndChannelIsTornDown() { + AtomicInteger releaseCount = new AtomicInteger(); + ProxyNodeWebsockets.DirectForwardingListener listener = + new ProxyNodeWebsockets.DirectForwardingListener( + countingReleaseNode(releaseCount), + id -> {}, + new SessionId(UUID.randomUUID()), + new AtomicBoolean(false), + new NoopHttpClient(), + new AtomicBoolean(false)); + + // Upstream queued a frame and then closed, all before the client-side handshake landed. + listener.onText("buffered"); + listener.onClose(4001, "browser gone"); + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + // The buffered text is gone (released on close so the ref-counted buffer cannot leak when + // the handshake never completes), and the client sees a proper close frame followed by a + // channel teardown. + CloseWebSocketFrame close = channel.readOutbound(); + assertThat(close.statusCode()).isEqualTo(4001); + assertThat(close.reasonText()).isEqualTo("browser gone"); + close.release(); + assertThat((Object) channel.readOutbound()).isNull(); + assertThat(channel.isOpen()).isFalse(); + assertThat(releaseCount.get()).isEqualTo(1); + } + + @Test + void closeReasonIsTruncatedSafelyEvenWhenItContainsMultiByteCharacters() { + ProxyNodeWebsockets.DirectForwardingListener listener = + new ProxyNodeWebsockets.DirectForwardingListener( + new NoopNode(), + id -> {}, + new SessionId(UUID.randomUUID()), + new AtomicBoolean(false), + new NoopHttpClient(), + new AtomicBoolean(false)); + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + // RFC 6455 §5.5.1 caps close-frame reasons at 123 bytes UTF-8. Build a 200-byte reason out + // of a two-byte UTF-8 character ('é') so that any cut at an arbitrary byte offset would + // split a codepoint and produce a U+FFFD replacement on decode — which re-encodes to three + // bytes and would push the result back over the limit. + StringBuilder tooLong = new StringBuilder(); + for (int i = 0; i < 100; i++) { + tooLong.append('é'); + } + listener.onClose(4001, tooLong.toString()); + + CloseWebSocketFrame close = channel.readOutbound(); + assertThat(close.reasonText().getBytes(java.nio.charset.StandardCharsets.UTF_8)) + .hasSizeLessThanOrEqualTo(123); + assertThat(close.reasonText()).doesNotContain("�"); + close.release(); + } + + @Test + void preHandshakeBufferOverflowRecordsCleanCloseAndReleasesSlot() { + AtomicInteger releaseCount = new AtomicInteger(); + ProxyNodeWebsockets.DirectForwardingListener listener = + new ProxyNodeWebsockets.DirectForwardingListener( + countingReleaseNode(releaseCount), + id -> {}, + new SessionId(UUID.randomUUID()), + new AtomicBoolean(false), + new NoopHttpClient(), + new AtomicBoolean(false)); + + // Cap is 128 frames; 200 frames is comfortably past it. + for (int i = 0; i < 200; i++) { + listener.onBinary(new byte[] {1}); + } + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + CloseWebSocketFrame close = channel.readOutbound(); + assertThat(close.statusCode()).isEqualTo(1009); + close.release(); + assertThat((Object) channel.readOutbound()).isNull(); + assertThat(channel.isOpen()).isFalse(); + + // Overflow must release the connection slot so the limit counter does not leak. + assertThat(releaseCount.get()).isEqualTo(1); + } + + private static Node countingReleaseNode(AtomicInteger releaseCount) { + return new NoopNode() { + @Override + public void releaseConnection(SessionId id) { + releaseCount.incrementAndGet(); + } + }; + } + + // ------------------------------------------------------------------ + // Stubs + // ------------------------------------------------------------------ + + private static final class NoopHttpClient implements HttpClient { + @Override + public HttpResponse execute(HttpRequest request) { + throw new UnsupportedOperationException(); + } + + @Override + public WebSocket openSocket(HttpRequest request, WebSocket.Listener listener) { + throw new UnsupportedOperationException(); + } + + @Override + public + java.util.concurrent.CompletableFuture> sendAsyncNative( + java.net.http.HttpRequest request, java.net.http.HttpResponse.BodyHandler handler) { + throw new UnsupportedOperationException(); + } + + @Override + public java.net.http.HttpResponse sendNative( + java.net.http.HttpRequest request, java.net.http.HttpResponse.BodyHandler handler) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + } + + private static class NoopNode extends Node { + NoopNode() { + super( + DefaultTestTracer.createTracer(), + new NodeId(UUID.randomUUID()), + URI.create("http://localhost:5555"), + SECRET, + Duration.ofSeconds(30)); + } + + @Override + public Either newSession( + CreateSessionRequest sessionRequest) { + throw new UnsupportedOperationException(); + } + + @Override + public HttpResponse executeWebDriverCommand(HttpRequest req) { + throw new UnsupportedOperationException(); + } + + @Override + public Session getSession(SessionId id) { + throw new UnsupportedOperationException(); + } + + @Override + public HttpResponse uploadFile(HttpRequest req, SessionId id) { + throw new UnsupportedOperationException(); + } + + @Override + public HttpResponse downloadFile(HttpRequest req, SessionId id) { + throw new UnsupportedOperationException(); + } + + @Override + public void stop(SessionId id) throws NoSuchSessionException { + throw new UnsupportedOperationException(); + } + + @Override + public boolean isSessionOwner(SessionId id) { + return true; + } + + @Override + public boolean isSupporting(Capabilities capabilities) { + return false; + } + + @Override + public NodeStatus getStatus() { + throw new UnsupportedOperationException(); + } + + @Override + public HealthCheck getHealthCheck() { + throw new UnsupportedOperationException(); + } + + @Override + public void drain() {} + + @Override + public boolean isReady() { + return true; + } + + @Override + public boolean tryAcquireConnection(SessionId id) { + return true; + } + + @Override + public void releaseConnection(SessionId id) {} + } +} diff --git a/java/test/org/openqa/selenium/netty/server/WebSocketFrameProxyTest.java b/java/test/org/openqa/selenium/netty/server/WebSocketFrameProxyTest.java index ce7e404a865cb..267e4a315606a 100644 --- a/java/test/org/openqa/selenium/netty/server/WebSocketFrameProxyTest.java +++ b/java/test/org/openqa/selenium/netty/server/WebSocketFrameProxyTest.java @@ -17,6 +17,7 @@ package org.openqa.selenium.netty.server; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; import io.netty.channel.embedded.EmbeddedChannel; @@ -54,4 +55,28 @@ public void close() {} // First failure should latch upstreamClosing so subsequent frames short-circuit. assertThat(upstreamClosing.get()).isTrue(); } + + @Test + void truncateCloseReasonReturnsShortReasonUnchanged() { + assertThat(WebSocketFrameProxy.truncateCloseReason("upstream gone")).isEqualTo("upstream gone"); + assertThat(WebSocketFrameProxy.truncateCloseReason("")).isEqualTo(""); + assertThat(WebSocketFrameProxy.truncateCloseReason(null)).isEqualTo(""); + } + + @Test + void truncateCloseReasonStaysWithinTheWebSocketByteLimit() { + // RFC 6455 §5.5.1 caps close-frame reasons at 123 bytes UTF-8. Build a 200-byte reason out + // of a two-byte UTF-8 character ('é') so the truncation is forced to cut where a naïve + // byte-level approach would split a codepoint and produce a U+FFFD replacement. + StringBuilder tooLong = new StringBuilder(); + for (int i = 0; i < 100; i++) { + tooLong.append('é'); + } + + String truncated = WebSocketFrameProxy.truncateCloseReason(tooLong.toString()); + + assertThat(truncated.getBytes(UTF_8)).hasSizeLessThanOrEqualTo(123); + assertThat(truncated).doesNotContain("�"); + assertThat(truncated).endsWith("..."); + } }