make threat intel async (#703)

Signed-off-by: Subhobrata Dey <[email protected]>
(cherry picked from commit 0dd9787)

src/main/java/org/opensearch/securityanalytics/threatIntel/DetectorThreatIntelService.java

@@ -158,7 +158,6 @@ private static String constructId(Detector detector, String iocType) {
 
     /** Updates all detectors having threat intel detection enabled with the latest threat intel feed data*/
     public void updateDetectorsWithLatestThreatIntelRules() {
-        try {
             QueryBuilder queryBuilder =
                     QueryBuilders.nestedQuery("detector",
                             QueryBuilders.boolQuery().must(
@@ -168,7 +167,6 @@ public void updateDetectorsWithLatestThreatIntelRules() {
             SearchSourceBuilder ssb = searchRequest.source();
             ssb.query(queryBuilder);
             ssb.size(9999);
-            CountDownLatch countDownLatch = new CountDownLatch(1);
             client.execute(SearchDetectorAction.INSTANCE, new SearchDetectorRequest(searchRequest),
                     ActionListener.wrap(searchResponse -> {
                         List<Detector> detectors = getDetectors(searchResponse, xContentRegistry);
@@ -181,22 +179,15 @@ public void updateDetectorsWithLatestThreatIntelRules() {
                                             ActionListener.wrap(
                                                     indexDetectorResponse -> {
                                                         log.debug("updated {} with latest threat intel info", indexDetectorResponse.getDetector().getId());
-                                                        countDownLatch.countDown();
                                                     },
                                                     e -> {
                                                         log.error(() -> new ParameterizedMessage("Failed to update detector {} with latest threat intel info", detector.getId()), e);
-                                                        countDownLatch.countDown();
                                                     }));
                                 }
                         );
                     }, e -> {
                         log.error("Failed to fetch detectors to update with threat intel queries.", e);
-                        countDownLatch.countDown();
                     }));
-            countDownLatch.await(5, TimeUnit.MINUTES);
-        } catch (InterruptedException e) {
-            log.error("");
-        }
 
 
     }

src/main/java/org/opensearch/securityanalytics/threatIntel/ThreatIntelFeedDataService.java

@@ -9,12 +9,15 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchException;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.DocWriteRequest;
 import org.opensearch.action.admin.indices.create.CreateIndexRequest;
+import org.opensearch.action.admin.indices.create.CreateIndexResponse;
 import org.opensearch.action.bulk.BulkRequest;
 import org.opensearch.action.bulk.BulkResponse;
 import org.opensearch.action.index.IndexRequest;
 import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.support.GroupedActionListener;
 import org.opensearch.action.support.IndicesOptions;
 import org.opensearch.action.support.WriteRequest;
 import org.opensearch.action.support.master.AcknowledgedResponse;
@@ -33,6 +36,7 @@
 import org.opensearch.securityanalytics.model.ThreatIntelFeedData;
 import org.opensearch.securityanalytics.threatIntel.action.PutTIFJobAction;
 import org.opensearch.securityanalytics.threatIntel.action.PutTIFJobRequest;
+import org.opensearch.securityanalytics.threatIntel.action.ThreatIntelIndicesResponse;
 import org.opensearch.securityanalytics.threatIntel.common.TIFMetadata;
 import org.opensearch.securityanalytics.threatIntel.common.StashedThreadContext;
 import org.opensearch.securityanalytics.settings.SecurityAnalyticsSettings;
@@ -47,6 +51,7 @@
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -102,17 +107,17 @@ public void getThreatIntelFeedData(
 
             String tifdIndex = getLatestIndexByCreationDate();
             if (tifdIndex == null) {
-                createThreatIntelFeedData();
-                tifdIndex = getLatestIndexByCreationDate();
+                createThreatIntelFeedData(listener);
+            } else {
+                SearchRequest searchRequest = new SearchRequest(tifdIndex);
+                searchRequest.source().size(9999); //TODO: convert to scroll
+                String finalTifdIndex = tifdIndex;
+                client.search(searchRequest, ActionListener.wrap(r -> listener.onResponse(ThreatIntelFeedDataUtils.getTifdList(r, xContentRegistry)), e -> {
+                    log.error(String.format(
+                            "Failed to fetch threat intel feed data from system index %s", finalTifdIndex), e);
+                    listener.onFailure(e);
+                }));
             }
-            SearchRequest searchRequest = new SearchRequest(tifdIndex);
-            searchRequest.source().size(9999); //TODO: convert to scroll
-            String finalTifdIndex = tifdIndex;
-            client.search(searchRequest, ActionListener.wrap(r -> listener.onResponse(ThreatIntelFeedDataUtils.getTifdList(r, xContentRegistry)), e -> {
-                log.error(String.format(
-                        "Failed to fetch threat intel feed data from system index %s", finalTifdIndex), e);
-                listener.onFailure(e);
-            }));
         } catch (InterruptedException e) {
             log.error("Failed to get threat intel feed data", e);
             listener.onFailure(e);
@@ -136,15 +141,30 @@ private String getLatestIndexByCreationDate() {
      *
      * @param indexName index name
      */
-    public void createIndexIfNotExists(final String indexName) {
+    public void createIndexIfNotExists(final String indexName, final ActionListener<CreateIndexResponse> listener) {
         if (clusterService.state().metadata().hasIndex(indexName) == true) {
+            listener.onResponse(new CreateIndexResponse(true, true, indexName));
             return;
         }
         final CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName).settings(INDEX_SETTING_TO_CREATE)
-                .mapping(getIndexMapping());
+                .mapping(getIndexMapping()).timeout(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT));
         StashedThreadContext.run(
                 client,
-                () -> client.admin().indices().create(createIndexRequest).actionGet(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT))
+                () -> client.admin().indices().create(createIndexRequest, new ActionListener<>() {
+                    @Override
+                    public void onResponse(CreateIndexResponse response) {
+                        if (response.isAcknowledged()) {
+                            listener.onResponse(response);
+                        } else {
+                            onFailure(new OpenSearchStatusException("Threat intel feed index creation failed", RestStatus.INTERNAL_SERVER_ERROR));
+                        }
+                    }
+
+                    @Override
+                    public void onFailure(Exception e) {
+                        listener.onFailure(e);
+                    }
+                })
         );
     }
 
@@ -159,16 +179,20 @@ public void parseAndSaveThreatIntelFeedDataCSV(
             final String indexName,
             final Iterator<CSVRecord> iterator,
             final Runnable renewLock,
-            final TIFMetadata tifMetadata
+            final TIFMetadata tifMetadata,
+            final ActionListener<ThreatIntelIndicesResponse> listener
     ) throws IOException {
         if (indexName == null || iterator == null || renewLock == null) {
             throw new IllegalArgumentException("Parameters cannot be null, failed to save threat intel feed data");
         }
 
         TimeValue timeout = clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT);
         Integer batchSize = clusterSettings.get(SecurityAnalyticsSettings.BATCH_SIZE);
-        final BulkRequest bulkRequest = new BulkRequest();
+
+        List<BulkRequest> bulkRequestList = new ArrayList<>();
+        BulkRequest bulkRequest = new BulkRequest();
         bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+
         List<ThreatIntelFeedData> tifdList = new ArrayList<>();
         while (iterator.hasNext()) {
             CSVRecord record = iterator.next();
@@ -192,10 +216,39 @@ public void parseAndSaveThreatIntelFeedDataCSV(
             bulkRequest.add(indexRequest);
 
             if (bulkRequest.requests().size() == batchSize) {
-                saveTifds(bulkRequest, timeout);
+                bulkRequestList.add(bulkRequest);
+                bulkRequest = new BulkRequest();
+                bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+            }
+        }
+        bulkRequestList.add(bulkRequest);
+
+        GroupedActionListener<BulkResponse> bulkResponseListener = new GroupedActionListener<>(new ActionListener<>() {
+            @Override
+            public void onResponse(Collection<BulkResponse> bulkResponses) {
+                int idx = 0;
+                for (BulkResponse response: bulkResponses) {
+                    BulkRequest request = bulkRequestList.get(idx);
+                    if (response.hasFailures()) {
+                        throw new OpenSearchException(
+                                "error occurred while ingesting threat intel feed data in {} with an error {}",
+                                StringUtils.join(request.getIndices()),
+                                response.buildFailureMessage()
+                        );
+                    }
+                }
+                listener.onResponse(new ThreatIntelIndicesResponse(true, List.of(indexName)));
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                listener.onFailure(e);
             }
+        }, bulkRequestList.size());
+
+        for (int i = 0; i < bulkRequestList.size(); ++i) {
+            saveTifds(bulkRequestList.get(i), timeout, bulkResponseListener);
         }
-        saveTifds(bulkRequest, timeout);
         renewLock.run();
     }
 
@@ -206,19 +259,9 @@ public static boolean isValidIp(String ip) {
         return matcher.matches();
     }
 
-    public void saveTifds(BulkRequest bulkRequest, TimeValue timeout) {
+    public void saveTifds(BulkRequest bulkRequest, TimeValue timeout, ActionListener<BulkResponse> listener) {
         try {
-            BulkResponse response = StashedThreadContext.run(client, () -> {
-                return client.bulk(bulkRequest).actionGet(timeout);
-            });
-            if (response.hasFailures()) {
-                throw new OpenSearchException(
-                        "error occurred while ingesting threat intel feed data in {} with an error {}",
-                        StringUtils.join(bulkRequest.getIndices()),
-                        response.buildFailureMessage()
-                );
-            }
-            bulkRequest.requests().clear();
+            StashedThreadContext.run(client, () -> client.bulk(bulkRequest, listener));
         } catch (OpenSearchException e) {
             log.error("failed to save threat intel feed data", e);
         }
@@ -241,31 +284,49 @@ public void deleteThreatIntelDataIndex(final List<String> indices) {
             );
         }
 
-        AcknowledgedResponse response = StashedThreadContext.run(
+        StashedThreadContext.run(
                 client,
                 () -> client.admin()
                         .indices()
                         .prepareDelete(indices.toArray(new String[0]))
                         .setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED_HIDDEN)
-                        .execute()
-                        .actionGet(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT))
-        );
+                        .setTimeout(clusterSettings.get(SecurityAnalyticsSettings.THREAT_INTEL_TIMEOUT))
+                        .execute(new ActionListener<>() {
+                            @Override
+                            public void onResponse(AcknowledgedResponse response) {
+                                if (response.isAcknowledged() == false) {
+                                    onFailure(new OpenSearchException("failed to delete data[{}]", String.join(",", indices)));
+                                }
+                            }
 
-        if (response.isAcknowledged() == false) {
-            throw new OpenSearchException("failed to delete data[{}]", String.join(",", indices));
-        }
+                            @Override
+                            public void onFailure(Exception e) {
+                                log.error("unknown exception:", e);
+                            }
+                        })
+        );
     }
 
-    private void createThreatIntelFeedData() throws InterruptedException {
+    private void createThreatIntelFeedData(ActionListener<List<ThreatIntelFeedData>> listener) throws InterruptedException {
         CountDownLatch countDownLatch = new CountDownLatch(1);
         client.execute(
                 PutTIFJobAction.INSTANCE,
                 new PutTIFJobRequest("feed_updater", clusterSettings.get(SecurityAnalyticsSettings.TIF_UPDATE_INTERVAL)),
-                new ActionListener<AcknowledgedResponse>() {
+                new ActionListener<>() {
                     @Override
                     public void onResponse(AcknowledgedResponse acknowledgedResponse) {
                         log.debug("Acknowledged threat intel feed updater job created");
                         countDownLatch.countDown();
+                        String tifdIndex = getLatestIndexByCreationDate();
+
+                        SearchRequest searchRequest = new SearchRequest(tifdIndex);
+                        searchRequest.source().size(9999); //TODO: convert to scroll
+                        String finalTifdIndex = tifdIndex;
+                        client.search(searchRequest, ActionListener.wrap(r -> listener.onResponse(ThreatIntelFeedDataUtils.getTifdList(r, xContentRegistry)), e -> {
+                            log.error(String.format(
+                                    "Failed to fetch threat intel feed data from system index %s", finalTifdIndex), e);
+                            listener.onFailure(e);
+                        }));
                     }
 
                     @Override

src/main/java/org/opensearch/securityanalytics/threatIntel/action/ThreatIntelIndicesResponse.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.securityanalytics.threatIntel.action;
+
+import org.opensearch.core.action.ActionResponse;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+import java.util.List;
+
+public class ThreatIntelIndicesResponse extends ActionResponse {
+
+    private Boolean isAcknowledged;
+
+    private List<String> indices;
+
+    public ThreatIntelIndicesResponse(Boolean isAcknowledged, List<String> indices) {
+        super();
+        this.isAcknowledged = isAcknowledged;
+        this.indices = indices;
+    }
+
+    public ThreatIntelIndicesResponse(StreamInput sin) throws IOException {
+        this(sin.readBoolean(), sin.readStringList());
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeBoolean(isAcknowledged);
+        out.writeStringCollection(indices);
+    }
+
+    public Boolean isAcknowledged() {
+        return isAcknowledged;
+    }
+
+    public List<String> getIndices() {
+        return indices;
+    }
+}

src/main/java/org/opensearch/securityanalytics/threatIntel/action/TransportPutTIFJobAction.java

@@ -7,6 +7,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.ResourceAlreadyExistsException;
 import org.opensearch.action.StepListener;
 import org.opensearch.action.index.IndexResponse;
@@ -127,12 +128,22 @@ protected ActionListener<IndexResponse> postIndexingTifJobParameter(
             @Override
             public void onResponse(final IndexResponse indexResponse) {
                 AtomicReference<LockModel> lockReference = new AtomicReference<>(lock);
-                try {
-                    createThreatIntelFeedData(tifJobParameter, lockService.getRenewLockRunnable(lockReference));
-                } finally {
-                    lockService.releaseLock(lockReference.get());
-                }
-                listener.onResponse(new AcknowledgedResponse(true));
+                createThreatIntelFeedData(tifJobParameter, lockService.getRenewLockRunnable(lockReference), new ActionListener<>() {
+                    @Override
+                    public void onResponse(ThreatIntelIndicesResponse threatIntelIndicesResponse) {
+                        if (threatIntelIndicesResponse.isAcknowledged()) {
+                            lockService.releaseLock(lockReference.get());
+                            listener.onResponse(new AcknowledgedResponse(true));
+                        } else {
+                            onFailure(new OpenSearchStatusException("creation of threat intel feed data failed", RestStatus.INTERNAL_SERVER_ERROR));
+                        }
+                    }
+
+                    @Override
+                    public void onFailure(Exception e) {
+                        listener.onFailure(e);
+                    }
+                });
             }
 
             @Override
@@ -149,26 +160,26 @@ public void onFailure(final Exception e) {
         };
     }
 
-    protected void createThreatIntelFeedData(final TIFJobParameter tifJobParameter, final Runnable renewLock) {
+    protected void createThreatIntelFeedData(final TIFJobParameter tifJobParameter, final Runnable renewLock, final ActionListener<ThreatIntelIndicesResponse> listener) {
         if (TIFJobState.CREATING.equals(tifJobParameter.getState()) == false) {
             log.error("Invalid tifJobParameter state. Expecting {} but received {}", TIFJobState.CREATING, tifJobParameter.getState());
-            markTIFJobAsCreateFailed(tifJobParameter);
+            markTIFJobAsCreateFailed(tifJobParameter, listener);
             return;
         }
 
         try {
-            tifJobUpdateService.createThreatIntelFeedData(tifJobParameter, renewLock);
+            tifJobUpdateService.createThreatIntelFeedData(tifJobParameter, renewLock, listener);
         } catch (Exception e) {
             log.error("Failed to create tifJobParameter for {}", tifJobParameter.getName(), e);
-            markTIFJobAsCreateFailed(tifJobParameter);
+            markTIFJobAsCreateFailed(tifJobParameter, listener);
         }
     }
 
-    private void markTIFJobAsCreateFailed(final TIFJobParameter tifJobParameter) {
+    private void markTIFJobAsCreateFailed(final TIFJobParameter tifJobParameter, final ActionListener<ThreatIntelIndicesResponse> listener) {
         tifJobParameter.getUpdateStats().setLastFailedAt(Instant.now());
         tifJobParameter.setState(TIFJobState.CREATE_FAILED);
         try {
-            tifJobParameterService.updateJobSchedulerParameter(tifJobParameter);
+            tifJobParameterService.updateJobSchedulerParameter(tifJobParameter, listener);
         } catch (Exception e) {
             log.error("Failed to mark tifJobParameter state as CREATE_FAILED for {}", tifJobParameter.getName(), e);
         }