diff --git a/java/src/org/openqa/selenium/grid/router/BUILD.bazel b/java/src/org/openqa/selenium/grid/router/BUILD.bazel index d3d7ca5acf2ab..6888c76e99e72 100644 --- a/java/src/org/openqa/selenium/grid/router/BUILD.bazel +++ b/java/src/org/openqa/selenium/grid/router/BUILD.bazel @@ -25,6 +25,7 @@ java_library( "//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/router/ProxyWebsocketsIntoGrid.java b/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java index 9fb1d4a91e7cc..e63ea28b0b855 100644 --- a/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java +++ b/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java @@ -19,20 +19,26 @@ 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.util.ArrayDeque; +import java.util.Deque; import java.util.Objects; import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Consumer; import java.util.logging.Level; import java.util.logging.Logger; +import org.jspecify.annotations.Nullable; import org.openqa.selenium.NoSuchSessionException; import org.openqa.selenium.grid.sessionmap.SessionMap; import org.openqa.selenium.internal.Require; @@ -40,13 +46,11 @@ import org.openqa.selenium.netty.server.WebSocketFrameProxy; import org.openqa.selenium.remote.HttpSessionId; 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.WebSocket; /** @@ -91,20 +95,14 @@ public Optional> apply(String uri, Consumer downstrea } AtomicBoolean upstreamClosing = new AtomicBoolean(false); - // Holds the client channel once onUpgradeComplete fires; used by DirectForwardingListener - // to write frames without going through MessageOutboundConverter. - AtomicReference clientChannelRef = new AtomicReference<>(); HttpClient client = clientFactory.createClient(ClientConfig.defaultConfig().baseUri(sessionUri)); + DirectForwardingListener listener = new DirectForwardingListener(upstreamClosing, client); try { - WebSocket upstream = - client.openSocket( - new HttpRequest(GET, uri), - new DirectForwardingListener(downstream, clientChannelRef, upstreamClosing, client)); + WebSocket upstream = client.openSocket(new HttpRequest(GET, uri), listener); - return Optional.of( - new FrameProxyConsumer(upstream, client, clientChannelRef, upstreamClosing)); + return Optional.of(new FrameProxyConsumer(upstream, client, listener, upstreamClosing)); } catch (Exception e) { LOG.log(Level.WARNING, "Connecting to upstream websocket failed", e); @@ -121,30 +119,32 @@ private static class FrameProxyConsumer implements Consumer, PostUpgrad private final WebSocket upstream; private final HttpClient client; - private final AtomicReference clientChannelRef; + private final DirectForwardingListener listener; private final AtomicBoolean upstreamClosing; FrameProxyConsumer( WebSocket upstream, HttpClient client, - AtomicReference clientChannelRef, + DirectForwardingListener listener, AtomicBoolean upstreamClosing) { this.upstream = upstream; this.client = client; - this.clientChannelRef = clientChannelRef; + this.listener = listener; this.upstreamClosing = upstreamClosing; } /** * Called by {@link org.openqa.selenium.netty.server.WebSocketUpgradeHandler} on the Netty IO - * thread after the client-side handshake completes. Install {@link WebSocketFrameProxy} and - * strip the three {@code Message}-layer handlers so subsequent data frames never pass through - * the full Selenium handler chain. + * thread after the client-side handshake completes. Hand the channel to the listener (which + * drains any frames received before the handshake landed), then install {@link + * WebSocketFrameProxy} and strip the three {@code Message}-layer handlers so subsequent data + * frames never pass through the full Selenium handler chain. */ @Override public void onUpgradeComplete(ChannelHandlerContext ctx) { Channel ch = ctx.channel(); - clientChannelRef.set(ch); + // Drain any pre-handshake buffer in order before the listener starts taking the fast path. + listener.onUpgrade(ch); WebSocketFrameProxy proxy = new WebSocketFrameProxy(upstream, upstreamClosing); ChannelPipeline pipeline = ctx.pipeline(); @@ -208,84 +208,190 @@ private void closeClient() { /** * Writes node-side messages directly to the client {@link Channel} as Netty WebSocket frames, - * bypassing {@code MessageOutboundConverter}. Falls back to the {@code downstream} consumer - * before the client channel reference is set (i.e. before {@link - * PostUpgradeHook#onUpgradeComplete} fires, which is rare). + * bypassing {@code MessageOutboundConverter}. + * + *

Frames received from the upstream before {@link #onUpgrade(Channel)} fires are buffered in + * arrival order; the buffer is then drained on the Netty IO thread before any subsequent listener + * call takes the fast path. This makes the pre-handshake → post-handshake transition + * deterministic: a frame can never land in a pipeline that has already had its Message-layer + * handlers removed. */ - private static class DirectForwardingListener implements WebSocket.Listener { + 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 (a stalled or hostile client), the buffer is dropped + // and the listener latches a 1009 ("Message Too Big" / overflow) terminal state 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<>(); + // Volatile so the post-handover fast path needs no synchronization. + private volatile Channel clientChannel; + // Pre-handshake terminal state. When `closed` is true at onUpgrade time, the listener + // surfaces the recorded close to the client instead of publishing the channel for normal + // forwarding — this stops the downstream from sitting open indefinitely when the upstream + // has already gone away. Guarded by lock. + private boolean closed; + private int closeCode; + private @Nullable String closeReason; - private final Consumer fallbackDownstream; - private final AtomicReference clientChannelRef; private final AtomicBoolean upstreamClosing; private final HttpClient client; - DirectForwardingListener( - Consumer fallbackDownstream, - AtomicReference clientChannelRef, - AtomicBoolean upstreamClosing, - HttpClient client) { - this.fallbackDownstream = Objects.requireNonNull(fallbackDownstream); - this.clientChannelRef = Objects.requireNonNull(clientChannelRef); + DirectForwardingListener(AtomicBoolean upstreamClosing, HttpClient client) { this.upstreamClosing = Objects.requireNonNull(upstreamClosing); this.client = Objects.requireNonNull(client); } + /** + * Hand the client channel over after the WebSocket upgrade has completed. If the upstream + * already closed or errored, surface that to the client now and close the channel; otherwise + * drain any frames received in the meantime in arrival order 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 onText(CharSequence data) { - Channel ch = clientChannelRef.get(); + Channel ch = clientChannel; if (ch != null) { - // Fast path: write TextWebSocketFrame directly, skipping MessageOutboundConverter. WebSocketFrameProxy.writeTextFrame(ch, data); - } else { - fallbackDownstream.accept(new TextMessage(data)); + return; } + enqueueOrWrite(new TextWebSocketFrame(data.toString())); } @Override public void onBinary(byte[] data) { - Channel ch = clientChannelRef.get(); + Channel ch = clientChannel; if (ch != null) { - // Fast path: write BinaryWebSocketFrame directly, skipping MessageOutboundConverter. WebSocketFrameProxy.writeBinaryFrame(ch, data); - } else { - fallbackDownstream.accept(new BinaryMessage(data)); + return; } + enqueueOrWrite(new BinaryWebSocketFrame(Unpooled.wrappedBuffer(data))); + } + + 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 arm a terminal state so the next + // onUpgrade closes the client cleanly instead of letting memory grow unbounded. + 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 {0}: exceeded {1} pending frames", + new Object[] {client, MAX_PENDING_FRAMES}); + closeClient(); + return; + } + ch.writeAndFlush(frame); } @Override public void onClose(int code, String reason) { upstreamClosing.set(true); - // After onUpgradeComplete the pipeline no longer contains MessageOutboundConverter, so - // writing a CloseMessage object via fallbackDownstream would fail to encode. Write the - // Netty frame directly once the client channel reference is available. - Channel ch = clientChannelRef.get(); + 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 + // immediately even if the client-side handshake never lands (for instance because + // the client disconnected mid-handshake or the upgrade itself failed). Record the + // close so onUpgrade can still surface it to the client if the handshake does fire. + discardPendingLocked(); + closed = true; + closeCode = code; + closeReason = reason; + } + } + } if (ch != null && ch.isActive()) { ch.writeAndFlush(new CloseWebSocketFrame(code, reason)); - } else { - fallbackDownstream.accept(new CloseMessage(code, reason)); - } - try { - client.close(); - } catch (Exception e) { - LOG.log(Level.FINE, "Failed to close client on upstream WebSocket close", e); } + closeClient(); } @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) { + // Pre-handshake error: drop the buffer (frames may not be coherent) and record 1011. + discardPendingLocked(); + closed = true; + closeCode = 1011; + closeReason = "upstream error"; + } + } + } // Close the client channel so Playwright/BiDi clients see a clean disconnect rather than // hanging until the next keepalive ping fires. - Channel ch = clientChannelRef.get(); if (ch != null && ch.isActive()) { ch.writeAndFlush(new CloseWebSocketFrame(1011, "upstream error")) .addListener(ChannelFutureListener.CLOSE); } + closeClient(); + } + + private void discardPendingLocked() { + WebSocketFrame frame; + while ((frame = pending.pollFirst()) != null) { + frame.release(); + } + } + + private void closeClient() { 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/test/org/openqa/selenium/grid/router/BUILD.bazel b/java/test/org/openqa/selenium/grid/router/BUILD.bazel index 0960a650e59ff..5f0ddc0a61dc7 100644 --- a/java/test/org/openqa/selenium/grid/router/BUILD.bazel +++ b/java/test/org/openqa/selenium/grid/router/BUILD.bazel @@ -17,6 +17,10 @@ BIDI_TESTS = [ "RemoteWebDriverBiDiTest.java", ] +SMALL_TESTS = [ + "DirectForwardingListenerTest.java", +] + java_library( name = "support", testonly = True, @@ -128,12 +132,27 @@ java_selenium_test_suite( ] + CDP_DEPS + JUNIT5_DEPS, ) +java_test_suite( + name = "small-tests", + size = "small", + srcs = SMALL_TESTS, + deps = [ + "//java/src/org/openqa/selenium/grid/router", + "//java/src/org/openqa/selenium/remote/http", + artifact("io.netty:netty-codec-http"), + artifact("io.netty:netty-transport"), + artifact("org.junit.jupiter:junit-jupiter-api"), + artifact("org.assertj:assertj-core"), + artifact("org.jspecify:jspecify"), + ] + JUNIT5_DEPS, +) + java_test_suite( name = "medium-tests", size = "medium", srcs = glob( ["*Test.java"], - exclude = LARGE_TESTS + BIDI_TESTS + STRESS_TESTS, + exclude = LARGE_TESTS + BIDI_TESTS + STRESS_TESTS + SMALL_TESTS, ), tags = [ "requires-network", diff --git a/java/test/org/openqa/selenium/grid/router/DirectForwardingListenerTest.java b/java/test/org/openqa/selenium/grid/router/DirectForwardingListenerTest.java new file mode 100644 index 0000000000000..1c1eb5a91d070 --- /dev/null +++ b/java/test/org/openqa/selenium/grid/router/DirectForwardingListenerTest.java @@ -0,0 +1,157 @@ +// 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.router; + +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.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.Test; +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; + +class DirectForwardingListenerTest { + + @Test + void framesReceivedBeforeUpgradeAreDrainedInOrder() { + HttpClient noopClient = new NoopHttpClient(); + ProxyWebsocketsIntoGrid.DirectForwardingListener listener = + new ProxyWebsocketsIntoGrid.DirectForwardingListener(new AtomicBoolean(false), noopClient); + + // Two frames arrive from the upstream before the WebSocket upgrade has completed. + listener.onText("first"); + listener.onBinary(new byte[] {1, 2, 3}); + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + Object firstWritten = channel.readOutbound(); + Object secondWritten = channel.readOutbound(); + Object nothingMore = channel.readOutbound(); + + assertThat(firstWritten).isInstanceOf(TextWebSocketFrame.class); + assertThat(((TextWebSocketFrame) firstWritten).text()).isEqualTo("first"); + assertThat(secondWritten).isInstanceOf(BinaryWebSocketFrame.class); + assertThat(((BinaryWebSocketFrame) secondWritten).content().readableBytes()).isEqualTo(3); + assertThat(nothingMore).isNull(); + + ((TextWebSocketFrame) firstWritten).release(); + ((BinaryWebSocketFrame) secondWritten).release(); + } + + @Test + void preHandshakeCloseSurfacesCloseAndReleasesBufferedFrames() { + ProxyWebsocketsIntoGrid.DirectForwardingListener listener = + new ProxyWebsocketsIntoGrid.DirectForwardingListener( + new AtomicBoolean(false), new NoopHttpClient()); + + // Frames arrive, then the upstream closes — all before the client-side handshake landed. + listener.onText("buffered"); + listener.onBinary(new byte[] {1, 2, 3}); + listener.onClose(4001, "upstream gone"); + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + // The only thing on the wire is the close frame: pending was released on the close so the + // ref-counted buffers do not leak when the client-side handshake never completes. The + // close + channel teardown also runs, so the client doesn't sit open indefinitely. + CloseWebSocketFrame close = channel.readOutbound(); + assertThat(close.statusCode()).isEqualTo(4001); + assertThat(close.reasonText()).isEqualTo("upstream gone"); + close.release(); + assertThat((Object) channel.readOutbound()).isNull(); + assertThat(channel.isOpen()).isFalse(); + } + + @Test + void overflowOfPreHandshakeBufferArmsCloseAndDrains() { + ProxyWebsocketsIntoGrid.DirectForwardingListener listener = + new ProxyWebsocketsIntoGrid.DirectForwardingListener( + new AtomicBoolean(false), new NoopHttpClient()); + + // Cap is 128; sending 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); + + // Buffer was discarded on overflow; the only thing on the wire is the 1009 close. + Object out = channel.readOutbound(); + assertThat(out).isInstanceOf(CloseWebSocketFrame.class); + CloseWebSocketFrame close = (CloseWebSocketFrame) out; + assertThat(close.statusCode()).isEqualTo(1009); + close.release(); + + assertThat((Object) channel.readOutbound()).isNull(); + assertThat(channel.isOpen()).isFalse(); + } + + @Test + void postUpgradeFramesGoStraightToTheChannel() { + ProxyWebsocketsIntoGrid.DirectForwardingListener listener = + new ProxyWebsocketsIntoGrid.DirectForwardingListener( + new AtomicBoolean(false), new NoopHttpClient()); + + EmbeddedChannel channel = new EmbeddedChannel(); + listener.onUpgrade(channel); + + listener.onText("after-handshake"); + + Object written = channel.readOutbound(); + assertThat(written).isInstanceOf(TextWebSocketFrame.class); + assertThat(((TextWebSocketFrame) written).text()).isEqualTo("after-handshake"); + ((TextWebSocketFrame) written).release(); + } + + /** Stand-in HttpClient so we don't pull a network factory into a unit test. */ + 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() {} + } +}