Add security check

Author: smalyshev

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlAsyncStopAction.java

@@ -7,10 +7,12 @@
 
 package org.elasticsearch.xpack.esql.plugin;
 
+import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionListenerResponseHandler;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
+import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
@@ -19,14 +21,22 @@
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.xpack.core.XPackPlugin;
 import org.elasticsearch.xpack.core.async.AsyncExecutionId;
+import org.elasticsearch.xpack.core.async.AsyncSearchSecurity;
 import org.elasticsearch.xpack.core.async.AsyncStopRequest;
 import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
+import org.elasticsearch.xpack.core.security.SecurityContext;
 import org.elasticsearch.xpack.esql.action.EsqlAsyncStopAction;
 import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
+import org.elasticsearch.xpack.esql.action.EsqlQueryTask;
 
+import java.io.IOException;
 import java.util.concurrent.TimeUnit;
 
+import static org.elasticsearch.xpack.core.ClientHelper.ASYNC_SEARCH_ORIGIN;
+import static org.elasticsearch.xpack.core.async.AsyncTaskIndexService.getTask;
+
 /**
  * This action will stop running async request and collect the results.
  * If the request is already finished, it will do the same thing as the regular async get.
@@ -38,6 +48,7 @@ public class TransportEsqlAsyncStopAction extends HandledTransportAction<AsyncSt
     private final BlockFactory blockFactory;
     private final ClusterService clusterService;
     private final TransportService transportService;
+    private final AsyncSearchSecurity security;
 
     @Inject
     public TransportEsqlAsyncStopAction(
@@ -46,6 +57,7 @@ public TransportEsqlAsyncStopAction(
         ActionFilters actionFilters,
         TransportEsqlQueryAction queryAction,
         TransportEsqlAsyncGetResultsAction getResultsAction,
+        Client client,
         BlockFactory blockFactory
     ) {
         super(EsqlAsyncStopAction.NAME, transportService, actionFilters, AsyncStopRequest::new, EsExecutors.DIRECT_EXECUTOR_SERVICE);
@@ -54,6 +66,12 @@ public TransportEsqlAsyncStopAction(
         this.blockFactory = blockFactory;
         this.transportService = transportService;
         this.clusterService = clusterService;
+        this.security = new AsyncSearchSecurity(
+            XPackPlugin.ASYNC_RESULTS_INDEX,
+            new SecurityContext(clusterService.getSettings(), client.threadPool().getThreadContext()),
+            client,
+            ASYNC_SEARCH_ORIGIN
+        );
     }
 
     @Override
@@ -63,7 +81,7 @@ protected void doExecute(Task task, AsyncStopRequest request, ActionListener<Esq
         if (clusterService.localNode().getId().equals(searchId.getTaskId().getNodeId()) || node == null) {
             // Don't use original request ID here because base64 decoding may not need some padding, but we want to match the original ID
             // for the map lookup
-            stopQueryAndReturnResult(task, searchId.getEncoded(), listener);
+            stopQueryAndReturnResult(task, searchId, listener);
         } else {
             transportService.sendRequest(
                 node,
@@ -74,18 +92,27 @@ protected void doExecute(Task task, AsyncStopRequest request, ActionListener<Esq
         }
     }
 
-    private void stopQueryAndReturnResult(Task task, String asyncId, ActionListener<EsqlQueryResponse> listener) {
-        var asyncListener = queryAction.getAsyncListener(asyncId);
+    private void stopQueryAndReturnResult(Task task, AsyncExecutionId asyncId, ActionListener<EsqlQueryResponse> listener) {
+        String asyncIdStr = asyncId.getEncoded();
+        var asyncListener = queryAction.getAsyncListener(asyncIdStr);
         if (asyncListener == null) {
             // This should mean one of the two things: either bad request ID, or the query has already finished
             // In both cases, let regular async get deal with it.
-            var getAsyncResultRequest = new GetAsyncResultRequest(asyncId);
+            var getAsyncResultRequest = new GetAsyncResultRequest(asyncIdStr);
             // TODO: this should not be happening, but if the listener is not registered and the query is not finished,
             // we give it some time to finish
             getAsyncResultRequest.setWaitForCompletionTimeout(new TimeValue(1, TimeUnit.SECONDS));
             getResultsAction.execute(task, getAsyncResultRequest, listener);
             return;
         }
+        try {
+            EsqlQueryTask asyncTask = getTask(taskManager, asyncId, EsqlQueryTask.class);
+            if (false == security.currentUserHasAccessToTask(asyncTask)) {
+                throw new ResourceNotFoundException(asyncId + " not found");
+            }
+        } catch (IOException e) {
+            throw new ResourceNotFoundException(asyncId + " not found", e);
+        }
         asyncListener.addListener(listener);
         // TODO: send the finish signal to the source
     }