A few improvements in websocket rejection: the ServerWebSocket#reject…

…() is overloaded to provide a specific status code, the reject() method is documented that it will send 502 by default and not 404, the client can be aware of the rejection status with the new WebsocketRejectedException. fixes #1996

src/main/java/io/vertx/core/http/ServerWebSocket.java

@@ -107,12 +107,17 @@ public interface ServerWebSocket extends WebSocketBase {
    * <p>
    * Calling this method from the websocket handler when it is first passed to you gives you the opportunity to reject
    * the websocket, which will cause the websocket handshake to fail by returning
-   * a 404 response code.
+   * a {@literal 502} response code.
    * <p>
    * You might use this method, if for example you only want to accept WebSockets with a particular path.
    */
   void reject();
 
+  /**
+   * Like {@link #reject()} but with a {@code status}.
+   */
+  void reject(int status);
+
   /**
    * @return an array of the peer certificates. Returns null if connection is
    *         not SSL.

src/main/java/io/vertx/core/http/WebsocketRejectedException.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2011-2014 The original author or authors
+ * ------------------------------------------------------
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Eclipse Public License v1.0
+ * and Apache License v2.0 which accompanies this distribution.
+ *
+ *     The Eclipse Public License is available at
+ *     http://www.eclipse.org/legal/epl-v10.html
+ *
+ *     The Apache License v2.0 is available at
+ *     http://www.opensource.org/licenses/apache2.0.php
+ *
+ * You may elect to redistribute this code under either of these licenses.
+ */
+package io.vertx.core.http;
+
+import io.vertx.core.VertxException;
+
+/**
+ * @author <a href="mailto:julien@julienviet.com">Julien Viet</a>
+ */
+public class WebsocketRejectedException extends VertxException {
+
+  private final int status;
+
+  WebsocketRejectedException(int status) {
+    super("Websocket connection attempt returned HTTP status code " + status);
+    this.status = status;
+  }
+
+  public int getStatus() {
+    return status;
+  }
+}

src/main/java/io/vertx/core/http/impl/ClientConnection.java

@@ -177,13 +177,14 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
       if (handshaker != null && handshaking) {
         if (msg instanceof HttpResponse) {
           HttpResponse resp = (HttpResponse) msg;
-          if (resp.getStatus().code() != 101) {
+          HttpResponseStatus status = resp.status();
+          if (status.code() != 101) {
             handshaker = null;
             close();
-            handleException(new WebSocketHandshakeException("Websocket connection attempt returned HTTP status code " + resp.getStatus().code()));
+            handleException(new WebsocketRejectedException(status.code()));
             return;
           }
-          response = new DefaultFullHttpResponse(resp.getProtocolVersion(), resp.getStatus());
+          response = new DefaultFullHttpResponse(resp.protocolVersion(), status);
           response.headers().add(resp.headers());
         }
 
@@ -214,7 +215,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
       }
     }
 
-    private void handleException(WebSocketHandshakeException e) {
+    private void handleException(Exception e) {
       handshaking = false;
       buffered.clear();
       Handler<Throwable> handler = exceptionHandler();

src/main/java/io/vertx/core/http/impl/HttpServerImpl.java

@@ -731,7 +731,11 @@ protected void handshake(ServerConnection conn, FullHttpRequest request, Channel
             }
             ws.connectNow();
           } else {
-            ch.writeAndFlush(new DefaultFullHttpResponse(HTTP_1_1, BAD_GATEWAY));
+            HttpResponseStatus status = ws.getRejectedStatus();
+            if (status == null) {
+              status = BAD_GATEWAY;
+            }
+            ch.writeAndFlush(new DefaultFullHttpResponse(HTTP_1_1, status));
           }
         });
       }

src/main/java/io/vertx/core/http/impl/ServerWebSocketImpl.java

@@ -16,6 +16,7 @@
 
 package io.vertx.core.http.impl;
 
+import io.netty.handler.codec.http.HttpResponseStatus;
 import io.vertx.core.MultiMap;
 import io.vertx.core.http.ServerWebSocket;
 import io.vertx.core.http.WebSocketFrame;
@@ -25,6 +26,8 @@
 import javax.net.ssl.SSLPeerUnverifiedException;
 import javax.security.cert.X509Certificate;
 
