From a2ce1857163f1690843035baf566847cf44e135f Mon Sep 17 00:00:00 2001
From: Julien Viet <julien@julienviet.com>
Date: Mon, 19 Jul 2021 09:47:33 +0200
Subject: [PATCH] Server WebSocket handshake should set the HttpServerResponse
 as written - see #4020 - see #4016

---
 .../vertx/core/http/impl/Http1xServerResponse.java   | 12 ++++++++++++
 .../io/vertx/core/http/impl/ServerWebSocketImpl.java |  8 +-------
 src/test/java/io/vertx/core/http/WebSocketTest.java  | 12 ++++++++++++
 3 files changed, 25 insertions(+), 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
index 90689dbd7b4..6dbe295a1cc 100644
--- 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
index 4be8c38e5f0..198e16d25bf 100644
--- 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
index 02ab60ac0e3..11405cf5099 100644
--- 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();