feat: check app enclave configuration in preflight checks

src/main/java/com/iexec/worker/compute/pre/PreComputeService.java

@@ -23,7 +23,6 @@
 import com.iexec.commons.containers.DockerRunRequest;
 import com.iexec.commons.containers.DockerRunResponse;
 import com.iexec.commons.poco.task.TaskDescription;
-import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
 import com.iexec.sms.api.config.TeeAppProperties;
 import com.iexec.sms.api.config.TeeServicesProperties;
 import com.iexec.worker.compute.ComputeExitCauseService;
@@ -37,7 +36,6 @@
 import com.iexec.worker.workflow.WorkflowError;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Service;
-import org.springframework.util.unit.DataSize;
 
 import java.time.Duration;
 import java.util.List;
@@ -81,31 +79,6 @@ public PreComputeResponse runTeePreCompute(final TaskDescription taskDescription
         final String chainTaskId = taskDescription.getChainTaskId();
         final PreComputeResponse.PreComputeResponseBuilder preComputeResponseBuilder = PreComputeResponse.builder();
 
-        // verify enclave configuration for compute stage
-        final TeeEnclaveConfiguration enclaveConfig = taskDescription.getAppEnclaveConfiguration();
-        if (enclaveConfig == null) {
-            log.error("No enclave configuration found for task [chainTaskId:{}]", chainTaskId);
-            return preComputeResponseBuilder
-                    .exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION)))
-                    .build();
-        }
-        if (!enclaveConfig.getValidator().isValid()) {
-            log.error("Invalid enclave configuration [chainTaskId:{}, violations:{}]",
-                    chainTaskId, enclaveConfig.getValidator().validate().toString());
-            return preComputeResponseBuilder
-                    .exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_CONFIGURATION)))
-                    .build();
-        }
-        long teeComputeMaxHeapSize = DataSize
-                .ofGigabytes(workerConfigService.getTeeComputeMaxHeapSizeGb())
-                .toBytes();
-        if (enclaveConfig.getHeapSize() > teeComputeMaxHeapSize) {
-            log.error("Enclave configuration should define a proper heap size [chainTaskId:{}, heapSize:{}, maxHeapSize:{}]",
-                    chainTaskId, enclaveConfig.getHeapSize(), teeComputeMaxHeapSize);
-            preComputeResponseBuilder.exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_HEAP_CONFIGURATION)));
-            return preComputeResponseBuilder.build();
-        }
-
         // run TEE pre-compute container if needed
         if (taskDescription.requiresPreCompute()) {
             log.info("Task contains TEE input data [chainTaskId:{}, containsDataset:{}, containsInputFiles:{}, isBulkRequest:{}]",

src/main/java/com/iexec/worker/tee/TeeService.java

@@ -17,6 +17,7 @@
 package com.iexec.worker.tee;
 
 import com.github.dockerjava.api.model.Device;
+import com.iexec.common.replicate.ReplicateStatusCause;
 import com.iexec.commons.poco.chain.WorkerpoolAuthorization;
 import com.iexec.commons.poco.task.TaskDescription;
 import com.iexec.sms.api.SmsClientCreationException;
@@ -32,8 +33,6 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
-import static com.iexec.common.replicate.ReplicateStatusCause.*;
-
 @Slf4j
 public abstract class TeeService {
     private final SmsService smsService;
@@ -49,7 +48,7 @@ protected TeeService(final SmsService smsService,
 
     public List<WorkflowError> areTeePrerequisitesMetForTask(final String chainTaskId) {
         if (!isTeeEnabled()) {
-            return List.of(new WorkflowError(TEE_NOT_SUPPORTED));
+            return List.of(new WorkflowError(ReplicateStatusCause.TEE_NOT_SUPPORTED));
         }
 
         try {
@@ -58,21 +57,12 @@ public List<WorkflowError> areTeePrerequisitesMetForTask(final String chainTaskI
             smsService.getSmsClient(chainTaskId);
         } catch (SmsClientCreationException e) {
             log.error("Couldn't get SmsClient [chainTaskId: {}]", chainTaskId, e);
-            return List.of(new WorkflowError(UNKNOWN_SMS));
-        }
-        try {
-            // Try to load the `TeeServicesProperties` relative to the task.
-            // If it can't be loaded, then we won't be able to run the task.
-            teeServicesPropertiesService.getTeeServicesProperties(chainTaskId);
-        } catch (NullPointerException e) {
-            log.error("TEE enclave configuration is null [chainTaskId: {}]", chainTaskId, e);
-            return List.of(new WorkflowError(PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION));
-        } catch (RuntimeException e) {
-            log.error("Couldn't get TeeServicesProperties [chainTaskId: {}]", chainTaskId, e);
-            return List.of(new WorkflowError(GET_TEE_SERVICES_CONFIGURATION_FAILED));
+            return List.of(new WorkflowError(ReplicateStatusCause.UNKNOWN_SMS));
         }
 
-        return List.of();
+        // Try to load the `TeeServicesProperties` relative to the task.
+        // If it can't be loaded, then we won't be able to run the task.
+        return teeServicesPropertiesService.retrieveTeeServicesProperties(chainTaskId);
     }
 
     public void createTeeSession(final WorkerpoolAuthorization workerpoolAuthorization) throws TeeSessionGenerationException {

src/main/java/com/iexec/worker/tee/TeeServicesPropertiesService.java

@@ -18,21 +18,26 @@
 
 import com.iexec.common.lifecycle.purge.ExpiringTaskMapFactory;
 import com.iexec.common.lifecycle.purge.Purgeable;
+import com.iexec.common.replicate.ReplicateStatusCause;
 import com.iexec.commons.containers.client.DockerClientInstance;
-import com.iexec.commons.poco.chain.IexecHubAbstractService;
 import com.iexec.commons.poco.task.TaskDescription;
 import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
 import com.iexec.commons.poco.tee.TeeFramework;
 import com.iexec.sms.api.SmsClient;
 import com.iexec.sms.api.config.TeeServicesProperties;
+import com.iexec.worker.chain.IexecHubService;
+import com.iexec.worker.config.WorkerConfigurationService;
 import com.iexec.worker.docker.DockerService;
 import com.iexec.worker.sms.SmsService;
+import com.iexec.worker.workflow.WorkflowError;
 import jakarta.annotation.PreDestroy;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Service;
+import org.springframework.util.unit.DataSize;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 
 /**
  * Manages the {@link TeeServicesProperties}, providing an easy way to get properties for a task
@@ -43,66 +48,89 @@
 public class TeeServicesPropertiesService implements Purgeable {
     private final SmsService smsService;
     private final DockerService dockerService;
-    private final IexecHubAbstractService iexecHubService;
+    private final IexecHubService iexecHubService;
+    private final WorkerConfigurationService workerConfigurationService;
 
     private final Map<String, TeeServicesProperties> propertiesForTask = ExpiringTaskMapFactory.getExpiringTaskMap();
 
-    public TeeServicesPropertiesService(SmsService smsService,
-                                        DockerService dockerService,
-                                        IexecHubAbstractService iexecHubService) {
+    public TeeServicesPropertiesService(final SmsService smsService,
+                                        final DockerService dockerService,
+                                        final IexecHubService iexecHubService,
+                                        final WorkerConfigurationService workerConfigurationService) {
         this.smsService = smsService;
         this.dockerService = dockerService;
         this.iexecHubService = iexecHubService;
+        this.workerConfigurationService = workerConfigurationService;
     }
 
     public TeeServicesProperties getTeeServicesProperties(final String chainTaskId) {
-        return propertiesForTask.computeIfAbsent(chainTaskId, this::retrieveTeeServicesProperties);
+        return propertiesForTask.get(chainTaskId);
     }
 
-    <T extends TeeServicesProperties> T retrieveTeeServicesProperties(final String chainTaskId) {
+    public List<WorkflowError> retrieveTeeServicesProperties(final String chainTaskId) {
         final TaskDescription taskDescription = iexecHubService.getTaskDescription(chainTaskId);
 
+        // TODO errors could be renamed for APP enclave checks
+        final TeeEnclaveConfiguration teeEnclaveConfiguration = taskDescription.getAppEnclaveConfiguration();
+        if (teeEnclaveConfiguration == null) {
+            log.error("No enclave configuration found for task [chainTaskId:{}]", chainTaskId);
+            return List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION));
+        }
+        if (!teeEnclaveConfiguration.getValidator().isValid()) {
+            log.error("Invalid enclave configuration [chainTaskId:{}, violations:{}]",
+                    chainTaskId, teeEnclaveConfiguration.getValidator().validate().toString());
+            return List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_CONFIGURATION));
+        }
+        long teeComputeMaxHeapSize = DataSize
+                .ofGigabytes(workerConfigurationService.getTeeComputeMaxHeapSizeGb())
+                .toBytes();
+        if (teeEnclaveConfiguration.getHeapSize() > teeComputeMaxHeapSize) {
+            log.error("Enclave configuration should define a proper heap size [chainTaskId:{}, heapSize:{}, maxHeapSize:{}]",
+                    chainTaskId, teeEnclaveConfiguration.getHeapSize(), teeComputeMaxHeapSize);
+            return List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_HEAP_CONFIGURATION));
+        }
+
         // SMS client should already have been created once before.
         // If it couldn't be created, then the task would have been aborted.
         // So the following won't throw an exception.
         final SmsClient smsClient = smsService.getSmsClient(chainTaskId);
         final TeeFramework teeFramework = taskDescription.getTeeFramework();
         final TeeFramework smsTeeFramework = smsClient.getTeeFramework();
         if (smsTeeFramework != teeFramework) {
-            throw new TeeServicesPropertiesCreationException(
-                    "SMS is configured for another TEE framework" +
-                            " [chainTaskId:" + chainTaskId +
-                            ", requiredFramework:" + teeFramework +
-                            ", actualFramework:" + smsTeeFramework + "]");
+            return List.of(new WorkflowError(ReplicateStatusCause.GET_TEE_SERVICES_CONFIGURATION_FAILED,
+                    String.format("SMS is configured for another TEE framework [chainTaskId:%s, requiredFramework:%s, actualFramework:%s]",
+                            chainTaskId, teeFramework, smsTeeFramework)));
         }
 
-        final TeeEnclaveConfiguration teeEnclaveConfiguration = taskDescription.getAppEnclaveConfiguration();
-        Objects.requireNonNull(teeEnclaveConfiguration, "Missing TEE enclave configuration [chainTaskId:" + chainTaskId + "]");
-
-        final T properties = smsClient.getTeeServicesPropertiesVersion(teeFramework, teeEnclaveConfiguration.getVersion());
-        log.info("Received TEE services properties [properties:{}]", properties);
+        final TeeServicesProperties properties = smsClient.getTeeServicesPropertiesVersion(teeFramework, teeEnclaveConfiguration.getVersion());
         if (properties == null) {
-            throw new TeeServicesPropertiesCreationException(
-                    "Missing TEE services properties [chainTaskId:" + chainTaskId + "]");
+            return List.of(new WorkflowError(ReplicateStatusCause.GET_TEE_SERVICES_CONFIGURATION_FAILED,
+                    String.format("Missing TEE services properties [chainTaskId:%s]", chainTaskId)));
         }
+        log.info("TEE services properties received [chainTaskId:{}]", chainTaskId);
 
         final String preComputeImage = properties.getPreComputeProperties().getImage();
         final String postComputeImage = properties.getPostComputeProperties().getImage();
+        final List<WorkflowError> errors = new ArrayList<>();
 
-        checkImageIsPresentOrDownload(preComputeImage, chainTaskId, "preComputeImage");
-        checkImageIsPresentOrDownload(postComputeImage, chainTaskId, "postComputeImage");
+        errors.addAll(checkImageIsPresentOrDownload(preComputeImage, chainTaskId, "preComputeImage"));
+        errors.addAll(checkImageIsPresentOrDownload(postComputeImage, chainTaskId, "postComputeImage"));
 
-        return properties;
+        if (errors.isEmpty()) {
+            propertiesForTask.put(chainTaskId, properties);
+            log.info("TEE services properties storage in cache [chainTaskId:{}, contains-key:{}]",
+                    chainTaskId, propertiesForTask.containsKey(chainTaskId));
+        }
+        return List.copyOf(errors);
     }
 
-    private void checkImageIsPresentOrDownload(final String image, final String chainTaskId, final String imageType) {
+    private List<WorkflowError> checkImageIsPresentOrDownload(final String image, final String chainTaskId, final String imageType) {
         final DockerClientInstance client = dockerService.getClient(image);
-        if (!client.isImagePresent(image)
-                && !client.pullImage(image)) {
-            throw new TeeServicesPropertiesCreationException(
-                    "Failed to download image " +
-                            "[chainTaskId:" + chainTaskId + ", " + imageType + ":" + image + "]");
+        if (!client.isImagePresent(image) && !client.pullImage(image)) {
+            return List.of(new WorkflowError(ReplicateStatusCause.GET_TEE_SERVICES_CONFIGURATION_FAILED,
+                    String.format("Failed to download image [chainTaskId:%s, %s:%s]", chainTaskId, imageType, image)));
         }
+        return List.of();
     }
 
     /**
@@ -114,8 +142,9 @@ private void checkImageIsPresentOrDownload(final String image, final String chai
      */
     @Override
     public boolean purgeTask(final String chainTaskId) {
-        log.debug("purgeTask [chainTaskId:{}]", chainTaskId);
         propertiesForTask.remove(chainTaskId);
+        log.info("TEE services properties removal from cache [chainTaskId:{}, contains-key:{}]",
+                chainTaskId, propertiesForTask.containsKey(chainTaskId));
         return !propertiesForTask.containsKey(chainTaskId);
     }
 

src/test/java/com/iexec/worker/compute/pre/PreComputeServiceTests.java

@@ -23,7 +23,6 @@
 import com.iexec.commons.containers.client.DockerClientInstance;
 import com.iexec.commons.poco.chain.DealParams;
 import com.iexec.commons.poco.task.TaskDescription;
-import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
 import com.iexec.commons.poco.tee.TeeFramework;
 import com.iexec.commons.poco.utils.BytesUtils;
 import com.iexec.sms.api.config.TeeAppProperties;
@@ -47,14 +46,13 @@
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
-import org.springframework.util.unit.DataSize;
 
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Stream;
 
-import static com.iexec.common.replicate.ReplicateStatusCause.*;
+import static com.iexec.common.replicate.ReplicateStatusCause.PRE_COMPUTE_FAILED_UNKNOWN_ISSUE;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
@@ -73,12 +71,7 @@ class PreComputeServiceTests {
             .datasetAddress("datasetAddress")
             .datasetUri(datasetUri)
             .datasetChecksum("datasetChecksum")
-            .teeFramework(TeeFramework.SCONE)
-            .appEnclaveConfiguration(TeeEnclaveConfiguration.builder()
-                    .fingerprint("01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b")
-                    .heapSize(1024)
-                    .entrypoint("python /app/app.py")
-                    .build());
+            .teeFramework(TeeFramework.SCONE);
     private final TeeAppProperties preComputeProperties = TeeAppProperties.builder()
             .image(PRE_COMPUTE_IMAGE)
             .entrypoint(PRE_COMPUTE_ENTRYPOINT)
@@ -122,7 +115,6 @@ void prepareMockWhenPreComputeShouldRunForTask(final TaskDescription taskDescrip
     }
 
     void prepareMocksForPreCompute(final TaskDescription taskDescription, DockerRunResponse dockerRunResponse) {
-        when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
         when(dockerService.getClient()).thenReturn(dockerClientInstanceMock);
         when(teeServicesManager.getTeeService(any())).thenReturn(teeMockedService);
         when(teeServicesPropertiesService.getTeeServicesProperties(chainTaskId)).thenReturn(properties);
@@ -216,46 +208,9 @@ void shouldRunTeePreComputeAndPrepareInputDataWhenBulkProcessingRequested() {
         verifyDockerRun();
     }
 
-    @Test
-    void shouldFailToRunTeePreComputeSinceMissingEnclaveConfiguration() {
-        final TaskDescription taskDescription = taskDescriptionBuilder.appEnclaveConfiguration(null).build();
-
-        final PreComputeResponse response = preComputeService.runTeePreCompute(taskDescription);
-        assertThat(response.isSuccessful()).isFalse();
-        assertThat(response.getExitCauses())
-                .containsExactly(new WorkflowError(PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION));
-    }
-
-    @Test
-    void shouldFailToRunTeePreComputeSinceInvalidEnclaveConfiguration() {
-        final TeeEnclaveConfiguration enclaveConfig = TeeEnclaveConfiguration.builder().build();
-        final TaskDescription taskDescription = taskDescriptionBuilder.appEnclaveConfiguration(enclaveConfig).build();
-        assertThat(enclaveConfig.getValidator().isValid()).isFalse();
-
-        final PreComputeResponse response = preComputeService.runTeePreCompute(taskDescription);
-        assertThat(response.isSuccessful()).isFalse();
-        assertThat(response.getExitCauses())
-                .containsExactly(new WorkflowError(PRE_COMPUTE_INVALID_ENCLAVE_CONFIGURATION));
-    }
-
-    @Test
-    void shouldFailToRunTeePreComputeSinceTooHighComputeHeapSize() {
-        when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
-        final TeeEnclaveConfiguration enclaveConfiguration = TeeEnclaveConfiguration.builder()
-                .fingerprint("01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b")
-                .heapSize(DataSize.ofGigabytes(8).toBytes() + 1)
-                .entrypoint("python /app/app.py")
-                .build();
-        final TaskDescription taskDescription = taskDescriptionBuilder.appEnclaveConfiguration(enclaveConfiguration).build();
-
-        assertThat(preComputeService.runTeePreCompute(taskDescription).isSuccessful())
-                .isFalse();
-    }
-
     @Test
     void shouldNotRunTeePreComputeSinceDockerImageNotFoundLocally() {
         final TaskDescription taskDescription = taskDescriptionBuilder.build();
-        when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
         when(dockerService.getClient()).thenReturn(dockerClientInstanceMock);
         when(teeServicesPropertiesService.getTeeServicesProperties(chainTaskId)).thenReturn(properties);
         when(properties.getPreComputeProperties()).thenReturn(preComputeProperties);
@@ -326,8 +281,6 @@ void shouldNotRunPreComputeWhenNotRequired() {
                 .dealParams(DealParams.builder().build())
                 .build();
 
-        when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
-
         assertThat(taskDescription.containsDataset()).isFalse();
         assertThat(taskDescription.containsInputFiles()).isFalse();
         assertThat(taskDescription.isBulkRequest()).isFalse();