[CDAP-21105] Address comment : Make builder

cdap-common/src/main/java/io/cdap/cdap/common/auditlogging/AuditLogSetterHook.java

@@ -64,14 +64,17 @@ public void postCall(HttpRequest request, HttpResponseStatus status, HandlerInfo
     long startTime = endTime - (endTimeNanos - startTimeNanos);
 
     LOG.trace("Setting a Audit Log event to published in SecurityRequestContext");
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.OP_RESP_CODE, status.code());
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.URI, request.uri());
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.HANDLER,
-                                                 getSimpleName(handlerInfo.getHandlerName()));
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.METHOD, handlerInfo.getMethodName());
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.METHOD_TYPE, request.method().name());
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.START_TIME_NANOS, startTime);
-    SecurityRequestContext.setAuditMetaDataInMap(AuditLogRequest.PropKey.END_TIME_NANOS, endTime);
+
+    AuditLogRequest.Builder auditLogRequestBuilder = new AuditLogRequest.Builder()
+        .operationResponseCode(status.code())
+        .uri(request.uri())
+        .handler(getSimpleName(handlerInfo.getHandlerName()))
+        .method(handlerInfo.getMethodName())
+        .methodType(request.method().name())
+        .startTimeNanos(startTime)
+        .endTimeNanos(endTime);
+
+    SecurityRequestContext.setAuditLogRequestBuilder(auditLogRequestBuilder);
   }
 
   private String getSimpleName(String className) {

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

@@ -38,8 +38,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
-import java.util.Map;
 import java.util.Queue;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -55,7 +55,9 @@ public class AuthenticationChannelHandler extends ChannelDuplexHandler {
       "CDAP-empty-user-credential",
       Credential.CredentialType.INTERNAL);
   private static final String EMPTY_USER_IP = "CDAP-empty-user-ip";
-  private static final String AUDIT_METADATA_MAP_ATTR_NAME = "AUDIT_LOG_METADATA_MAP";
+  static final String AUDIT_LOG_REQ_BUILDER_ATTR = "AUDIT_LOG_REQ_BUILDER";
+  static final String AUDIT_LOG_USER_IP_ATTR = "AUDIT_LOG_USER_IP";
+  static final String AUDIT_LOG_CONTEXT_QUEUE_ATTR = "AUDIT_LOG_CONTEXT_QUEUE";
 
   private final boolean internalAuthEnabled;
   private final boolean auditLoggingEnabled;
@@ -133,7 +135,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
       SecurityRequestContext.setUserCredential(currentUserCredential);
       SecurityRequestContext.setUserIp(currentUserIp);
       //Also set userIp in ATTR , to be used in audit logging incase it was replaced at a later stage
-      ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.USER_IP)).set(currentUserIp);
+      ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_USER_IP_ATTR)).set(currentUserIp);
     }
 
     try {
@@ -190,34 +192,41 @@ private void publishAuditLogRequest(ChannelHandlerContext ctx) throws IOExceptio
   private AuditLogRequest getAuditLogRequest(ChannelHandlerContext ctx) {
 
     Object auditLogContextsQueueAttr =
-      ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.AUDIT_LOG_CONTEXT_QUEUE)).get();
+      ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_CONTEXT_QUEUE_ATTR)).get();
 
     // If NO audit logs, then return NULL.
     if (auditLogContextsQueueAttr == null) {
       return null;
     }
 
-    Object userIpObj = ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.USER_IP)).get();
+    Queue<AuditLogContext> auditLogContextsQueue =  (Queue<AuditLogContext>)auditLogContextsQueueAttr;
+
+    Object userIpObj = ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_USER_IP_ATTR)).get();
     String userIp = userIpObj == null ? SecurityRequestContext.getUserIp() : (String) userIpObj;
 
-    // Check Attr for Audit Metadata
-    Object auditMetadataMapObj = ctx.channel().attr(AttributeKey.valueOf(AUDIT_METADATA_MAP_ATTR_NAME)).get();
-    Map<String, Object> auditMetadataMap = auditMetadataMapObj == null ? SecurityRequestContext.getAuditMetaDataMap() :
-      (Map<String, Object>) auditMetadataMapObj;
+    // Check Attr for AuditLogRequest Builder.
+    Object builderObj = ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_REQ_BUILDER_ATTR)).get();
+    AuditLogRequest.Builder builder = builderObj == null ? SecurityRequestContext.getAuditLogRequestBuilder() :
+      (AuditLogRequest.Builder) builderObj;
+
+    //If the AuditLogContextsQueue is NOT empty, then the AuditLogRequest.Builder should be NEVER be null.
+    //It should either come from ATTR of pipeline or SecurityRequestContext.
+    //Ideally we will never encounter this.
+    if (builder == null) {
+      LOG.error("The meta data required to publish audit logs are missing or null. Following Audit logs will be " +
+                  "skipped: {}", auditLogContextsQueue.stream().map(String::valueOf)
+                  .collect(Collectors.joining(", ")));
+      //Returning null as Operation is already completed.
+      return null;
+    }
 
     //Clear Attributes
     clearAttributes(ctx);
 
