From 097bedcb9b22e56e7fa6da548f22f11b059f5264 Mon Sep 17 00:00:00 2001 From: Jon Chambers Date: Wed, 22 May 2024 13:11:13 -0400 Subject: [PATCH] Add a health check handler to the Noise-over-WebSocket pipeline --- .../grpc/net/NoiseWebSocketTunnelServer.java | 3 ++- .../net/WebSocketOpeningHandshakeHandler.java | 9 ++++++++- .../WebSocketOpeningHandshakeHandlerTest.java | 18 ++++++++++++++---- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/NoiseWebSocketTunnelServer.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/NoiseWebSocketTunnelServer.java index c997eb9f6..4cc0d1dee 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/NoiseWebSocketTunnelServer.java +++ b/service/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/service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandler.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandler.java index 453ccb390..ddca6c4fe 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandler.java +++ b/service/src/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/service/src/test/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandlerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandlerTest.java index 0f27e5efb..1dd906bb2 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/net/WebSocketOpeningHandshakeHandlerTest.java +++ b/service/src/test/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()); } }