Add a health check handler to the Noise-over-WebSocket pipeline

signalapp · May 22, 2024 · 097bedc · 097bedc
1 parent 907ff89
commit 097bedc
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 6 deletions.
diff --git a/...e/src/main/java/org/whispersystems/textsecuregcm/grpc/net/NoiseWebSocketTunnelServer.java b/...e/src/main/java/org/whispersystems/textsecuregcm/grpc/net/NoiseWebSocketTunnelServer.java
@@ -41,6 +41,7 @@ public class NoiseWebSocketTunnelServer implements Managed {
 
   static final String AUTHENTICATED_SERVICE_PATH = "/authenticated";
   static final String ANONYMOUS_SERVICE_PATH = "/anonymous";
+  static final String HEALTH_CHECK_PATH = "/health-check";
 
   private static final Logger log = LoggerFactory.getLogger(NoiseWebSocketTunnelServer.class);
 
@@ -96,7 +97,7 @@ protected void initChannel(SocketChannel socketChannel) {
                 .addLast(new HttpObjectAggregator(Noise.MAX_PACKET_LEN))
                 // The WebSocket opening handshake handler will remove itself from the pipeline once it has received a valid WebSocket upgrade
                 // request and passed it down the pipeline
-                .addLast(new WebSocketOpeningHandshakeHandler(AUTHENTICATED_SERVICE_PATH, ANONYMOUS_SERVICE_PATH))
+                .addLast(new WebSocketOpeningHandshakeHandler(AUTHENTICATED_SERVICE_PATH, ANONYMOUS_SERVICE_PATH, HEALTH_CHECK_PATH))
                 .addLast(new WebSocketServerProtocolHandler("/", true))
                 .addLast(new RejectUnsupportedMessagesHandler())
                 // The WebSocket handshake complete listener will replace itself with an appropriate Noise handshake handler once

diff --git a/...main/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandler.java b/...main/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandler.java
@@ -19,10 +19,15 @@ class WebSocketOpeningHandshakeHandler extends ChannelInboundHandlerAdapter {
 
   private final String authenticatedPath;
   private final String anonymousPath;
+  private final String healthCheckPath;
+
+  WebSocketOpeningHandshakeHandler(final String authenticatedPath,
+      final String anonymousPath,
+      final String healthCheckPath) {
 
-  WebSocketOpeningHandshakeHandler(final String authenticatedPath, final String anonymousPath) {
     this.authenticatedPath = authenticatedPath;
     this.anonymousPath = anonymousPath;
+    this.healthCheckPath = healthCheckPath;
   }
 
   @Override
@@ -43,6 +48,8 @@ public void channelRead(final ChannelHandlerContext context, final Object messag
               } else {
                 closeConnectionWithStatus(context, request, HttpResponseStatus.UPGRADE_REQUIRED);
               }
+            } else if (healthCheckPath.equals(request.uri())) {
+              closeConnectionWithStatus(context, request, HttpResponseStatus.NO_CONTENT);
             } else {
               closeConnectionWithStatus(context, request, HttpResponseStatus.NOT_FOUND);
             }

diff --git a/.../java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandlerTest.java b/.../java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandlerTest.java
@@ -15,7 +15,6 @@
 import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.HttpVersion;
-import io.netty.util.ReferenceCountUtil;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -27,10 +26,12 @@ class WebSocketOpeningHandshakeHandlerTest extends AbstractLeakDetectionTest {
 
   private static final String AUTHENTICATED_PATH = "/authenticated";
   private static final String ANONYMOUS_PATH = "/anonymous";
+  private static final String HEALTH_CHECK_PATH = "/health-check";
 
   @BeforeEach
   void setUp() {
-    embeddedChannel = new EmbeddedChannel(new WebSocketOpeningHandshakeHandler(AUTHENTICATED_PATH, ANONYMOUS_PATH));
+    embeddedChannel =
+        new EmbeddedChannel(new WebSocketOpeningHandshakeHandler(AUTHENTICATED_PATH, ANONYMOUS_PATH, HEALTH_CHECK_PATH));
   }
 
   @ParameterizedTest
@@ -50,6 +51,16 @@ void handleValidRequest(final String path) {
     }
   }
 
+  @Test
+  void handleHealthCheckRequest() {
+    final FullHttpRequest request = buildRequest(HttpMethod.GET, HEALTH_CHECK_PATH, new DefaultHttpHeaders());
+
+    embeddedChannel.writeOneInbound(request);
+
+    assertEquals(0, request.refCnt());
+    assertHttpResponse(HttpResponseStatus.NO_CONTENT);
+  }
+
   @ParameterizedTest
   @ValueSource(strings = { AUTHENTICATED_PATH, ANONYMOUS_PATH })
   void handleUpgradeRequired(final String path) {
@@ -89,7 +100,6 @@ private void assertHttpResponse(final HttpResponseStatus expectedStatus) {
 
     final FullHttpResponse response = assertInstanceOf(FullHttpResponse.class, embeddedChannel.outboundMessages().poll());
 
-    //noinspection DataFlowIssue
     assertEquals(expectedStatus, response.status());
   }
 
@@ -99,6 +109,6 @@ private FullHttpRequest buildRequest(final HttpMethod method, final String path,
         path,
         Unpooled.buffer(0),
         headers,
-        new DefaultHttpHeaders(true));
+        new DefaultHttpHeaders());
   }
 }