Server WebSocket handshake should set the HttpServerResponse as writt…

…en - see #4020 - see #4016
eclipse-vertx · Jul 19, 2021 · a2ce185 · a2ce185
1 parent 2c870a9
commit a2ce185
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 7 deletions.
diff --git a/src/main/java/io/vertx/core/http/impl/Http1xServerResponse.java b/src/main/java/io/vertx/core/http/impl/Http1xServerResponse.java
@@ -420,6 +420,18 @@ private void end(Buffer chunk, PromiseInternal<Void> listener) {
     }
   }
 
+  void completeHandshake() {
+    if (conn.metrics != null) {
+      conn.metrics.responseBegin(requestMetric, this);
+    }
+    setStatusCode(101);
+    synchronized (conn) {
+      headWritten = true;
+      written = true;
+    }
+    conn.responseComplete();
+  }
+
   @Override
   public void close() {
     synchronized (conn) {

diff --git a/src/main/java/io/vertx/core/http/impl/ServerWebSocketImpl.java b/src/main/java/io/vertx/core/http/impl/ServerWebSocketImpl.java
@@ -185,22 +185,16 @@ private void handleHandshake(int sc) {
 
   private void doHandshake() {
     Channel channel = conn.channel();
-    Object metric;
     Http1xServerResponse response = request.response();
     try {
       handshaker.handshake(channel, request.nettyRequest());
-      metric = request.metric;
     } catch (Exception e) {
       response.setStatusCode(BAD_REQUEST.code()).end();
       throw e;
     } finally {
       request = null;
     }
-    response.setStatusCode(101);
-    if (conn.metrics != null) {
-      conn.metrics.responseBegin(metric, response);
-    }
-    conn.responseComplete();
+    response.completeHandshake();
     status = SWITCHING_PROTOCOLS.code();
     subProtocol(handshaker.selectedSubprotocol());
     // remove compressor as its not needed anymore once connection was upgraded to websockets

diff --git a/src/test/java/io/vertx/core/http/WebSocketTest.java b/src/test/java/io/vertx/core/http/WebSocketTest.java
@@ -1956,6 +1956,18 @@ private void testUpgrade(boolean delayed) {
     server.requestHandler(request -> {
       Runnable runner = () -> {
         request.toWebSocket().onComplete(onSuccess(ws -> {
+          HttpServerResponse response = request.response();
+          assertTrue(response.ended());
+          try {
+            response.putHeader("foo", "bar");
+            fail();
+          } catch (IllegalStateException ignore) {
+          }
+          try {
+            response.end();
+            fail();
+          } catch (IllegalStateException ignore) {
+          }
           ws.handler(buff -> {
             ws.write(Buffer.buffer("helloworld"));
             ws.close();