From 2cf4412e04c044825e229258cb37587e3ff9e100 Mon Sep 17 00:00:00 2001
From: Julien Viet <julien@julienviet.com>
Date: Thu, 24 Oct 2024 15:37:39 +0200
Subject: [PATCH] Port ClientWebSocketInternal from 4.x

---
 .../core/http/impl/ClientWebSocketImpl.java   | 36 +++++++++++++------
 .../http/ClientWebSocketInternal.java         | 23 ++++++++++++
 2 files changed, 49 insertions(+), 10 deletions(-)
 create mode 100644 vertx-core/src/main/java/io/vertx/core/internal/http/ClientWebSocketInternal.java

diff --git a/vertx-core/src/main/java/io/vertx/core/http/impl/ClientWebSocketImpl.java b/vertx-core/src/main/java/io/vertx/core/http/impl/ClientWebSocketImpl.java
index cfa845fbde8..8d9a3a17b2d 100644
--- a/vertx-core/src/main/java/io/vertx/core/http/impl/ClientWebSocketImpl.java
+++ b/vertx-core/src/main/java/io/vertx/core/http/impl/ClientWebSocketImpl.java
@@ -10,11 +10,14 @@
  */
 package io.vertx.core.http.impl;
 
+import io.netty.channel.ChannelHandlerContext;
 import io.vertx.codegen.annotations.Nullable;
 import io.vertx.core.*;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.*;
 import io.vertx.core.internal.ContextInternal;
+import io.vertx.core.internal.http.ClientWebSocketInternal;
+import io.vertx.core.internal.http.WebSocketInternal;
 import io.vertx.core.net.SocketAddress;
 
 import javax.net.ssl.SSLPeerUnverifiedException;
@@ -27,11 +30,11 @@
 /**
  * Client WebSocket implementation
  */
-public class ClientWebSocketImpl implements ClientWebSocket {
+public class ClientWebSocketImpl implements ClientWebSocketInternal {
 
-  private WebSocketClientImpl client;
+  private final WebSocketClientImpl client;
   private final AtomicReference<Promise<WebSocket>> connect = new AtomicReference<>();
-  private volatile WebSocket ws;
+  private volatile WebSocketInternal ws;
   private Handler<Throwable> exceptionHandler;
   private Handler<Buffer> dataHandler;
   private Handler<Void> endHandler;
@@ -49,17 +52,25 @@ public class ClientWebSocketImpl implements ClientWebSocket {
 
   @Override
   public Future<WebSocket> connect(WebSocketConnectOptions options) {
-    ContextInternal ctx = client.vertx().getOrCreateContext();
-    Promise<WebSocket> promise = ctx.promise();
+    return connect(client.vertx().getOrCreateContext(), options);
+  }
+
+  @Override
+  public Future<WebSocket> connect(Context context, WebSocketConnectOptions options) {
+    return connect((ContextInternal) context, options);
+  }
+
+  private Future<WebSocket> connect(ContextInternal context, WebSocketConnectOptions options) {
+    Promise<WebSocket> promise = context.promise();
     if (!connect.compareAndSet(null, promise)) {
-      return ctx.failedFuture("Already connecting");
+      return context.failedFuture("Already connecting");
     }
-    client.webSocket(ctx, options, promise);
+    client.webSocket(context, options, promise);
     return promise
       .future()
       .andThen(ar -> {
         if (ar.succeeded()) {
-          WebSocket w = ar.result();
+          WebSocketInternal w = (WebSocketInternal) ar.result();
           ws = w;
           w.handler(dataHandler);
           w.binaryMessageHandler(binaryMessageHandler);
@@ -118,6 +129,11 @@ public ClientWebSocket endHandler(Handler<Void> handler) {
     return this;
   }
 
+  @Override
+  public ChannelHandlerContext channelHandlerContext() {
+    return delegate().channelHandlerContext();
+  }
+
   @Override
   public ClientWebSocket setWriteQueueMaxSize(int maxSize) {
     delegate().setWriteQueueMaxSize(maxSize);
@@ -309,8 +325,8 @@ public boolean writeQueueFull() {
     return delegate().writeQueueFull();
   }
 
-  private WebSocket delegate() {
-    WebSocket w = ws;
+  private WebSocketInternal delegate() {
+    WebSocketInternal w = ws;
     if (w == null) {
       throw new IllegalStateException("Not connected");
     }
diff --git a/vertx-core/src/main/java/io/vertx/core/internal/http/ClientWebSocketInternal.java b/vertx-core/src/main/java/io/vertx/core/internal/http/ClientWebSocketInternal.java
new file mode 100644
index 00000000000..b97ede691a5
--- /dev/null
+++ b/vertx-core/src/main/java/io/vertx/core/internal/http/ClientWebSocketInternal.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2011-2024 Contributors to the Eclipse Foundation
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
+ * which is available at https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
+ */
+package io.vertx.core.internal.http;
+
+import io.vertx.core.Context;
+import io.vertx.core.Future;
+import io.vertx.core.http.ClientWebSocket;
+import io.vertx.core.http.WebSocket;
+import io.vertx.core.http.WebSocketConnectOptions;
+
+public interface ClientWebSocketInternal extends ClientWebSocket, WebSocketInternal {
+
+  Future<WebSocket> connect(Context context, WebSocketConnectOptions options);
+
+}