Add more tests for coverage

kfaraz · Dec 19, 2023 · 33d8a22 · 33d8a22
1 parent c87fb13
commit 33d8a22
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 34 deletions.
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java
@@ -140,7 +140,7 @@ public class OverlordResource
   private AtomicReference<WorkerBehaviorConfig> workerConfigRef = null;
   private static final List<String> API_TASK_STATES = ImmutableList.of("pending", "waiting", "running", "complete");
   private static final Set<String> AUDITED_TASK_TYPES
-      = ImmutableSet.of("index", "index_parallel", "compact", "index_hadoop");
+      = ImmutableSet.of("index", "index_parallel", "compact", "index_hadoop", "kill");
 
   private enum TaskStateLookup
   {

diff --git a/...g-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java b/...g-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java
@@ -25,13 +25,16 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import org.apache.druid.audit.AuditEntry;
+import org.apache.druid.audit.AuditManager;
 import org.apache.druid.common.config.JacksonConfigManager;
 import org.apache.druid.indexer.RunnerTaskState;
 import org.apache.druid.indexer.TaskInfo;
 import org.apache.druid.indexer.TaskLocation;
 import org.apache.druid.indexer.TaskState;
 import org.apache.druid.indexer.TaskStatus;
 import org.apache.druid.indexer.TaskStatusPlus;
+import org.apache.druid.indexing.common.task.KillUnusedSegmentsTask;
 import org.apache.druid.indexing.common.task.NoopTask;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.ImmutableWorkerInfo;
@@ -67,6 +70,7 @@
 import org.apache.druid.server.security.Resource;
 import org.apache.druid.server.security.ResourceAction;
 import org.apache.druid.server.security.ResourceType;
+import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
 import org.joda.time.DateTime;
@@ -104,7 +108,9 @@ public class OverlordResourceTest
   private IndexerMetadataStorageAdapter indexerMetadataStorageAdapter;
   private HttpServletRequest req;
   private TaskRunner taskRunner;
+  private TaskQueue taskQueue;
   private WorkerTaskRunnerQueryAdapter workerTaskRunnerQueryAdapter;
+  private AuditManager auditManager;
 
   @Rule
   public ExpectedException expectedException = ExpectedException.none();
@@ -113,6 +119,7 @@ public class OverlordResourceTest
   public void setUp()
   {
     taskRunner = EasyMock.createMock(TaskRunner.class);
+    taskQueue = EasyMock.createMock(TaskQueue.class);
     configManager = EasyMock.createMock(JacksonConfigManager.class);
     provisioningStrategy = EasyMock.createMock(ProvisioningStrategy.class);
     authConfig = EasyMock.createMock(AuthConfig.class);
@@ -121,6 +128,7 @@ public void setUp()
     indexerMetadataStorageAdapter = EasyMock.createStrictMock(IndexerMetadataStorageAdapter.class);
     req = EasyMock.createStrictMock(HttpServletRequest.class);
     workerTaskRunnerQueryAdapter = EasyMock.createStrictMock(WorkerTaskRunnerQueryAdapter.class);
+    auditManager = EasyMock.createMock(AuditManager.class);
 
     EasyMock.expect(taskMaster.getTaskRunner()).andReturn(
         Optional.of(taskRunner)
@@ -165,7 +173,7 @@ public Access authorize(AuthenticationResult authenticationResult, Resource reso
         indexerMetadataStorageAdapter,
         null,
         configManager,
-        null,
+        auditManager,
         authMapper,
         workerTaskRunnerQueryAdapter,
         provisioningStrategy,
@@ -880,6 +888,53 @@ public void testSecuredTaskPost()
     overlordResource.taskPost(task, req);
   }
 
+  @Test
+  public void testKillTaskIsAudited()
+  {
+    EasyMock.expect(authConfig.isEnableInputSourceSecurity()).andReturn(false);
+
+    final String username = Users.DRUID;
+    expectAuthorizationTokenCheck(username);
+    EasyMock.expect(req.getMethod()).andReturn("POST").once();
+    EasyMock.expect(req.getRequestURI()).andReturn("/indexer/v2/task").once();
+    EasyMock.expect(req.getQueryString()).andReturn("").once();
+    EasyMock.expect(req.getHeader(AuditManager.X_DRUID_AUTHOR)).andReturn(username).once();
+    EasyMock.expect(req.getHeader(AuditManager.X_DRUID_COMMENT)).andReturn("killing segments").once();
+    EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT))
+            .andReturn(new AuthenticationResult(username, "druid", null, null))
+            .once();
+    EasyMock.expect(req.getRemoteAddr()).andReturn("127.0.0.1").once();
+
+    EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
+    EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).once();
+
+    final Capture<AuditEntry> auditEntryCapture = EasyMock.newCapture();
+    auditManager.doAudit(EasyMock.capture(auditEntryCapture));
+    EasyMock.expectLastCall().once();
+
+    EasyMock.replay(
+        taskRunner,
+        taskMaster,
+        taskQueue,
+        taskStorageQueryAdapter,
+        indexerMetadataStorageAdapter,
+        req,
+        workerTaskRunnerQueryAdapter,
+        authConfig,
+        auditManager
+    );
+
+    Task task = new KillUnusedSegmentsTask("kill_all", "allow", Intervals.ETERNITY, null, false, 10, null);
+    overlordResource.taskPost(task, req);
+
+    Assert.assertTrue(auditEntryCapture.hasCaptured());
+    AuditEntry auditEntry = auditEntryCapture.getValue();
+    Assert.assertEquals(username, auditEntry.getAuditInfo().getAuthor());
+    Assert.assertEquals("killing segments", auditEntry.getAuditInfo().getComment());
+    Assert.assertEquals("druid", auditEntry.getAuditInfo().getIdentity());
+    Assert.assertEquals("127.0.0.1", auditEntry.getAuditInfo().getIp());
+  }
+
   @Test
   public void testTaskPostDeniesDatasourceReadUser()
   {

diff --git a/server/src/main/java/org/apache/druid/guice/SQLMetadataStorageDruidModule.java b/server/src/main/java/org/apache/druid/guice/SQLMetadataStorageDruidModule.java
@@ -22,8 +22,6 @@
 import com.google.inject.Binder;
 import com.google.inject.Key;
 import com.google.inject.Module;
-import com.google.inject.multibindings.MapBinder;
-import org.apache.druid.audit.AuditManager;
 import org.apache.druid.indexer.MetadataStorageUpdaterJobHandler;
 import org.apache.druid.indexer.SQLMetadataStorageUpdaterJobHandler;
 import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator;
@@ -46,10 +44,6 @@
 import org.apache.druid.metadata.SegmentsMetadataManagerProvider;
 import org.apache.druid.metadata.SqlSegmentsMetadataManager;
 import org.apache.druid.metadata.SqlSegmentsMetadataManagerProvider;
-import org.apache.druid.server.audit.AuditManagerConfig;
-import org.apache.druid.server.audit.AuditSerdeHelper;
-import org.apache.druid.server.audit.LoggingAuditManager;
-import org.apache.druid.server.audit.SQLAuditManager;
 
 public class SQLMetadataStorageDruidModule implements Module
 {
@@ -84,8 +78,6 @@ public void createBindingChoices(Binder binder, String defaultValue)
     PolyBind.createChoiceWithDefault(binder, prop, Key.get(MetadataStorageActionHandlerFactory.class), defaultValue);
     PolyBind.createChoiceWithDefault(binder, prop, Key.get(MetadataStorageUpdaterJobHandler.class), defaultValue);
     PolyBind.createChoiceWithDefault(binder, prop, Key.get(MetadataSupervisorManager.class), defaultValue);
-
-    configureAuditManager(binder);
   }
 
   @Override
@@ -136,28 +128,4 @@ public void configure(Binder binder)
             .to(SQLMetadataSupervisorManager.class)
             .in(LazySingleton.class);
   }
