Use Escalator to determine system identity

kfaraz · Dec 18, 2023 · 0b73b5f · 0b73b5f
1 parent ed22509
commit 0b73b5f
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 21 deletions.
diff --git a/processing/src/main/java/org/apache/druid/audit/AuditManager.java b/processing/src/main/java/org/apache/druid/audit/AuditManager.java
@@ -32,20 +32,6 @@ public interface AuditManager
   String X_DRUID_AUTHOR = "X-Druid-Author";
   String X_DRUID_COMMENT = "X-Druid-Comment";
 
-  /**
-   * Value of header {@link #X_DRUID_AUTHOR} used by Druid services so that they
-   * can be distinguished from external requests.
-   */
-  String AUTHOR_DRUID_SYSTEM = "druid_system";
-
-  /**
-   * @return true if the audited event was initiated by the Druid system itself.
-   */
-  default boolean isSystemRequest(AuditInfo auditInfo)
-  {
-    return AUTHOR_DRUID_SYSTEM.equals(auditInfo.getAuthor());
-  }
-
   void doAudit(AuditEntry event);
 
   /**

diff --git a/server/src/main/java/org/apache/druid/rpc/indexing/OverlordClientImpl.java b/server/src/main/java/org/apache/druid/rpc/indexing/OverlordClientImpl.java
@@ -24,7 +24,6 @@
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.audit.AuditManager;
 import org.apache.druid.client.JsonParserIterator;
 import org.apache.druid.client.indexing.IndexingTotalWorkerCapacityInfo;
 import org.apache.druid.client.indexing.IndexingWorkerInfo;
@@ -97,8 +96,7 @@ public ListenableFuture<Void> runTask(final String taskId, final Object taskObje
     return FutureUtils.transform(
         client.asyncRequest(
             new RequestBuilder(HttpMethod.POST, "/druid/indexer/v1/task")
-                .jsonContent(jsonMapper, taskObject)
-                .header(AuditManager.X_DRUID_AUTHOR, AuditManager.AUTHOR_DRUID_SYSTEM),
+                .jsonContent(jsonMapper, taskObject),
             new BytesFullResponseHandler()
         ),
         holder -> {

diff --git a/server/src/main/java/org/apache/druid/server/audit/AuditSerdeHelper.java b/server/src/main/java/org/apache/druid/server/audit/AuditSerdeHelper.java
@@ -26,6 +26,7 @@
 import org.apache.druid.guice.annotations.JsonNonNull;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.server.security.Escalator;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -50,18 +51,35 @@ public class AuditSerdeHelper
   private final ObjectMapper jsonMapper;
   private final ObjectMapper jsonMapperSkipNulls;
 
+  private final String systemIdentity;
   private final AuditManagerConfig config;
 
   @Inject
   public AuditSerdeHelper(
       AuditManagerConfig config,
+      Escalator escalator,
       @Json ObjectMapper jsonMapper,
       @JsonNonNull ObjectMapper jsonMapperSkipNulls
   )
   {
     this.config = config;
     this.jsonMapper = jsonMapper;
     this.jsonMapperSkipNulls = jsonMapperSkipNulls;
+    this.systemIdentity = escalator == null
+                          ? null : escalator.createEscalatedAuthenticationResult().getIdentity();
+  }
+
+  /**
+   * Checks if the given audit event needs to be handled.
+   *
+   * @return true only if the event was not initiated by the Druid system OR if
+   * system requests should be audited too.
+   */
+  public boolean shouldProcessAuditEntry(AuditEntry entry)
+  {
+    final boolean isSystemRequest = systemIdentity != null
+                                    && systemIdentity.equals(entry.getAuditInfo().getIdentity());
+    return config.isAuditSystemRequests() || !isSystemRequest;
   }
 
   /**

diff --git a/server/src/main/java/org/apache/druid/server/audit/LoggingAuditManager.java b/server/src/main/java/org/apache/druid/server/audit/LoggingAuditManager.java
@@ -56,7 +56,7 @@ public LoggingAuditManager(
   @Override
   public void doAudit(AuditEntry entry)
   {
-    if (managerConfig.isAuditSystemRequests() || !isSystemRequest(entry.getAuditInfo())) {
+    if (serdeHelper.shouldProcessAuditEntry(entry)) {
       auditLogger.log(serdeHelper.processAuditEntry(entry));
     }
   }

diff --git a/server/src/main/java/org/apache/druid/server/audit/SQLAuditManager.java b/server/src/main/java/org/apache/druid/server/audit/SQLAuditManager.java
@@ -136,7 +136,7 @@ private ServiceMetricEvent.Builder createMetricEventBuilder(AuditEntry entry)
   @Override
   public void doAudit(AuditEntry event, Handle handle) throws IOException
   {
-    if (isSystemRequest(event.getAuditInfo()) && !config.isAuditSystemRequests()) {
+    if (!serdeHelper.shouldProcessAuditEntry(event)) {
       return;
     }
 

diff --git a/server/src/test/java/org/apache/druid/metadata/SQLMetadataRuleManagerTest.java b/server/src/test/java/org/apache/druid/metadata/SQLMetadataRuleManagerTest.java
@@ -75,7 +75,7 @@ public void setUp()
     final SQLAuditManagerConfig auditManagerConfig = new SQLAuditManagerConfig(null, null, null, null, null);
     auditManager = new SQLAuditManager(
         auditManagerConfig,
-        new AuditSerdeHelper(auditManagerConfig, mapper, mapper),
+        new AuditSerdeHelper(auditManagerConfig, null, mapper, mapper),
         connector,
         Suppliers.ofInstance(tablesConfig),
         new NoopServiceEmitter(),

diff --git a/server/src/test/java/org/apache/druid/server/audit/SQLAuditManagerTest.java b/server/src/test/java/org/apache/druid/server/audit/SQLAuditManagerTest.java
@@ -73,7 +73,7 @@ private SQLAuditManager createAuditManager(SQLAuditManagerConfig config)
   {
     return new SQLAuditManager(
         config,
-        new AuditSerdeHelper(config, mapper, mapperSkipNull),
+        new AuditSerdeHelper(config, null, mapper, mapperSkipNull),
         connector,
         derbyConnectorRule.metadataTablesConfigSupplier(),
         serviceEmitter,