Fix issues with the new handshake API.

tsegismont · Oct 11, 2024 · 6d98edb · 6d98edb
1 parent cb65296
commit 6d98edb
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 16 deletions.
diff --git a/vertx-core/src/main/java/io/vertx/core/http/impl/Http1xServerRequest.java b/vertx-core/src/main/java/io/vertx/core/http/impl/Http1xServerRequest.java
@@ -386,22 +386,22 @@ public String getFormAttribute(String attributeName) {
 
   @Override
   public Future<ServerWebSocket> toWebSocket() {
-    return webSocket().compose(handshake -> handshake.accept());
+    return webSocketHandshake().compose(handshake -> handshake.accept());
   }
 
   /**
    * @return a future of the un-accepted WebSocket
    */
-  Future<ServerWebSocketHandshake> webSocket() {
+  Future<ServerWebSocketHandshake> webSocketHandshake() {
     PromiseInternal<ServerWebSocketHandshake> promise = context.promise();
-    webSocket(promise);
+    webSocketHandshake(promise);
     return promise.future();
   }
 
   /**
    * Handle the request when a WebSocket upgrade header is present.
    */
-  private void webSocket(PromiseInternal<ServerWebSocketHandshake> promise) {
+  private void webSocketHandshake(PromiseInternal<ServerWebSocketHandshake> promise) {
     BufferInternal body = BufferInternal.buffer();
     boolean[] failed = new boolean[1];
     handler(buff -> {

diff --git a/vertx-core/src/main/java/io/vertx/core/http/impl/Http1xServerRequestHandler.java b/vertx-core/src/main/java/io/vertx/core/http/impl/Http1xServerRequestHandler.java
@@ -41,19 +41,20 @@ public void handle(HttpServerRequest req) {
     Handler<ServerWebSocket> wsHandler = handlers.webSocketHandler;
     Handler<ServerWebSocketHandshake> wsHandshakeHandler = handlers.webSocketHandshakeHandler;
     Handler<HttpServerRequest> reqHandler = handlers.requestHandler;
-    if (wsHandler != null ) {
+    if (wsHandler != null || wsHandshakeHandler != null) {
       if (req.headers().contains(UPGRADE, WEBSOCKET, true) && handlers.server.wsAccept()) {
         // Missing upgrade header + null request handler will be handled when creating the handshake by sending a 400 error
-        ((Http1xServerRequest)req).webSocket().onComplete(ar -> {
+        ((Http1xServerRequest)req).webSocketHandshake().onComplete(ar -> {
           if (ar.succeeded()) {
-            ServerWebSocketHandshake handshake = ar.result();
-            if (wsHandshakeHandler != null) {
-              wsHandshakeHandler.handle(handshake);
+            ServerWebSocketHandshaker handshake = (ServerWebSocketHandshaker) ar.result();
+            if (wsHandshakeHandler == null) {
+              handshake.accept();
             } else {
-              handshake.accept().onSuccess(wsHandler);
+              wsHandshakeHandler.handle(handshake);
+            }
+            if (wsHandler != null) {
+              handshake.onSuccess(wsHandler);
             }
-          } else {
-            // ????
           }
         });
       } else {

diff --git a/vertx-core/src/main/java/io/vertx/core/http/impl/HttpServerImpl.java b/vertx-core/src/main/java/io/vertx/core/http/impl/HttpServerImpl.java
@@ -172,7 +172,7 @@ public Future<HttpServer> listen() {
 
   @Override
   public synchronized Future<HttpServer> listen(SocketAddress address) {
-    if (requestHandler == null && webSocketHandler == null) {
+    if (requestHandler == null && webSocketHandler == null && webSocketHandhakeHandler == null) {
       throw new IllegalStateException("Set request or WebSocket handler first");
     }
     if (tcpServer != null) {

diff --git a/vertx-core/src/main/java/io/vertx/core/http/impl/ServerWebSocketHandshaker.java b/vertx-core/src/main/java/io/vertx/core/http/impl/ServerWebSocketHandshaker.java
@@ -22,6 +22,7 @@
 import io.vertx.core.*;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.*;
+import io.vertx.core.impl.future.FutureImpl;
 import io.vertx.core.internal.ContextInternal;
 import io.vertx.core.net.HostAndPort;
 import io.vertx.core.net.SocketAddress;
@@ -33,6 +34,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 
 import static io.netty.handler.codec.http.HttpResponseStatus.BAD_REQUEST;
@@ -43,14 +45,15 @@
  *
  * @author <a href="mailto:[email protected]">Julien Viet</a>
  */
-public class ServerWebSocketHandshaker implements ServerWebSocketHandshake, ServerWebSocket {
+public class ServerWebSocketHandshaker extends FutureImpl<ServerWebSocket> implements ServerWebSocketHandshake, ServerWebSocket {
 
   private final Http1xServerRequest request;
   private final HttpServerOptions options;
   private final WebSocketServerHandshaker handshaker;
   private boolean done;
 
   public ServerWebSocketHandshaker(Http1xServerRequest request, WebSocketServerHandshaker handshaker, HttpServerOptions options) {
+    super(request.context);
     this.request = request;
     this.handshaker = handshaker;
     this.options = options;
@@ -93,7 +96,6 @@ public Future<ServerWebSocket> accept() {
     try {
       ws = acceptHandshake();
     } catch (Exception e) {
-      e.printStackTrace(System.out);
       return rejectHandshake(BAD_REQUEST.code())
         .transform(ar -> {
           if (ar.succeeded()) {
@@ -104,7 +106,8 @@ public Future<ServerWebSocket> accept() {
           }
         });
     }
-    return request.context.succeededFuture(ws);
+    tryComplete(ws);
+    return this;
   }
 
   @Override
@@ -116,6 +119,7 @@ public Future<Void> reject(int sc) {
       }
       done = true;
     }
+    tryFail(new RejectedExecutionException()); // Not great but for now OK
     return rejectHandshake(sc);
   }