-
-  private void configureAuditManager(Binder binder)
-  {
-    JsonConfigProvider.bind(binder, "druid.audit.manager", AuditManagerConfig.class);
-
-    PolyBind.createChoice(
-        binder,
-        "druid.audit.manager.type",
-        Key.get(AuditManager.class),
-        Key.get(SQLAuditManager.class)
-    );
-    final MapBinder<String, AuditManager> auditManagerBinder
-        = PolyBind.optionBinder(binder, Key.get(AuditManager.class));
-    auditManagerBinder
-        .addBinding("log")
-        .to(LoggingAuditManager.class)
-        .in(LazySingleton.class);
-    auditManagerBinder
-        .addBinding("sql")
-        .to(SQLAuditManager.class)
-        .in(LazySingleton.class);
-
-    binder.bind(AuditSerdeHelper.class).in(LazySingleton.class);
-  }
 }
diff --git a/server/src/main/java/org/apache/druid/guice/ServerModule.java b/server/src/main/java/org/apache/druid/guice/ServerModule.java
@@ -23,7 +23,10 @@
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Binder;
+import com.google.inject.Key;
 import com.google.inject.Provides;
+import com.google.inject.multibindings.MapBinder;
+import org.apache.druid.audit.AuditManager;
 import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.initialization.DruidModule;
 import org.apache.druid.jackson.DruidServiceSerializerModifier;
@@ -33,6 +36,10 @@
 import org.apache.druid.java.util.common.concurrent.ScheduledExecutors;
 import org.apache.druid.java.util.common.lifecycle.Lifecycle;
 import org.apache.druid.server.DruidNode;
+import org.apache.druid.server.audit.AuditManagerConfig;
+import org.apache.druid.server.audit.AuditSerdeHelper;
+import org.apache.druid.server.audit.LoggingAuditManager;
+import org.apache.druid.server.audit.SQLAuditManager;
 import org.apache.druid.server.initialization.ZkPathsConfig;
 
 import java.util.List;
@@ -48,6 +55,8 @@ public void configure(Binder binder)
   {
     JsonConfigProvider.bind(binder, ZK_PATHS_PROPERTY_BASE, ZkPathsConfig.class);
     JsonConfigProvider.bind(binder, "druid", DruidNode.class, Self.class);
+
+    configureAuditManager(binder);
   }
 
   @Provides @LazySingleton
@@ -65,4 +74,28 @@ public List<? extends Module> getJacksonModules()
             .setSerializerModifier(new DruidServiceSerializerModifier())
     );
   }
+
+  private void configureAuditManager(Binder binder)
+  {
+    JsonConfigProvider.bind(binder, "druid.audit.manager", AuditManagerConfig.class);
+
+    PolyBind.createChoice(
+        binder,
+        "druid.audit.manager.type",
+        Key.get(AuditManager.class),
+        Key.get(SQLAuditManager.class)
+    );
+    final MapBinder<String, AuditManager> auditManagerBinder
+        = PolyBind.optionBinder(binder, Key.get(AuditManager.class));
+    auditManagerBinder
+        .addBinding("log")
+        .to(LoggingAuditManager.class)
+        .in(LazySingleton.class);
+    auditManagerBinder
+        .addBinding("sql")
+        .to(SQLAuditManager.class)
+        .in(LazySingleton.class);
+
+    binder.bind(AuditSerdeHelper.class).in(LazySingleton.class);
+  }
 }