[grid] Close pre-handshake race in WebSocket proxy

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"),

java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java

@@ -19,34 +19,38 @@
 
 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;
 import org.openqa.selenium.netty.server.PostUpgradeHook;
 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<Consumer<Message>> apply(String uri, Consumer<Message> downstrea
     }
 
     AtomicBoolean upstreamClosing = new AtomicBoolean(false);
-    // Holds the client channel once onUpgradeComplete fires; used by DirectForwardingListener
-    // to write frames without going through MessageOutboundConverter.
-    AtomicReference<Channel> 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<Message>, PostUpgrad
 
     private final WebSocket upstream;
     private final HttpClient client;
-    private final AtomicReference<Channel> clientChannelRef;
+    private final DirectForwardingListener listener;
     private final AtomicBoolean upstreamClosing;
 
     FrameProxyConsumer(
         WebSocket upstream,
         HttpClient client,
-        AtomicReference<Channel> 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}.
+   *
+   * <p>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<WebSocketFrame> 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<Message> fallbackDownstream;
-    private final AtomicReference<Channel> clientChannelRef;
     private final AtomicBoolean upstreamClosing;
     private final HttpClient client;
 
-    DirectForwardingListener(
-        Consumer<Message> fallbackDownstream,
-        AtomicReference<Channel> 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);
       }
     }
   }

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",