diff --git a/verified-inference/host/api/domain/tee/shutdown_enclave_use_case.py b/verified-inference/host/api/domain/tee/shutdown_enclave_use_case.py index 558a552..ff795ec 100644 --- a/verified-inference/host/api/domain/tee/shutdown_enclave_use_case.py +++ b/verified-inference/host/api/domain/tee/shutdown_enclave_use_case.py @@ -9,6 +9,6 @@ async def execute(enclave_name: str) -> bool: enclave = enclaves[enclave_name] print(f"Requesting shutdown from enclave: {enclave_name}") - shutdown_result = shutdown(enclave.cid) + shutdown_result = await shutdown(enclave.cid) print(f"Enclave shut down: {shutdown_result}") return shutdown_result diff --git a/verified-inference/host/api/repository/nitro_cli_repository.py b/verified-inference/host/api/repository/nitro_cli_repository.py index e3cdb3a..2fb370e 100644 --- a/verified-inference/host/api/repository/nitro_cli_repository.py +++ b/verified-inference/host/api/repository/nitro_cli_repository.py @@ -69,7 +69,7 @@ def _create_run_enclave_command(enclave_name: str, memory_mb: int) -> List[str]: f"{TEE_CPU_COUNT}", "--memory", f"{memory_mb}", - "--debug-mode", + # "--debug-mode", "--eif-path", f"{enclave_name}.eif", ] diff --git a/verified-inference/host/api/repository/vsock_repository.py b/verified-inference/host/api/repository/vsock_repository.py index f949a67..559a4f6 100644 --- a/verified-inference/host/api/repository/vsock_repository.py +++ b/verified-inference/host/api/repository/vsock_repository.py @@ -1,3 +1,4 @@ +import errno import socket from enum import Enum from typing import Optional @@ -22,9 +23,15 @@ async def request_attestation(enclave_cid: int) -> str: async def shutdown(enclave_cid: int) -> bool: - response = _send_request( - enclave_cid, VsockRequestType.SHUTDOWN, receive_response=True - ) + try: + response = _send_request( + enclave_cid, VsockRequestType.SHUTDOWN, receive_response=True + ) + except OSError as e: + if e.errno == errno.ENOTCONN: + # The enclave closed the connection, which is expected + return True + return False if response is None: raise Exception("No response received from shutdown request") return True diff --git a/verified-inference/host/api/service/tee/entities.py b/verified-inference/host/api/service/tee/entities.py index 4fef591..cdf154a 100644 --- a/verified-inference/host/api/service/tee/entities.py +++ b/verified-inference/host/api/service/tee/entities.py @@ -20,7 +20,7 @@ class TeeTerminateRequest(BaseModel): class TeeTerminateResponse(BaseModel): - result: Dict = Field(description="Termination result") + result: bool = Field(description="Termination result") class TeeGetEnclaveResponse(BaseModel):