TeeService.java

/*
 * Copyright 2022-2025 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.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;
import com.iexec.sms.api.TeeSessionGenerationError;
import com.iexec.sms.api.TeeSessionGenerationResponse;
import com.iexec.worker.sms.SmsService;
import com.iexec.worker.sms.TeeSessionGenerationException;
import com.iexec.worker.workflow.WorkflowError;
import lombok.extern.slf4j.Slf4j;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public abstract class TeeService {
    private final SmsService smsService;
    protected final TeeServicesPropertiesService teeServicesPropertiesService;

    private final Map<String, TeeSessionGenerationResponse> teeSessions = new ConcurrentHashMap<>();

    protected TeeService(final SmsService smsService,
                         final TeeServicesPropertiesService teeServicesPropertiesService) {
        this.smsService = smsService;
        this.teeServicesPropertiesService = teeServicesPropertiesService;
    }

    public List<WorkflowError> areTeePrerequisitesMetForTask(final String chainTaskId) {
        if (!isTeeEnabled()) {
            return List.of(new WorkflowError(ReplicateStatusCause.TEE_NOT_SUPPORTED));
        }

        try {
            // Try to load the `SmsClient` relative to the task.
            // If it can't be loaded, then we won't be able to run the task.
            smsService.getSmsClient(chainTaskId);
        } catch (SmsClientCreationException e) {
            log.error("Couldn't get SmsClient [chainTaskId:{}]", chainTaskId, e);
            return List.of(new WorkflowError(ReplicateStatusCause.UNKNOWN_SMS));
        }

        // 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 {
        if (teeSessions.containsKey(workerpoolAuthorization.getChainTaskId())) {
            log.warn("TEE session already exists for task [chainTaskId:{}]", workerpoolAuthorization.getChainTaskId());
            return;
        }
        final TeeSessionGenerationResponse secureSession = smsService.createTeeSession(workerpoolAuthorization);
        if (secureSession == null) {
            throw new TeeSessionGenerationException(TeeSessionGenerationError.UNKNOWN_ISSUE);
        }
        teeSessions.put(workerpoolAuthorization.getChainTaskId(), secureSession);
    }

    public TeeSessionGenerationResponse getTeeSession(final String chainTaskId) {
        return teeSessions.get(chainTaskId);
    }

    public abstract boolean isTeeEnabled();

    /**
     * Start any required service(s) to use TEE with selected technology for given task.
     *
     * @param chainTaskId Task whose service(s) should be started.
     * @return {@literal true} if all services have been correctly started, {@literal false} otherwise.
     */
    public abstract boolean prepareTeeForTask(String chainTaskId);

    public abstract List<String> buildPreComputeDockerEnv(TaskDescription taskDescription);

    public abstract List<String> buildComputeDockerEnv(TaskDescription taskDescription);

    public abstract List<String> buildPostComputeDockerEnv(TaskDescription taskDescription);

    public abstract Collection<String> getAdditionalBindings();

    public abstract List<Device> getDevices();

    // region Purge

    /**
     * Purge description of given task.
     *
     * @param chainTaskId ID of the task to purge.
     * @return {@literal true} if no TEE session remains in cache after purge, {@literal false} otherwise.
     */
    protected boolean purgeTask(final String chainTaskId) {
        teeSessions.remove(chainTaskId);
        return !teeSessions.containsKey(chainTaskId);
    }

    /**
     * Purge all cached task descriptions.
     */
    protected void purgeAllTasksData() {
        teeSessions.clear();
    }

    // endregion
}