-    return new AuditLogRequest(
-      (Integer) auditMetadataMap.get(AuditLogRequest.PropKey.OP_RESP_CODE),
-      userIp,
-      (String) auditMetadataMap.get(AuditLogRequest.PropKey.URI),
-      (String) auditMetadataMap.get(AuditLogRequest.PropKey.HANDLER),
-      (String) auditMetadataMap.get(AuditLogRequest.PropKey.METHOD),
-      (String) auditMetadataMap.get(AuditLogRequest.PropKey.METHOD_TYPE),
-      (Queue<AuditLogContext>) auditLogContextsQueueAttr,
-      (Long) auditMetadataMap.get(AuditLogRequest.PropKey.START_TIME_NANOS),
-      (Long) auditMetadataMap.get(AuditLogRequest.PropKey.END_TIME_NANOS));
+    return builder
+      .userIp(userIp)
+      .auditLogContextQueue(auditLogContextsQueue)
+      .build();
   }
 
   /**
@@ -231,22 +240,22 @@ private void setAuditLogMetaDataInChannel(ChannelHandlerContext ctx) {
 
     // In either case, if Queue from SecurityRequestContext is empty, then no Audit Logs.
     if (!auditLogContextQueue.isEmpty()){
-      ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.AUDIT_LOG_CONTEXT_QUEUE))
+      ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_CONTEXT_QUEUE_ATTR))
         .set(auditLogContextQueue);
     }
 
-    // Store all audit metadata info in ATTR from  AuditLogSetterHook#postCall
+    // Store all audit metadata info stored in AuditLogRequest.Builder in ATTR from  AuditLogSetterHook#postCall
     // This is just to ensure we don't lose metadata information if already populated because of some RESET call on
     // SecurityRequestContext. This Also ensures that if Thread changes in Close / Write , then this info is preserved.
-    Map<String, Object> auditMetaDataMap = SecurityRequestContext.getAuditMetaDataMap();
-    if (!auditMetaDataMap.isEmpty()) {
-      ctx.channel().attr(AttributeKey.valueOf(AUDIT_METADATA_MAP_ATTR_NAME)).set(auditMetaDataMap);
+    AuditLogRequest.Builder builder = SecurityRequestContext.getAuditLogRequestBuilder();
+    if (builder != null) {
+      ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_REQ_BUILDER_ATTR)).set(builder);
     }
   }
 
   private void clearAttributes(ChannelHandlerContext ctx){
-    ctx.channel().attr(AttributeKey.valueOf(AUDIT_METADATA_MAP_ATTR_NAME)).set(null);
-    ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.USER_IP)).set(null);
-    ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.AUDIT_LOG_CONTEXT_QUEUE)).set(null);
+    ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_REQ_BUILDER_ATTR)).set(null);
+    ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_USER_IP_ATTR)).set(null);
+    ctx.channel().attr(AttributeKey.valueOf(AUDIT_LOG_CONTEXT_QUEUE_ATTR)).set(null);
   }
 }

cdap-common/src/test/java/io/cdap/cdap/common/http/AuthenticationChannelHandlerTest.java

@@ -39,16 +39,13 @@
 import org.mockito.Mockito;
 
 import java.util.ArrayDeque;
-import java.util.HashMap;
-import java.util.Map;
 import java.util.Queue;
 
 public class AuthenticationChannelHandlerTest {
 
   private DefaultHttpRequest req;
   private AuthenticationChannelHandler handler;
   private ChannelHandlerContext ctx;
-  private static final String AUDIT_METADATA_MAP_ATTR_NAME = "AUDIT_LOG_METADATA_MAP";
 
 
   @Before
@@ -108,12 +105,12 @@ public void testWellFormedValidCredentialCallsFireChannelReader() throws Excepti
   public void testWriteWithAuditLogging() throws Exception {
     boolean internalAuthEnabled = true;
     AuditLogWriter auditLogWriterMock = Mockito.mock(AuditLogWriter.class);
-    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.AUDIT_LOG_CONTEXT_QUEUE)).get())
-      .thenReturn(getAuditLogContexts());
-    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.USER_IP)).get())
+    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuthenticationChannelHandler.AUDIT_LOG_CONTEXT_QUEUE_ATTR))
+                   .get()).thenReturn(getAuditLogContexts());
+    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuthenticationChannelHandler.AUDIT_LOG_USER_IP_ATTR)).get())
       .thenReturn("testuserIp");
-    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AUDIT_METADATA_MAP_ATTR_NAME)).get())
-      .thenReturn(getAuditLogMetaMap());
+    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuthenticationChannelHandler.AUDIT_LOG_REQ_BUILDER_ATTR))
+                   .get()).thenReturn(getAuditLogRequestBuilder());
     handler = new AuthenticationChannelHandler(internalAuthEnabled, true, auditLogWriterMock);
     handler.write(ctx, "msg", new DefaultChannelPromise(ctx.channel()));
 
@@ -124,12 +121,12 @@ public void testWriteWithAuditLogging() throws Exception {
   public void testCloseWithAuditLogging() throws Exception {
     boolean internalAuthEnabled = true;
     AuditLogWriter auditLogWriterMock = Mockito.mock(AuditLogWriter.class);
-    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.AUDIT_LOG_CONTEXT_QUEUE)).get())
-      .thenReturn(getAuditLogContexts());
-    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuditLogRequest.PropKey.USER_IP)).get())
+    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuthenticationChannelHandler.AUDIT_LOG_CONTEXT_QUEUE_ATTR))
+                   .get()).thenReturn(getAuditLogContexts());
+    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuthenticationChannelHandler.AUDIT_LOG_USER_IP_ATTR)).get())
       .thenReturn("testuserIp");
-    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AUDIT_METADATA_MAP_ATTR_NAME)).get())
-      .thenReturn(getAuditLogMetaMap());
+    Mockito.when(ctx.channel().attr(AttributeKey.valueOf(AuthenticationChannelHandler.AUDIT_LOG_REQ_BUILDER_ATTR))
+                   .get()).thenReturn(getAuditLogRequestBuilder());
     handler = new AuthenticationChannelHandler(internalAuthEnabled, true, auditLogWriterMock);
     handler.close(ctx, new DefaultChannelPromise(ctx.channel()));
 
@@ -145,17 +142,14 @@ private Queue<AuditLogContext> getAuditLogContexts(){
                            .build());
     return auditLogContexts;
   }
-  private Map<String, Object> getAuditLogMetaMap() {
-
-    Map<String, Object> auditMetadataMap = new HashMap<>();
-    auditMetadataMap.put(AuditLogRequest.PropKey.OP_RESP_CODE, 200);
-    auditMetadataMap.put(AuditLogRequest.PropKey.URI, "v3/test");
-    auditMetadataMap.put(AuditLogRequest.PropKey.HANDLER, "Testhandler");
-    auditMetadataMap.put(AuditLogRequest.PropKey.METHOD, "create");
-    auditMetadataMap.put(AuditLogRequest.PropKey.METHOD_TYPE, "POST");
-    auditMetadataMap.put(AuditLogRequest.PropKey.START_TIME_NANOS, 1000000L);
-    auditMetadataMap.put(AuditLogRequest.PropKey.END_TIME_NANOS, 1000002L);
-
-    return auditMetadataMap;
+  private AuditLogRequest.Builder getAuditLogRequestBuilder() {
+    return new AuditLogRequest.Builder()
+      .operationResponseCode(200)
+      .uri("v3/test")
+      .handler("Testhandler")
+      .method("create")
+      .methodType("POST")
+      .startTimeNanos(1000000L)
+      .endTimeNanos(1000002L);
   }
 }

cdap-security-spi/src/main/java/io/cdap/cdap/security/spi/authentication/SecurityRequestContext.java

@@ -20,6 +20,7 @@
 import io.cdap.cdap.proto.security.Credential;
 import io.cdap.cdap.proto.security.Principal;
 import io.cdap.cdap.security.spi.authorization.AuditLogContext;
+import io.cdap.cdap.security.spi.authorization.AuditLogRequest;
 import io.cdap.cdap.security.spi.authorization.AuthorizationResponse;
 
 import java.util.ArrayDeque;
@@ -41,7 +42,7 @@ public final class SecurityRequestContext {
   private static final ThreadLocal<Queue<AuditLogContext>> auditLogContextQueue = new InheritableThreadLocal<>();
   private static final ThreadLocal<Map<? extends EntityId, AuthorizationResponse>> entityToAuthResponseMap =
     new InheritableThreadLocal<>();
-  private static final ThreadLocal<Map<String, Object>> auditMetaDataMap = new InheritableThreadLocal<>();
+  private static final ThreadLocal<AuditLogRequest.Builder> auditLogRequestBuilder = new InheritableThreadLocal<>();
 
   private SecurityRequestContext() {
   }
@@ -119,7 +120,7 @@ public static void reset() {
     userCredential.remove();
     auditLogContextQueue.remove();
     entityToAuthResponseMap.remove();
-    auditMetaDataMap.remove();
+    auditLogRequestBuilder.remove();
   }
 
   /**
@@ -202,25 +203,17 @@ public static void clearEntityToAuthResponseMap() {
   }
 
   /**
-   *  Set a property of the Map of String <> Object for audit logging.
+   *  Set the Builder for AuditLogRequest.
    */