+import static io.netty.handler.codec.http.HttpResponseStatus.BAD_GATEWAY;
+
 /**
  * This class is optimised for performance when used on the same event loop. However it can be used safely from other threads.
  *
@@ -44,6 +47,7 @@ public class ServerWebSocketImpl extends WebSocketImplBase<ServerWebSocket> impl
 
   private boolean connected;
   private boolean rejected;
+  private HttpResponseStatus rejectedStatus;
 
   public ServerWebSocketImpl(VertxInternal vertx, String uri, String path, String query, MultiMap headers,
                              ConnectionBase conn, boolean supportsContinuation, Runnable connectRunnable,
@@ -78,6 +82,15 @@ public MultiMap headers() {
 
   @Override
   public void reject() {
+    reject(BAD_GATEWAY);
+  }
+
+  @Override
+  public void reject(int status) {
+    reject(HttpResponseStatus.valueOf(status));
+  }
+
+  private void reject(HttpResponseStatus status) {
     synchronized (conn) {
       checkClosed();
       if (connectRunnable == null) {
@@ -87,6 +100,7 @@ public void reject() {
         throw new IllegalStateException("Cannot reject websocket, it has already been written to");
       }
       rejected = true;
+      rejectedStatus = status;
     }
   }
 
@@ -101,7 +115,7 @@ public void close() {
       checkClosed();
       if (connectRunnable != null) {
         // Server side
-        if (rejected) {
+        if (isRejected()) {
           throw new IllegalStateException("Cannot close websocket, it has been rejected");
         }
         if (!connected && !closed) {
@@ -116,7 +130,7 @@ public void close() {
   public ServerWebSocket writeFrame(WebSocketFrame frame) {
     synchronized (conn) {
       if (connectRunnable != null) {
-        if (rejected) {
+        if (isRejected()) {
           throw new IllegalStateException("Cannot write to websocket, it has been rejected");
         }
         if (!connected && !closed) {
@@ -135,15 +149,18 @@ private void connect() {
   // Connect if not already connected
   void connectNow() {
     synchronized (conn) {
-      if (!connected && !rejected) {
+      if (!connected && !isRejected()) {
         connect();
       }
     }
   }
 
-  boolean isRejected() {
-    synchronized (conn) {
-      return rejected;
-    }
+  synchronized boolean isRejected() {
+    return rejected;
   }
+
+  synchronized HttpResponseStatus getRejectedStatus() {
+    return rejectedStatus;
+  }
+
 }

src/test/java/io/vertx/test/core/WebsocketTest.java

@@ -28,6 +28,7 @@
 import io.vertx.core.VertxOptions;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.*;
+import io.vertx.core.http.WebsocketRejectedException;
 import io.vertx.core.impl.ConcurrentHashSet;
 import io.vertx.core.net.NetServer;
 import io.vertx.core.net.NetSocket;
@@ -36,7 +37,6 @@
 import io.vertx.test.core.tls.Trust;
 import org.junit.Test;
 
-import javax.net.ssl.SSLPeerUnverifiedException;
 import javax.security.cert.X509Certificate;
 import java.io.UnsupportedEncodingException;
 import java.security.MessageDigest;
@@ -103,14 +103,18 @@ protected VertxOptions getOptions() {
 
   @Test
   public void testRejectHybi00() throws Exception {
-    testReject(WebsocketVersion.V00);
+    testReject(WebsocketVersion.V00, null, 502);
   }
 
   @Test
   public void testRejectHybi08() throws Exception {
-    testReject(WebsocketVersion.V08);
+    testReject(WebsocketVersion.V08, null, 502);
   }
 
+  @Test
+  public void testRejectWithStatusCode() throws Exception {
+    testReject(WebsocketVersion.V08, 404, 404);
+  }
 
   @Test
   public void testWSBinaryHybi00() throws Exception {
@@ -1005,19 +1009,27 @@ public void testInvalidMethod() throws Exception {
     await();
   }
 
-  private void testReject(WebsocketVersion version) throws Exception {
+  private void testReject(WebsocketVersion version, Integer rejectionStatus, int expectedRejectionStatus) throws Exception {
 
     String path = "/some/path";
 
     server = vertx.createHttpServer(new HttpServerOptions().setPort(HttpTestBase.DEFAULT_HTTP_PORT)).websocketHandler(ws -> {
       assertEquals(path, ws.path());
-      ws.reject();
+      if (rejectionStatus != null) {
+        ws.reject(rejectionStatus);
+      } else {
+        ws.reject();
+      }
     });
 
     server.listen(ar -> {
       assertTrue(ar.succeeded());
       client.websocketStream(HttpTestBase.DEFAULT_HTTP_PORT, HttpTestBase.DEFAULT_HTTP_HOST, path, null, version).
-          exceptionHandler(t -> testComplete()).
+          exceptionHandler(t -> {
+            assertTrue(t instanceof WebsocketRejectedException);
+            assertEquals(expectedRejectionStatus, ((WebsocketRejectedException)t).getStatus());
+            testComplete();
+          }).
           handler(ws -> fail("Should not be called"));
     });
     await();
@@ -1607,10 +1619,9 @@ public void testRaceConditionWithWebsocketClientWorker2() throws Exception {
 
   @Test
   public void httpClientWebsocketConnectionFailureHandlerShouldBeCalled() throws Exception {
-    String nonExistingHost = "idont.even.exist";
     int port = 7867;
     HttpClient client = vertx.createHttpClient();
-    client.websocket(port, nonExistingHost, "", websocket -> {
+    client.websocket(port, "localhost", "", websocket -> {
       websocket.handler(data -> {
         fail("connection should not succeed");
       });
@@ -1716,8 +1727,8 @@ private void doTestClientWebsocketConnectionCloseOnBadResponse(boolean keepAlive
         throw new AssertionError("Timed out waiting for expected state, current: serverGotClose = " + serverGotClose + ", clientGotCorrectException = " + clientGotCorrectException);
       } else if (result == serverGotCloseException) {
         serverGotClose = true;
-      } else if (result instanceof WebSocketHandshakeException
-              && result.getMessage().equals("Websocket connection attempt returned HTTP status code 200")) {
+      } else if (result instanceof WebsocketRejectedException
+              && ((WebsocketRejectedException)result).getStatus() == 200) {
         clientGotCorrectException = true;
       } else {
         throw result;