Fetching AuditLogRequest from Security Request Context in Write and C…

…lose.

cdap-common/src/main/java/io/cdap/cdap/common/http/AuthenticationChannelHandler.java

@@ -132,15 +132,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
       SecurityRequestContext.setUserIp(currentUserIp);
     }
 
-    try {
-      ctx.fireChannelRead(msg);
-    } finally {
-      // Set the audit log info onto the ongoing channel so it is ensured to be reused later in the same channel, making
-      // it independent of Thread local.
-      ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_REQ_ATTR_NAME))
-        .set(SecurityRequestContext.getAuditLogRequest());
-      SecurityRequestContext.reset();
-    }
+
+    ctx.fireChannelRead(msg);
   }
 
   /**
@@ -149,17 +142,25 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
    */
   @Override
   public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
-    publishAuditLogRequest(ctx);
-    super.write(ctx, msg, promise);
+    try {
+      publishAuditLogRequest();
+      super.write(ctx, msg, promise);
+    } finally {
+      SecurityRequestContext.reset();
+    }
   }
 
   /**
    * Need to handle for the case when "write" is not called in the netty channel.
    */
   @Override
   public void close(ChannelHandlerContext ctx, ChannelPromise promise) throws Exception {
-    publishAuditLogRequest(ctx);
-    super.close(ctx, promise);
+    try {
+      publishAuditLogRequest();
+      super.close(ctx, promise);
+    } finally {
+      SecurityRequestContext.reset();
+    }
   }
 
   @Override
@@ -174,14 +175,14 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
   }
 
   /**
-   * Check the audit log attribute attached to a channel.
-   * It's not null, publish it and set it to Null.
+   * When write / close is called , it is expected that {@link io.cdap.cdap.common.auditlogging.AuditLogSetterHook}'s
+   * postCall method is triggered, which will set AuditLogRequest in SecurityRequestContext.
+   * TODO :  CDAP-21085
    */
-  private void publishAuditLogRequest(ChannelHandlerContext ctx) throws IOException {
-    Object auditLogRequestObj = ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_REQ_ATTR_NAME)).get();
-    if (auditLoggingEnabled && auditLogRequestObj != null) {
-      auditLogWriter.publish((AuditLogRequest) auditLogRequestObj);
-      ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_REQ_ATTR_NAME)).set(null);
+  private void publishAuditLogRequest() throws IOException {
+    AuditLogRequest auditLogRequest = SecurityRequestContext.getAuditLogRequest();
+    if (auditLoggingEnabled && auditLogRequest != null ) {
+      auditLogWriter.publish(auditLogRequest);
     }
   }
 }