When undefined, set final deadline after retention-duration for up to batch-size TEE challenges during cleanup (#281)

Author: jbern0rd

CHANGELOG.md

@@ -11,6 +11,7 @@ All notable changes to this project will be documented in this file.
 - Configure the SMS at startup to generate Scone sessions in Hardware or MAA mode. (#275)
 - Add configurable cron job to delete expired tasks TEE challenges and Ethereum credentials. (#278)
 - Use new `FileHashUtils` API. (#280)
+- When undefined, set final deadline after `retention-duration` for up to `batch-size` TEE challenges during cleanup. (#281)
 
 ### Quality
 
@@ -24,7 +25,7 @@ All notable changes to this project will be documented in this file.
 
 - Upgrade to `eclipse-temurin:11.0.24_8-jre-focal`. (#270)
 - Upgrade to Gradle 8.10.2. (#271)
-- Upgrade to `iexec-commons-poco` 4.1.0-NEXT-SNAPSHOT. (#272)
+- Upgrade to H2 database 2.2.224. (#281)
 
 ## [[8.6.0]](https://github.com/iExecBlockchainComputing/iexec-sms/releases/tag/v8.6.0) 2024-06-18
 

README.md

@@ -66,6 +66,8 @@ To support:
 | `IEXEC_SECRET_PROVISIONER_ENCLAVE_HOSTNAME` | Secret provisioner server host for retrieving secrets from attested enclaves. Typically used by workers to execute TEE tasks. | Positive integer | `localhost` | `localhost` |
 | `IEXEC_SECRET_PROVISIONER_ENCLAVE_PORT`|  Secret provisioner server port for retrieving secrets from attested enclaves. | Positive integer | `18765` | `4433` |
 | `IEXEC_TEE_CHALLENGE_CLEANUP_CRON` | Cron expression to configure TEE challenges cleanup policy. | String | `@hourly` | `@hourly` |
+| `IEXEC_TEE_CHALLENGE_CLEANUP_MAX_BATCH_SIZE` | Max number of TEE challenges whose missing deadline could be set at a given time. | Integer | `500` | `500` |
+| `IEXEC_TEE_CHALLENGE_CLEANUP_RETENTION_DURATION` | Retention duration when setting missing final deadline. | Duration | `P5D` | `P5D` |
 | `IEXEC_TEE_WORKER_PRE_COMPUTE_IMAGE` | TEE enabled OCI image name for worker pre-compute stage of TEE tasks. | String | | |
 | `IEXEC_TEE_WORKER_PRE_COMPUTE_FINGERPRINT` | Fingerprint (aka mrenclave) of the TEE enabled worker pre-compute image. | String | | |
 | `IEXEC_TEE_WORKER_PRE_COMPUTE_HEAP_SIZE_GB` | Required heap size for a worker pre-compute enclave (in Giga Bytes). | Positive integer | `3` | `3` |

build.gradle

@@ -57,7 +57,7 @@ dependencies {
     implementation 'org.springframework.retry:spring-retry'
     // H2
     implementation 'org.springframework.boot:spring-boot-starter-data-jpa'
-    implementation 'com.h2database:h2:2.2.222'
+    implementation 'com.h2database:h2:2.2.224'
 
     // Spring Doc
     implementation 'org.springdoc:springdoc-openapi-ui:1.7.0'

src/main/java/com/iexec/sms/tee/challenge/TeeChallengeRepository.java

@@ -27,4 +27,6 @@ public interface TeeChallengeRepository extends JpaRepository<TeeChallenge, Stri
 
     @Transactional
     void deleteByFinalDeadlineBefore(Instant now);
+
+    int countByFinalDeadlineIsNull();
 }

src/main/java/com/iexec/sms/tee/challenge/TeeChallengeService.java

@@ -19,7 +19,9 @@
 import com.iexec.sms.blockchain.IexecHubService;
 import com.iexec.sms.encryption.EncryptionService;
 import com.iexec.sms.secret.MeasuredSecretService;
+import com.iexec.sms.tee.config.TeeChallengeCleanupConfiguration;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.scheduling.annotation.Scheduled;
 import org.springframework.stereotype.Service;
 
@@ -30,22 +32,28 @@
 @Service
 public class TeeChallengeService {
 
+    private final JdbcTemplate jdbcTemplate;
     private final TeeChallengeRepository teeChallengeRepository;
     private final EncryptionService encryptionService;
     private final IexecHubService iexecHubService;
     private final MeasuredSecretService teeChallengesMeasuredSecretService;
     private final MeasuredSecretService ethereumCredentialsMeasuredSecretService;
+    private final TeeChallengeCleanupConfiguration teeChallengeCleanupConfiguration;
 
-    public TeeChallengeService(final TeeChallengeRepository teeChallengeRepository,
+    public TeeChallengeService(final JdbcTemplate jdbcTemplate,
+                               final TeeChallengeRepository teeChallengeRepository,
                                final EncryptionService encryptionService,
                                final IexecHubService iexecHubService,
                                final MeasuredSecretService teeChallengeMeasuredSecretService,
-                               final MeasuredSecretService ethereumCredentialsMeasuredSecretService) {
+                               final MeasuredSecretService ethereumCredentialsMeasuredSecretService,
+                               final TeeChallengeCleanupConfiguration teeChallengeCleanupConfiguration) {
+        this.jdbcTemplate = jdbcTemplate;
         this.teeChallengeRepository = teeChallengeRepository;
         this.encryptionService = encryptionService;
         this.iexecHubService = iexecHubService;
         this.teeChallengesMeasuredSecretService = teeChallengeMeasuredSecretService;
         this.ethereumCredentialsMeasuredSecretService = ethereumCredentialsMeasuredSecretService;
+        this.teeChallengeCleanupConfiguration = teeChallengeCleanupConfiguration;
     }
 
     public Optional<TeeChallenge> getOrCreate(String taskId, boolean shouldDecryptKeys) {
@@ -81,17 +89,17 @@ public Optional<TeeChallenge> getOrCreate(String taskId, boolean shouldDecryptKe
     }
 
     public void encryptChallengeKeys(TeeChallenge teeChallenge) {
-        EthereumCredentials credentials = teeChallenge.getCredentials();
+        final EthereumCredentials credentials = teeChallenge.getCredentials();
         if (!credentials.isEncrypted()) {
-            String encPrivateKey = encryptionService.encrypt(credentials.getPrivateKey());
+            final String encPrivateKey = encryptionService.encrypt(credentials.getPrivateKey());
             credentials.setEncryptedPrivateKey(encPrivateKey);
         }
     }
 
     public void decryptChallengeKeys(TeeChallenge teeChallenge) {
-        EthereumCredentials credentials = teeChallenge.getCredentials();
+        final EthereumCredentials credentials = teeChallenge.getCredentials();
         if (credentials.isEncrypted()) {
-            String privateKey = encryptionService.decrypt(credentials.getPrivateKey());
+            final String privateKey = encryptionService.decrypt(credentials.getPrivateKey());
             credentials.setPlainTextPrivateKey(privateKey);
         }
     }
@@ -101,10 +109,28 @@ public void decryptChallengeKeys(TeeChallenge teeChallenge) {
      * <p>
      * The interval between two consecutive executions is based on the {@code @Scheduled} annotation
      * and its {@code cron} attribute.
+     * <p>
+     * Look up for entries without deadlines in the database.
+     * If entries are found, set deadline to {@code retentionDuration} in the future
+     * for up to {@code batchSize} entries.
+     *
+     * @see <a href="https://protocol.docs.iex.ec/key-concepts/pay-per-task-model>categories in protocol documentation</a>
      */
     @Scheduled(cron = "${tee.challenge.cleanup.cron}")
     void cleanExpiredTasksTeeChallenges() {
-        log.debug("cleanExpiredTasksTeeChallenges");
+        final long start = System.currentTimeMillis();
         teeChallengeRepository.deleteByFinalDeadlineBefore(Instant.now());
+        final int remaining = teeChallengeRepository.countByFinalDeadlineIsNull();
+        log.info("cleanExpiredTasksTeeChallenges [duration:{}ms, remaining:{}]",
+                System.currentTimeMillis() - start, remaining);
+        if (remaining == 0) {
+            return;
+        }
+        final int updated = jdbcTemplate.update(
+                "UPDATE \"tee_challenge\" SET \"final_deadline\" = ? WHERE \"final_deadline\" IS NULL FETCH FIRST ? ROWS ONLY",
+                Instant.now().plus(teeChallengeCleanupConfiguration.getMissingDeadlineRetentionDuration()),
+                teeChallengeCleanupConfiguration.getMissingDeadlineMaxBatchSize());
+        log.info("cleanExpiredTasksTeeChallenges [duration:{}ms, updated:{}]",
+                System.currentTimeMillis() - start, updated);
     }
 }

src/main/java/com/iexec/sms/tee/config/TeeChallengeCleanupConfiguration.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2024 IEXEC BLOCKCHAIN TECH
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.iexec.sms.tee.config;
+
+import lombok.Value;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.context.properties.ConstructorBinding;
+
+import java.time.Duration;
+
+@Value
+@ConstructorBinding
+@ConfigurationProperties(prefix = "tee.challenge.cleanup")
+public class TeeChallengeCleanupConfiguration {
+    String cron;
+    int missingDeadlineMaxBatchSize;
+    Duration missingDeadlineRetentionDuration;
+}

src/main/resources/application.yml

@@ -51,3 +51,5 @@ springdoc:
 
 tee:
   challenge.cleanup.cron: ${IEXEC_TEE_CHALLENGE_CLEANUP_CRON:@hourly}
+  challenge.cleanup.missing-deadline-max-batch-size: ${IEXEC_TEE_CHALLENGE_CLEANUP_MAX_BATCH_SIZE:500}
+  challenge.cleanup.missing-deadline-retention-duration: ${IEXEC_TEE_CHALLENGE_CLEANUP_RETENTION_DURATION:P5D}

src/test/java/com/iexec/sms/config/OpenApiConfigTests.java

@@ -30,7 +30,7 @@
 
 @ExtendWith(SpringExtension.class)
 @Import(ProjectInfoAutoConfiguration.class)
-class OpenApiConfigTest {
+class OpenApiConfigTests {
 
     @Autowired
     private BuildProperties buildProperties;

src/test/java/com/iexec/sms/tee/challenge/TeeChallengeServiceTests.java

@@ -20,15 +20,19 @@
 import com.iexec.sms.blockchain.IexecHubService;
 import com.iexec.sms.encryption.EncryptionService;
 import com.iexec.sms.secret.MeasuredSecretService;
+import com.iexec.sms.tee.config.TeeChallengeCleanupConfiguration;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
+import org.springframework.jdbc.core.JdbcTemplate;
 
+import javax.persistence.EntityManager;
 import java.security.GeneralSecurityException;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Optional;
 
@@ -46,6 +50,10 @@ class TeeChallengeServiceTests {
 
     private final Instant finalDeadline = Instant.now().minusMillis(1000);
 
+    @Autowired
+    private EntityManager entityManager;
+    @Autowired
+    private JdbcTemplate jdbcTemplate;
     @Autowired
     private EthereumCredentialsRepository ethereumCredentialsRepository;
     @Autowired
@@ -68,12 +76,16 @@ class TeeChallengeServiceTests {
     @BeforeEach
     void beforeEach() {
         teeChallengeRepository.deleteAll();
+        final TeeChallengeCleanupConfiguration cleanupConfiguration = new TeeChallengeCleanupConfiguration(
+                "@hourly", 1, Duration.ofMinutes(1));
         teeChallengeService = new TeeChallengeService(
+                jdbcTemplate,
                 teeChallengeRepository,
                 encryptionService,
                 iexecHubService,
                 teeChallengeMeasuredSecretService,
-                ethereumCredentialsMeasuredSecretService
+                ethereumCredentialsMeasuredSecretService,
+                cleanupConfiguration
         );
     }
 
@@ -178,7 +190,7 @@ void shouldDecryptChallengeKeys() {
     // region cleanExpiredChallenge
     @Test
     void shouldPurgeExpiredChallenge() throws GeneralSecurityException {
-        TeeChallenge encryptedTeeChallengeStub = getEncryptedTeeChallengeStub();
+        final TeeChallenge encryptedTeeChallengeStub = getEncryptedTeeChallengeStub();
         teeChallengeRepository.save(encryptedTeeChallengeStub);
 
         teeChallengeService.cleanExpiredTasksTeeChallenges();
@@ -188,15 +200,26 @@ void shouldPurgeExpiredChallenge() throws GeneralSecurityException {
     }
 
     @Test
-    void shouldNotPurgeExpiredChallenge() throws GeneralSecurityException {
-        TeeChallenge teeChallenge = new TeeChallenge(TASK_ID, null);
+    void shouldSetChallengeFinalDeadlineWhenUnset() throws GeneralSecurityException {
+        final TeeChallenge teeChallenge = new TeeChallenge(TASK_ID, null);
         teeChallenge.getCredentials().setEncryptedPrivateKey(ENC_PRIVATE);
-        teeChallengeRepository.save(teeChallenge);
+        final TeeChallenge savedChallenge = teeChallengeRepository.save(teeChallenge);
+
+        assertThat(teeChallengeRepository.countByFinalDeadlineIsNull()).isOne();
 
         teeChallengeService.cleanExpiredTasksTeeChallenges();
+        // refresh entity to update cache with new database state
+        entityManager.refresh(savedChallenge);
 
         assertThat(teeChallengeRepository.count()).isOne();
         assertThat(ethereumCredentialsRepository.count()).isOne();
+
+        assertThat(teeChallengeRepository.countByFinalDeadlineIsNull()).isZero();
+        final TeeChallenge currentChallenge = teeChallengeRepository.findByTaskId(TASK_ID).orElseThrow();
+        assertThat(currentChallenge).isNotNull();
+        assertThat(currentChallenge.getFinalDeadline())
+                .isNotNull()
+                .isAfter(Instant.now());
     }
     // endregion
 }