Improving Websocket logging by adding API URI into logs (#12508)

wso2 · Aug 13, 2024 · 2af6297 · 2af6297
1 parent b9872c5
commit 2af6297
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/...ateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/WebsocketInboundHandler.java b/...ateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/WebsocketInboundHandler.java
@@ -79,6 +79,7 @@ public class WebsocketInboundHandler extends ChannelInboundHandlerAdapter {
     private WebSocketAnalyticsMetricsHandler metricsHandler;
     private InboundWebSocketProcessor webSocketProcessor;
     private final String API_PROPERTIES = "API_PROPERTIES";
+    private final String API_CONTEXT_URI = "API_CONTEXT_URI";
     private final String WEB_SC_API_UT = "api.ut.WS_SC";
 
     public WebsocketInboundHandler() {
@@ -161,6 +162,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
                 if (!responseDTO.isError()) {
                     setApiAuthPropertiesToChannel(ctx, inboundMessageContext);
                     setApiPropertiesMapToChannel(ctx, inboundMessageContext);
+                    setApiContextUriToChannel(ctx, inboundMessageContext);
                     if (StringUtils.isNotEmpty(inboundMessageContext.getToken())) {
                         String backendJwtHeader = null;
                         JWTConfigurationDto jwtConfigurationDto = ServiceReferenceHolder.getInstance()
@@ -419,6 +421,13 @@ private void setApiPropertiesMapToChannel(ChannelHandlerContext ctx, InboundMess
         ctx.channel().attr(AttributeKey.valueOf(API_PROPERTIES)).set(createApiPropertiesMap(inboundMessageContext));
     }
 
+    private void setApiContextUriToChannel(ChannelHandlerContext ctx, InboundMessageContext inboundMessageContext) {
+
+        Map<String, String> apiContextUriMap = new HashMap<>();
+        apiContextUriMap.put("apiContextUri", inboundMessageContext.getRequestPath());
+        ctx.channel().attr(AttributeKey.valueOf(API_CONTEXT_URI)).set(apiContextUriMap);
+    }
+
     private Map<String, Object> createApiPropertiesMap(InboundMessageContext inboundMessageContext) {
 
         Map<String, Object> apiPropertiesMap = new HashMap<>();
@@ -452,6 +461,18 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
             CorruptedWebSocketFrameException corruptedWebSocketFrameException = ((CorruptedWebSocketFrameException) cause);
             apiProperties.put(WEB_SC_API_UT, corruptedWebSocketFrameException.closeStatus().code());
         }
-        super.exceptionCaught(ctx, cause);
+
+        // Improve Websocket logging by adding API URI into log
+        Attribute<Object> apiContextUriAttributes = ctx.channel().attr(AttributeKey.valueOf(API_CONTEXT_URI));
+        HashMap apiContextUris = (HashMap) apiContextUriAttributes.get();
+        String apiContextUri = (String) apiContextUris.get("apiContextUri");
+
+        if (apiContextUri != null) {
+            Throwable newCause = new Throwable(cause.getMessage() + " For the URI: " + apiContextUri);
+            newCause.initCause(cause);
+            super.exceptionCaught(ctx, newCause);
+        } else {
+            super.exceptionCaught(ctx, cause);
+        }
     }
 }
diff --git a/...rc/test/java/org/wso2/carbon/apimgt/gateway/handlers/WebsocketInboundHandlerTestCase.java b/...rc/test/java/org/wso2/carbon/apimgt/gateway/handlers/WebsocketInboundHandlerTestCase.java
@@ -96,6 +96,7 @@ public void setup() throws Exception {
         Attribute attribute = Mockito.mock(Attribute.class);
         ChannelId channelId = Mockito.mock(ChannelId.class);
         Mockito.when(channel.attr(AttributeKey.valueOf("API_PROPERTIES"))).thenReturn(attribute);
+        Mockito.when(channel.attr(AttributeKey.valueOf("API_CONTEXT_URI"))).thenReturn(attribute);
         Mockito.when(channelHandlerContext.channel()).thenReturn(channel);
         Mockito.when(channel.id()).thenReturn(channelId);
         Mockito.when(channelId.asLongText()).thenReturn(channelIdString);
@@ -240,11 +241,20 @@ public void testWSFrameResponse() throws Exception {
     public void exceptionCaughtTest() throws Exception {
         Throwable cause = new CorruptedWebSocketFrameException(WebSocketCloseStatus.MESSAGE_TOO_BIG,
                 "Max frame length of 65536 has been exceeded.");
-        Attribute<Object> attributes = Mockito.mock(Attribute.class);
+        Attribute<Object> apiPropertiesAttributes = Mockito.mock(Attribute.class);
+        Attribute<Object> contextUriAttributes = Mockito.mock(Attribute.class);
+
         Mockito.when(channelHandlerContext.channel().attr(AttributeKey.valueOf("API_PROPERTIES")))
-                .thenReturn(attributes);
+                .thenReturn(apiPropertiesAttributes);
+        Mockito.when(channelHandlerContext.channel().attr(AttributeKey.valueOf("API_CONTEXT_URI")))
+                .thenReturn(contextUriAttributes);
+
         HashMap apiProperties = new HashMap();
-        Mockito.when((HashMap)attributes.get()).thenReturn(apiProperties);
+        HashMap apiContextUriAttributes = new HashMap();
+
+        Mockito.when((HashMap)apiPropertiesAttributes.get()).thenReturn(apiProperties);
+        Mockito.when((HashMap)contextUriAttributes.get()).thenReturn(apiContextUriAttributes);
+
         websocketInboundHandler.exceptionCaught(channelHandlerContext, cause);
         Assert.assertEquals(apiProperties.get("api.ut.WS_SC"), 1009);
     }