Repository: OpenAEV-Platform/openaev (java)

Issue:
My scenario can't run because injects stay in pending but it's executed
As you can see here sometimes into my scenario injects are successfully done but still pending and when I've got chaining that blocked my scenario


Discussion:

Patch:
diff --git a/openbas-api/src/main/java/io/openbas/rest/inject/InjectApi.java b/openbas-api/src/main/java/io/openbas/rest/inject/InjectApi.java
--- a/openbas-api/src/main/java/io/openbas/rest/inject/InjectApi.java
+++ b/openbas-api/src/main/java/io/openbas/rest/inject/InjectApi.java
@@ -325,6 +325,21 @@ public void injectExecutionCallback(
 
   @Secured(ROLE_ADMIN)
   @PostMapping(INJECT_URI + "/execution/{agentId}/callback/{injectId}")
+  @Operation(
+      summary = "Inject execution callback for implants",
+      description =
+          "This endpoint is invoked by implants to report the result of an inject execution. "
+              + "It is used to update the inject status and execution traces based on the implant's execution result."
+              + " If the requested action is 'complete', the inject must be in the 'PENDING' state. otherwise a 409"
+              + " is returned. This can sometimes happen if the payload executed by the implant was executed too quickly. ")
+  @ApiResponses(
+      value = {
+        @ApiResponse(responseCode = "200", description = "Execution callback was successful"),
+        @ApiResponse(
+            responseCode = "409",
+            description =
+                "The inject to update was not in a valid state in regards to the requested action. Retry in a few seconds."),
+      })
   public void injectExecutionCallback(
       @PathVariable
           String agentId, // must allow null because http injector used also this method to work.
diff --git a/openbas-api/src/main/java/io/openbas/rest/inject/service/ExecutableInjectService.java b/openbas-api/src/main/java/io/openbas/rest/inject/service/ExecutableInjectService.java
--- a/openbas-api/src/main/java/io/openbas/rest/inject/service/ExecutableInjectService.java
+++ b/openbas-api/src/main/java/io/openbas/rest/inject/service/ExecutableInjectService.java
@@ -13,6 +13,7 @@
 import io.openbas.rest.exception.ElementNotFoundException;
 import io.openbas.rest.payload.service.PayloadService;
 import jakarta.annotation.Resource;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Base64;
 import java.util.List;
@@ -171,9 +172,16 @@ private String processAndEncodeCommand(
 
   public Payload getExecutablePayloadAndUpdateInjectStatus(String injectId, String agentId)
       throws Exception {
+    // Need startTime to be defined before everything else to be the most accurate start time, as
+    // this whole process is
+    // called at the beginning of the implant execution. A better solution would be to have the
+    // implant send the start time
+    // but it would require more changes in the implant code and change this endpoint from a get to
+    // a post.
+    Instant startTime = Instant.now();
     Payload payloadToExecute = getExecutablePayloadInject(injectId);
     this.injectStatusService.addStartImplantExecutionTraceByInject(
-        injectId, agentId, "Implant is up and starting execution");
+        injectId, agentId, "Implant is up and starting execution", startTime);
     return payloadToExecute;
   }
 
diff --git a/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectExecutionService.java b/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectExecutionService.java
--- a/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectExecutionService.java
+++ b/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectExecutionService.java
@@ -10,6 +10,7 @@
 import io.openbas.database.repository.InjectRepository;
 import io.openbas.rest.exception.ElementNotFoundException;
 import io.openbas.rest.finding.FindingService;
+import io.openbas.rest.inject.form.InjectExecutionAction;
 import io.openbas.rest.inject.form.InjectExecutionInput;
 import io.openbas.rest.inject.form.InjectExpectationUpdateInput;
 import io.openbas.service.InjectExpectationService;
@@ -19,14 +20,14 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
-import java.util.logging.Level;
 import lombok.RequiredArgsConstructor;
-import lombok.extern.java.Log;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.dao.DataIntegrityViolationException;
 import org.springframework.stereotype.Service;
 
 @RequiredArgsConstructor
 @Service
-@Log
+@Slf4j
 public class InjectExecutionService {
 
   private final InjectRepository injectRepository;
@@ -45,6 +46,27 @@ public void handleInjectExecutionCallback(
 
     try {
       inject = loadInjectOrThrow(injectId);
+      // issue/3550: added this condition to ensure we only update statuses if the inject is in a
+      // coherent state.
+      // This prevents issues where the PENDING status took more time to persist than it took for
+      // the agent to send the complete action.
+      // FIXME: At the moment, this whole function is only called by our implant. These implant are
+      // launched with the async value to true, which force the implant to go from EXECUTING to
+      // PENDING, before going to EXECUTED.
+      // So if in the future, this function is called to update a synchronous inject, we will need
+      // to find a way to get the async boolean somehow and add it to this condition.
+      if (input.getAction().equals(InjectExecutionAction.complete)
+          && (inject.getStatus().isEmpty()
+              || !inject.getStatus().get().getName().equals(ExecutionStatus.PENDING))) {
+        // If we receive a status update with a terminal state status, we must first check that the
+        // current status is in the PENDING state
+        log.warn(
+            "Received a complete action for inject {} with status {}, but current status is not PENDING",
+            injectId,
+            inject.getStatus().map(is -> is.getName().toString()).orElse("unknown"));
+        throw new DataIntegrityViolationException(
+            "Cannot complete inject that is not in PENDING state");
+      }
       Agent agent = loadAgentIfPresent(agentId);
 
       Set<OutputParser> outputParsers = structuredOutputUtils.extractOutputParsers(inject);
@@ -181,7 +203,7 @@ private Inject loadInjectOrThrow(String injectId) {
   }
 
   private void handleInjectExecutionError(Inject inject, Exception e) {
-    log.log(Level.SEVERE, e.getMessage());
+    log.error(e.getMessage(), e);
     if (inject != null) {
       inject
           .getStatus()
diff --git a/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectStatusService.java b/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectStatusService.java
--- a/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectStatusService.java
+++ b/openbas-api/src/main/java/io/openbas/rest/inject/service/InjectStatusService.java
@@ -50,7 +50,7 @@ public Inject updateInjectStatus(String injectId, InjectUpdateStatusInput input)
   }
 
   public void addStartImplantExecutionTraceByInject(
-      String injectId, String agentId, String message) {
+      String injectId, String agentId, String message, Instant startTime) {
     InjectStatus injectStatus =
         injectStatusRepository.findByInjectId(injectId).orElseThrow(ElementNotFoundException::new);
     Agent agent = agentRepository.findById(agentId).orElseThrow(ElementNotFoundException::new);
@@ -62,7 +62,7 @@ public void addStartImplantExecutionTraceByInject(
             message,
             ExecutionTraceAction.START,
             agent,
-            null);
+            startTime);
     injectStatus.addTrace(trace);
     injectStatusRepository.save(injectStatus);
   }
@@ -80,13 +80,6 @@ private int getCompleteTrace(Inject inject) {
   public boolean isAllInjectAgentsExecuted(Inject inject) {
     int totalCompleteTrace = getCompleteTrace(inject);
     List<Agent> agents = this.injectService.getAgentsByInject(inject);
-    log.info(
-        "[issue/2797] isAllInjectAgentsExecuted "
-            + inject.getId()
-            + " totalCompleteTrace: "
-            + totalCompleteTrace
-            + " agents.size(); "
-            + agents.size());
     return agents.size() == totalCompleteTrace;
   }
 
@@ -108,11 +101,34 @@ public ExecutionTrace createExecutionTrace(
       InjectExecutionInput input,
       Agent agent,
       ObjectNode structuredOutput) {
+
+    // We start by computing the trace date. It should be qual to the START execution trace +
+    // input.duration.
+    // If the duration is 0 or if there is no START execution trace, we use the current time.
+    Instant traceCreationTime =
+        (injectStatus.getTraces().isEmpty() || input.getDuration() == 0)
+            ? Instant.now()
+            : injectStatus.getTraces().stream()
+                .filter(
+                    t ->
+                        ExecutionTraceAction.START.equals(t.getAction())
+                            && t.getAgent().getId().equals(agent.getId()))
+                .findFirst()
+                .map(startTrace -> startTrace.getTime().plusMillis(input.getDuration()))
+                .orElse(Instant.now());
+
     ExecutionTraceAction executionAction = convertExecutionAction(input.getAction());
     ExecutionTraceStatus traceStatus = ExecutionTraceStatus.valueOf(input.getStatus());
+
     ExecutionTrace base =
         new ExecutionTrace(
-            injectStatus, traceStatus, null, input.getMessage(), executionAction, agent, null);
+            injectStatus,
+            traceStatus,
+            null,
+            input.getMessage(),
+            executionAction,
+            agent,
+            traceCreationTime);
     return ExecutionTrace.from(base, structuredOutput);
   }
 
@@ -143,14 +159,6 @@ public void updateInjectStatus(
       if (executionTrace.getAction().equals(ExecutionTraceAction.COMPLETE)
           && (agent == null || isAllInjectAgentsExecuted(inject))) {
         updateFinalInjectStatus(injectStatus);
-      } else {
-        log.info(
-            "[issue/2797] updateInjectStatuselse injectId: "
-                + inject.getId()
-                + " executionTrace.getAction(): "
-                + executionTrace.getAction()
-                + " agent == null: "
-                + agent);
       }
 
       injectRepository.save(inject);


Test patch:
diff --git a/openbas-api/src/test/java/io/openbas/rest/inject/InjectApiTest.java b/openbas-api/src/test/java/io/openbas/rest/inject/InjectApiTest.java
--- a/openbas-api/src/test/java/io/openbas/rest/inject/InjectApiTest.java
+++ b/openbas-api/src/test/java/io/openbas/rest/inject/InjectApiTest.java
@@ -554,7 +554,7 @@ void getExecutablePayloadInjectWithArguments() throws Exception {
       Inject injectSaved = injectRepository.save(inject);
       doNothing()
           .when(injectStatusService)
-          .addStartImplantExecutionTraceByInject(any(), any(), any());
+          .addStartImplantExecutionTraceByInject(any(), any(), any(), any());
 
       // -- EXECUTE --
       String response =
@@ -643,7 +643,7 @@ void given_targetedAssetArgument_should_replaceByAssetIDs() throws Exception {
       Inject injectSaved = injectRepository.save(inject);
       doNothing()
           .when(injectStatusService)
-          .addStartImplantExecutionTraceByInject(any(), any(), any());
+          .addStartImplantExecutionTraceByInject(any(), any(), any(), any());
 
       // -- EXECUTE --
       String response =
@@ -689,7 +689,7 @@ void getExecutableObfuscatePayloadInject() throws Exception {
       Inject injectSaved = injectRepository.save(inject);
       doNothing()
           .when(injectStatusService)
-          .addStartImplantExecutionTraceByInject(any(), any(), any());
+          .addStartImplantExecutionTraceByInject(any(), any(), any(), any());
 
       // -- EXECUTE --
       String response =