-  public static void setAuditMetaDataInMap(String key, Object val) {
-    Map<String, Object> map = auditMetaDataMap.get();
-    if (map == null) {
-      map = new HashMap<>();
-    }
-    map.put(key, val);
-    auditMetaDataMap.set(map);
+  public static void setAuditLogRequestBuilder(AuditLogRequest.Builder builder) {
+    auditLogRequestBuilder.set(builder);
   }
 
   /**
-   *  Get the Map of EntityId <> AuthorizationResponse.
+   *  Get the Builder for AuditLogRequest.
    */
-  public static Map<String, Object> getAuditMetaDataMap() {
-    Map<String, Object> map = auditMetaDataMap.get();
-    if (map != null) {
-      return map;
-    }
-    return new HashMap<>();
+  @Nullable
+  public static AuditLogRequest.Builder getAuditLogRequestBuilder() {
+    return auditLogRequestBuilder.get();
   }
 }

cdap-security-spi/src/main/java/io/cdap/cdap/security/spi/authorization/AuditLogContext.java

@@ -103,4 +103,12 @@ public boolean equals(Object o) {
   public int hashCode() {
     return Objects.hash(auditLoggingRequired, auditLogBody);
   }
+
+  @Override
+  public String toString() {
+    return "AuditLogContext{" +
+      "auditLoggingRequired=" + auditLoggingRequired +
+      ", auditLogBody='" + auditLogBody + '\'' +
+      '}';
+  }
 }

cdap-security-spi/src/main/java/io/cdap/cdap/security/spi/authorization/AuditLogRequest.java

@@ -86,15 +86,65 @@ public long getEndTimeNanos() {
     return endTimeNanos;
   }
 
-  public static final class PropKey {
-    public static final String OP_RESP_CODE = "operationResponseCode";
-    public static final String USER_IP = "userIp";
-    public static final String URI = "uri";
-    public static final String HANDLER = "handler";
-    public static final String METHOD = "method";
-    public static final String METHOD_TYPE = "methodType";
-    public static final String AUDIT_LOG_CONTEXT_QUEUE = "auditLogContextQueue";
-    public static final String START_TIME_NANOS = "startTimeNanos";
-    public static final String END_TIME_NANOS = "endTimeNanos";
+  public static class Builder {
+    private int operationResponseCode;
+    private String userIp;
+    private String uri;
+    private String handler;
+    private String method;
+    private String methodType;
+    private Queue<AuditLogContext> auditLogContextQueue;
+    private Long startTimeNanos;
+    private Long endTimeNanos;
+
+    public Builder operationResponseCode(int operationResponseCode) {
+      this.operationResponseCode = operationResponseCode;
+      return this;
+    }
+
+    public Builder userIp(String userIp) {
+      this.userIp = userIp;
+      return this;
+    }
+
+    public Builder uri(String uri) {
+      this.uri = uri;
+      return this;
+    }
+
+    public Builder handler(String handler) {
+      this.handler = handler;
+      return this;
+    }
+
+    public Builder method(String method) {
+      this.method = method;
+      return this;
+    }
+
+    public Builder methodType(String methodType) {
+      this.methodType = methodType;
+      return this;
+    }
+
+    public Builder auditLogContextQueue(Queue<AuditLogContext> auditLogContextQueue) {
+      this.auditLogContextQueue = auditLogContextQueue;
+      return this;
+    }
+
+    public Builder startTimeNanos(Long startTimeNanos) {
+      this.startTimeNanos = startTimeNanos;
+      return this;
+    }
+
+    public Builder endTimeNanos(Long endTimeNanos) {
+      this.endTimeNanos = endTimeNanos;
+      return this;
+    }
+
+    public AuditLogRequest build() {
+      return new AuditLogRequest(operationResponseCode, userIp, uri, handler,  method, methodType, auditLogContextQueue,
+                                 startTimeNanos, endTimeNanos);
+    }
   }
 }