Removes all possible upgrade headers that don't belong in a failed re…

…sponse. Some other minor cleanup.

Signed-off-by: Santiago Pericas-Geertsen <[email protected]>

webserver/websocket/src/main/java/io/helidon/webserver/websocket/WebSocketUpgradeCodec.java

@@ -18,12 +18,12 @@
 import java.nio.charset.Charset;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import io.helidon.common.http.Http;
 import io.helidon.webserver.ForwardingHandler;
-
 import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelFutureListener;
@@ -41,9 +41,11 @@
 import org.glassfish.tyrus.core.TyrusUpgradeResponse;
 
 class WebSocketUpgradeCodec implements HttpServerUpgradeHandler.UpgradeCodec {
-
     private static final Logger LOGGER = Logger.getLogger(WebSocketUpgradeCodec.class.getName());
 
+    private static final String SEC_WEBSOCKET_ACCEPT = "Sec-WebSocket-Accept";
+    private static final String SEC_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol";
+
     private final WebSocketRouting webSocketRouting;
     private String path;
     private WebSocketHandler wsHandler;
@@ -73,23 +75,26 @@ public boolean prepareUpgradeResponse(ChannelHandlerContext ctx,
             TyrusUpgradeResponse upgradeResponse = wsHandler.upgradeResponse();
             if (upgradeResponse.getStatus() != Http.Status.SWITCHING_PROTOCOLS_101.code()) {
                 // prepare headers for failed response
-                upgradeResponse.getHeaders().remove(Http.Header.UPGRADE);
-                upgradeResponse.getHeaders().remove(Http.Header.CONNECTION);
-                upgradeResponse.getHeaders().remove("sec-websocket-accept");
+                Map<String, List<String>> upgradeHeaders = upgradeResponse.getHeaders();
+                upgradeHeaders.remove(Http.Header.UPGRADE);
+                upgradeHeaders.remove(Http.Header.CONNECTION);
+                upgradeHeaders.remove(SEC_WEBSOCKET_ACCEPT);
+                upgradeHeaders.remove(SEC_WEBSOCKET_PROTOCOL);
                 HttpHeaders headers = new DefaultHttpHeaders();
-                upgradeResponse.getHeaders().forEach(headers::add);
+                upgradeHeaders.forEach(headers::add);
 
                 // set payload as text/plain with reason phrase
                 headers.add(Http.Header.CONTENT_TYPE, "text/plain");
                 String reasonPhrase = upgradeResponse.getReasonPhrase() == null ? ""
                         : upgradeResponse.getReasonPhrase();
-                HttpResponse r = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1,
+                HttpResponse httpResponse = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1,
                         HttpResponseStatus.valueOf(upgradeResponse.getStatus()),
                         Unpooled.wrappedBuffer(reasonPhrase.getBytes(Charset.defaultCharset())),
-                        headers, EmptyHttpHeaders.INSTANCE);
+                        headers,
+                        EmptyHttpHeaders.INSTANCE);     // trailing headers
 
                 // write, flush and later close connection
-                ChannelFuture writeComplete = ctx.writeAndFlush(r);
+                ChannelFuture writeComplete = ctx.writeAndFlush(httpResponse);
                 writeComplete.addListener(ChannelFutureListener.CLOSE);
                 return false;
             }