Skip to content

Commit 7c11687

Browse files
jbern0rdjbern0rd
authored
feat: check app enclave configuration in preflight checks (#667)
1 parent 981f186 commit 7c11687

File tree

6 files changed

+161
-214
lines changed

6 files changed

+161
-214
lines changed

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

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import com.iexec.commons.containers.DockerRunRequest;
2424
import com.iexec.commons.containers.DockerRunResponse;
2525
import com.iexec.commons.poco.task.TaskDescription;
26-
import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
2726
import com.iexec.sms.api.config.TeeAppProperties;
2827
import com.iexec.sms.api.config.TeeServicesProperties;
2928
import com.iexec.worker.compute.ComputeExitCauseService;
@@ -37,7 +36,6 @@
3736
import com.iexec.worker.workflow.WorkflowError;
3837
import lombok.extern.slf4j.Slf4j;
3938
import org.springframework.stereotype.Service;
40-
import org.springframework.util.unit.DataSize;
4139

4240
import java.time.Duration;
4341
import java.util.List;
@@ -81,31 +79,6 @@ public PreComputeResponse runTeePreCompute(final TaskDescription taskDescription
8179
final String chainTaskId = taskDescription.getChainTaskId();
8280
final PreComputeResponse.PreComputeResponseBuilder preComputeResponseBuilder = PreComputeResponse.builder();
8381

84-
// verify enclave configuration for compute stage
85-
final TeeEnclaveConfiguration enclaveConfig = taskDescription.getAppEnclaveConfiguration();
86-
if (enclaveConfig == null) {
87-
log.error("No enclave configuration found for task [chainTaskId:{}]", chainTaskId);
88-
return preComputeResponseBuilder
89-
.exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION)))
90-
.build();
91-
}
92-
if (!enclaveConfig.getValidator().isValid()) {
93-
log.error("Invalid enclave configuration [chainTaskId:{}, violations:{}]",
94-
chainTaskId, enclaveConfig.getValidator().validate().toString());
95-
return preComputeResponseBuilder
96-
.exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_CONFIGURATION)))
97-
.build();
98-
}
99-
long teeComputeMaxHeapSize = DataSize
100-
.ofGigabytes(workerConfigService.getTeeComputeMaxHeapSizeGb())
101-
.toBytes();
102-
if (enclaveConfig.getHeapSize() > teeComputeMaxHeapSize) {
103-
log.error("Enclave configuration should define a proper heap size [chainTaskId:{}, heapSize:{}, maxHeapSize:{}]",
104-
chainTaskId, enclaveConfig.getHeapSize(), teeComputeMaxHeapSize);
105-
preComputeResponseBuilder.exitCauses(List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_HEAP_CONFIGURATION)));
106-
return preComputeResponseBuilder.build();
107-
}
108-
10982
// run TEE pre-compute container if needed
11083
if (taskDescription.requiresPreCompute()) {
11184
log.info("Task contains TEE input data [chainTaskId:{}, containsDataset:{}, containsInputFiles:{}, isBulkRequest:{}]",

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

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.iexec.worker.tee;
1818

1919
import com.github.dockerjava.api.model.Device;
20+
import com.iexec.common.replicate.ReplicateStatusCause;
2021
import com.iexec.commons.poco.chain.WorkerpoolAuthorization;
2122
import com.iexec.commons.poco.task.TaskDescription;
2223
import com.iexec.sms.api.SmsClientCreationException;
@@ -32,8 +33,6 @@
3233
import java.util.Map;
3334
import java.util.concurrent.ConcurrentHashMap;
3435

35-
import static com.iexec.common.replicate.ReplicateStatusCause.*;
36-
3736
@Slf4j
3837
public abstract class TeeService {
3938
private final SmsService smsService;
@@ -49,7 +48,7 @@ protected TeeService(final SmsService smsService,
4948

5049
public List<WorkflowError> areTeePrerequisitesMetForTask(final String chainTaskId) {
5150
if (!isTeeEnabled()) {
52-
return List.of(new WorkflowError(TEE_NOT_SUPPORTED));
51+
return List.of(new WorkflowError(ReplicateStatusCause.TEE_NOT_SUPPORTED));
5352
}
5453

5554
try {
@@ -58,21 +57,12 @@ public List<WorkflowError> areTeePrerequisitesMetForTask(final String chainTaskI
5857
smsService.getSmsClient(chainTaskId);
5958
} catch (SmsClientCreationException e) {
6059
log.error("Couldn't get SmsClient [chainTaskId: {}]", chainTaskId, e);
61-
return List.of(new WorkflowError(UNKNOWN_SMS));
62-
}
63-
try {
64-
// Try to load the `TeeServicesProperties` relative to the task.
65-
// If it can't be loaded, then we won't be able to run the task.
66-
teeServicesPropertiesService.getTeeServicesProperties(chainTaskId);
67-
} catch (NullPointerException e) {
68-
log.error("TEE enclave configuration is null [chainTaskId: {}]", chainTaskId, e);
69-
return List.of(new WorkflowError(PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION));
70-
} catch (RuntimeException e) {
71-
log.error("Couldn't get TeeServicesProperties [chainTaskId: {}]", chainTaskId, e);
72-
return List.of(new WorkflowError(GET_TEE_SERVICES_CONFIGURATION_FAILED));
60+
return List.of(new WorkflowError(ReplicateStatusCause.UNKNOWN_SMS));
7361
}
7462

75-
return List.of();
63+
// Try to load the `TeeServicesProperties` relative to the task.
64+
// If it can't be loaded, then we won't be able to run the task.
65+
return teeServicesPropertiesService.retrieveTeeServicesProperties(chainTaskId);
7666
}
7767

7868
public void createTeeSession(final WorkerpoolAuthorization workerpoolAuthorization) throws TeeSessionGenerationException {

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

Lines changed: 59 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,26 @@
1818

1919
import com.iexec.common.lifecycle.purge.ExpiringTaskMapFactory;
2020
import com.iexec.common.lifecycle.purge.Purgeable;
21+
import com.iexec.common.replicate.ReplicateStatusCause;
2122
import com.iexec.commons.containers.client.DockerClientInstance;
22-
import com.iexec.commons.poco.chain.IexecHubAbstractService;
2323
import com.iexec.commons.poco.task.TaskDescription;
2424
import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
2525
import com.iexec.commons.poco.tee.TeeFramework;
2626
import com.iexec.sms.api.SmsClient;
2727
import com.iexec.sms.api.config.TeeServicesProperties;
28+
import com.iexec.worker.chain.IexecHubService;
29+
import com.iexec.worker.config.WorkerConfigurationService;
2830
import com.iexec.worker.docker.DockerService;
2931
import com.iexec.worker.sms.SmsService;
32+
import com.iexec.worker.workflow.WorkflowError;
3033
import jakarta.annotation.PreDestroy;
3134
import lombok.extern.slf4j.Slf4j;
3235
import org.springframework.stereotype.Service;
36+
import org.springframework.util.unit.DataSize;
3337

38+
import java.util.ArrayList;
39+
import java.util.List;
3440
import java.util.Map;
35-
import java.util.Objects;
3641

3742
/**
3843
* Manages the {@link TeeServicesProperties}, providing an easy way to get properties for a task
@@ -43,66 +48,89 @@
4348
public class TeeServicesPropertiesService implements Purgeable {
4449
private final SmsService smsService;
4550
private final DockerService dockerService;
46-
private final IexecHubAbstractService iexecHubService;
51+
private final IexecHubService iexecHubService;
52+
private final WorkerConfigurationService workerConfigurationService;
4753

4854
private final Map<String, TeeServicesProperties> propertiesForTask = ExpiringTaskMapFactory.getExpiringTaskMap();
4955

50-
public TeeServicesPropertiesService(SmsService smsService,
51-
DockerService dockerService,
52-
IexecHubAbstractService iexecHubService) {
56+
public TeeServicesPropertiesService(final SmsService smsService,
57+
final DockerService dockerService,
58+
final IexecHubService iexecHubService,
59+
final WorkerConfigurationService workerConfigurationService) {
5360
this.smsService = smsService;
5461
this.dockerService = dockerService;
5562
this.iexecHubService = iexecHubService;
63+
this.workerConfigurationService = workerConfigurationService;
5664
}
5765

5866
public TeeServicesProperties getTeeServicesProperties(final String chainTaskId) {
59-
return propertiesForTask.computeIfAbsent(chainTaskId, this::retrieveTeeServicesProperties);
67+
return propertiesForTask.get(chainTaskId);
6068
}
6169

62-
<T extends TeeServicesProperties> T retrieveTeeServicesProperties(final String chainTaskId) {
70+
public List<WorkflowError> retrieveTeeServicesProperties(final String chainTaskId) {
6371
final TaskDescription taskDescription = iexecHubService.getTaskDescription(chainTaskId);
6472

73+
// TODO errors could be renamed for APP enclave checks
74+
final TeeEnclaveConfiguration teeEnclaveConfiguration = taskDescription.getAppEnclaveConfiguration();
75+
if (teeEnclaveConfiguration == null) {
76+
log.error("No enclave configuration found for task [chainTaskId:{}]", chainTaskId);
77+
return List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION));
78+
}
79+
if (!teeEnclaveConfiguration.getValidator().isValid()) {
80+
log.error("Invalid enclave configuration [chainTaskId:{}, violations:{}]",
81+
chainTaskId, teeEnclaveConfiguration.getValidator().validate().toString());
82+
return List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_CONFIGURATION));
83+
}
84+
long teeComputeMaxHeapSize = DataSize
85+
.ofGigabytes(workerConfigurationService.getTeeComputeMaxHeapSizeGb())
86+
.toBytes();
87+
if (teeEnclaveConfiguration.getHeapSize() > teeComputeMaxHeapSize) {
88+
log.error("Enclave configuration should define a proper heap size [chainTaskId:{}, heapSize:{}, maxHeapSize:{}]",
89+
chainTaskId, teeEnclaveConfiguration.getHeapSize(), teeComputeMaxHeapSize);
90+
return List.of(new WorkflowError(ReplicateStatusCause.PRE_COMPUTE_INVALID_ENCLAVE_HEAP_CONFIGURATION));
91+
}
92+
6593
// SMS client should already have been created once before.
6694
// If it couldn't be created, then the task would have been aborted.
6795
// So the following won't throw an exception.
6896
final SmsClient smsClient = smsService.getSmsClient(chainTaskId);
6997
final TeeFramework teeFramework = taskDescription.getTeeFramework();
7098
final TeeFramework smsTeeFramework = smsClient.getTeeFramework();
7199
if (smsTeeFramework != teeFramework) {
72-
throw new TeeServicesPropertiesCreationException(
73-
"SMS is configured for another TEE framework" +
74-
" [chainTaskId:" + chainTaskId +
75-
", requiredFramework:" + teeFramework +
76-
", actualFramework:" + smsTeeFramework + "]");
100+
return List.of(new WorkflowError(ReplicateStatusCause.GET_TEE_SERVICES_CONFIGURATION_FAILED,
101+
String.format("SMS is configured for another TEE framework [chainTaskId:%s, requiredFramework:%s, actualFramework:%s]",
102+
chainTaskId, teeFramework, smsTeeFramework)));
77103
}
78104

79-
final TeeEnclaveConfiguration teeEnclaveConfiguration = taskDescription.getAppEnclaveConfiguration();
80-
Objects.requireNonNull(teeEnclaveConfiguration, "Missing TEE enclave configuration [chainTaskId:" + chainTaskId + "]");
81-
82-
final T properties = smsClient.getTeeServicesPropertiesVersion(teeFramework, teeEnclaveConfiguration.getVersion());
83-
log.info("Received TEE services properties [properties:{}]", properties);
105+
final TeeServicesProperties properties = smsClient.getTeeServicesPropertiesVersion(teeFramework, teeEnclaveConfiguration.getVersion());
84106
if (properties == null) {
85-
throw new TeeServicesPropertiesCreationException(
86-
"Missing TEE services properties [chainTaskId:" + chainTaskId + "]");
107+
return List.of(new WorkflowError(ReplicateStatusCause.GET_TEE_SERVICES_CONFIGURATION_FAILED,
108+
String.format("Missing TEE services properties [chainTaskId:%s]", chainTaskId)));
87109
}
110+
log.info("TEE services properties received [chainTaskId:{}]", chainTaskId);
88111

89112
final String preComputeImage = properties.getPreComputeProperties().getImage();
90113
final String postComputeImage = properties.getPostComputeProperties().getImage();
114+
final List<WorkflowError> errors = new ArrayList<>();
91115

92-
checkImageIsPresentOrDownload(preComputeImage, chainTaskId, "preComputeImage");
93-
checkImageIsPresentOrDownload(postComputeImage, chainTaskId, "postComputeImage");
116+
errors.addAll(checkImageIsPresentOrDownload(preComputeImage, chainTaskId, "preComputeImage"));
117+
errors.addAll(checkImageIsPresentOrDownload(postComputeImage, chainTaskId, "postComputeImage"));
94118

95-
return properties;
119+
if (errors.isEmpty()) {
120+
propertiesForTask.put(chainTaskId, properties);
121+
log.info("TEE services properties storage in cache [chainTaskId:{}, contains-key:{}]",
122+
chainTaskId, propertiesForTask.containsKey(chainTaskId));
123+
}
124+
return List.copyOf(errors);
96125
}
97126

98-
private void checkImageIsPresentOrDownload(final String image, final String chainTaskId, final String imageType) {
127+
private List<WorkflowError> checkImageIsPresentOrDownload(final String image, final String chainTaskId, final String imageType) {
99128
final DockerClientInstance client = dockerService.getClient(image);
100-
if (!client.isImagePresent(image)
101-
&& !client.pullImage(image)) {
102-
throw new TeeServicesPropertiesCreationException(
103-
"Failed to download image " +
104-
"[chainTaskId:" + chainTaskId + ", " + imageType + ":" + image + "]");
129+
if (!client.isImagePresent(image) && !client.pullImage(image)) {
130+
return List.of(new WorkflowError(ReplicateStatusCause.GET_TEE_SERVICES_CONFIGURATION_FAILED,
131+
String.format("Failed to download image [chainTaskId:%s, %s:%s]", chainTaskId, imageType, image)));
105132
}
133+
return List.of();
106134
}
107135

108136
/**
@@ -114,8 +142,9 @@ private void checkImageIsPresentOrDownload(final String image, final String chai
114142
*/
115143
@Override
116144
public boolean purgeTask(final String chainTaskId) {
117-
log.debug("purgeTask [chainTaskId:{}]", chainTaskId);
118145
propertiesForTask.remove(chainTaskId);
146+
log.info("TEE services properties removal from cache [chainTaskId:{}, contains-key:{}]",
147+
chainTaskId, propertiesForTask.containsKey(chainTaskId));
119148
return !propertiesForTask.containsKey(chainTaskId);
120149
}
121150

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

Lines changed: 2 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import com.iexec.commons.containers.client.DockerClientInstance;
2424
import com.iexec.commons.poco.chain.DealParams;
2525
import com.iexec.commons.poco.task.TaskDescription;
26-
import com.iexec.commons.poco.tee.TeeEnclaveConfiguration;
2726
import com.iexec.commons.poco.tee.TeeFramework;
2827
import com.iexec.commons.poco.utils.BytesUtils;
2928
import com.iexec.sms.api.config.TeeAppProperties;
@@ -47,14 +46,13 @@
4746
import org.mockito.InjectMocks;
4847
import org.mockito.Mock;
4948
import org.mockito.junit.jupiter.MockitoExtension;
50-
import org.springframework.util.unit.DataSize;
5149

5250
import java.time.Duration;
5351
import java.util.List;
5452
import java.util.Map;
5553
import java.util.stream.Stream;
5654

57-
import static com.iexec.common.replicate.ReplicateStatusCause.*;
55+
import static com.iexec.common.replicate.ReplicateStatusCause.PRE_COMPUTE_FAILED_UNKNOWN_ISSUE;
5856
import static org.assertj.core.api.Assertions.assertThat;
5957
import static org.mockito.ArgumentMatchers.any;
6058
import static org.mockito.Mockito.*;
@@ -73,12 +71,7 @@ class PreComputeServiceTests {
7371
.datasetAddress("datasetAddress")
7472
.datasetUri(datasetUri)
7573
.datasetChecksum("datasetChecksum")
76-
.teeFramework(TeeFramework.SCONE)
77-
.appEnclaveConfiguration(TeeEnclaveConfiguration.builder()
78-
.fingerprint("01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b")
79-
.heapSize(1024)
80-
.entrypoint("python /app/app.py")
81-
.build());
74+
.teeFramework(TeeFramework.SCONE);
8275
private final TeeAppProperties preComputeProperties = TeeAppProperties.builder()
8376
.image(PRE_COMPUTE_IMAGE)
8477
.entrypoint(PRE_COMPUTE_ENTRYPOINT)
@@ -122,7 +115,6 @@ void prepareMockWhenPreComputeShouldRunForTask(final TaskDescription taskDescrip
122115
}
123116

124117
void prepareMocksForPreCompute(final TaskDescription taskDescription, DockerRunResponse dockerRunResponse) {
125-
when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
126118
when(dockerService.getClient()).thenReturn(dockerClientInstanceMock);
127119
when(teeServicesManager.getTeeService(any())).thenReturn(teeMockedService);
128120
when(teeServicesPropertiesService.getTeeServicesProperties(chainTaskId)).thenReturn(properties);
@@ -216,46 +208,9 @@ void shouldRunTeePreComputeAndPrepareInputDataWhenBulkProcessingRequested() {
216208
verifyDockerRun();
217209
}
218210

219-
@Test
220-
void shouldFailToRunTeePreComputeSinceMissingEnclaveConfiguration() {
221-
final TaskDescription taskDescription = taskDescriptionBuilder.appEnclaveConfiguration(null).build();
222-
223-
final PreComputeResponse response = preComputeService.runTeePreCompute(taskDescription);
224-
assertThat(response.isSuccessful()).isFalse();
225-
assertThat(response.getExitCauses())
226-
.containsExactly(new WorkflowError(PRE_COMPUTE_MISSING_ENCLAVE_CONFIGURATION));
227-
}
228-
229-
@Test
230-
void shouldFailToRunTeePreComputeSinceInvalidEnclaveConfiguration() {
231-
final TeeEnclaveConfiguration enclaveConfig = TeeEnclaveConfiguration.builder().build();
232-
final TaskDescription taskDescription = taskDescriptionBuilder.appEnclaveConfiguration(enclaveConfig).build();
233-
assertThat(enclaveConfig.getValidator().isValid()).isFalse();
234-
235-
final PreComputeResponse response = preComputeService.runTeePreCompute(taskDescription);
236-
assertThat(response.isSuccessful()).isFalse();
237-
assertThat(response.getExitCauses())
238-
.containsExactly(new WorkflowError(PRE_COMPUTE_INVALID_ENCLAVE_CONFIGURATION));
239-
}
240-
241-
@Test
242-
void shouldFailToRunTeePreComputeSinceTooHighComputeHeapSize() {
243-
when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
244-
final TeeEnclaveConfiguration enclaveConfiguration = TeeEnclaveConfiguration.builder()
245-
.fingerprint("01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b")
246-
.heapSize(DataSize.ofGigabytes(8).toBytes() + 1)
247-
.entrypoint("python /app/app.py")
248-
.build();
249-
final TaskDescription taskDescription = taskDescriptionBuilder.appEnclaveConfiguration(enclaveConfiguration).build();
250-
251-
assertThat(preComputeService.runTeePreCompute(taskDescription).isSuccessful())
252-
.isFalse();
253-
}
254-
255211
@Test
256212
void shouldNotRunTeePreComputeSinceDockerImageNotFoundLocally() {
257213
final TaskDescription taskDescription = taskDescriptionBuilder.build();
258-
when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
259214
when(dockerService.getClient()).thenReturn(dockerClientInstanceMock);
260215
when(teeServicesPropertiesService.getTeeServicesProperties(chainTaskId)).thenReturn(properties);
261216
when(properties.getPreComputeProperties()).thenReturn(preComputeProperties);
@@ -326,8 +281,6 @@ void shouldNotRunPreComputeWhenNotRequired() {
326281
.dealParams(DealParams.builder().build())
327282
.build();
328283

329-
when(workerConfigService.getTeeComputeMaxHeapSizeGb()).thenReturn(8);
330-
331284
assertThat(taskDescription.containsDataset()).isFalse();
332285
assertThat(taskDescription.containsInputFiles()).isFalse();
333286
assertThat(taskDescription.isBulkRequest()).isFalse();

0 commit comments

Comments
 (0)