Rustで実装されたAWS Nitro Enclaves CLIの実装を読む

この記事は [Rust Advent Calendar 2020](https://qiita.com/advent-calendar/2020/rust) 4日目の記事です。

当記事では、先日発表されたAWS Nitro Enclavesの[CLI実装](https://github.com/aws/aws-nitro-enclaves-cli)を見ていきます。CLIをはじめ、[ACMのPKCS#11プロバイダ](https://github.com/aws/aws-nitro-enclaves-acm)やAttestationやPCR、Enclave内での乱数取得のためにデバイスとやりとりするためのライブラリ（[NSM lib](https://github.com/aws/aws-nitro-enclaves-nsm-api)）など多くのライブラリがRustで実装されています。
そこで、ここでは実際のNitro Enclavesの使い方よりも操作するCLI内部でNitro Enclavesを起動するためにどのような処理が行われているかに着目していきます。

### AWS Nitro Enclavesとは
高度にセキュアな情報を扱うことを想定し、AWS Nitro Systemにより提供される隔離保護された実行環境です。

それぞれのEnclaveには、EC2のメモリやCPUリソースを割り当てることができ、独立したカーネルOS上で動作します。Enclaveは、外部のネットワークやストレージに接続することはできず、admin IAM権限でもユーザーアクセスはできません。また、全てのデータは親EC2とのvsockのみでやりとり可能となります。vsockについて、[VSOCK: VM ↔host socket with minimal configuration - DevConf.CZ 2020](https://youtu.be/R5DQWdPUQSY)の動画が参考になります。
<div align="center">
<img width="400" alt="image" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/135415/476e1d7c-378b-e9a2-79bc-adb63cce0767.png">
</dev>
（引用：[AWS Nitro Enclaves – Isolated EC2 Environments to Process Confidential Data](https://aws.amazon.com/jp/blogs/aws/aws-nitro-enclaves-isolated-ec2-environments-to-process-confidential-data/)）

つまり、Nitro Enclavesを利用するには、Nitro Systemをベースとしているインスタンスでのみ利用することができます。

### CLIでrun-enclave

enclaveを起動するために必要なEIF(Enclave Image Format)ファイルがdocker imageから[Nitro CLI](https://github.com/aws/aws-nitro-enclaves-cli)を利用してビルドすることができます。EIFをビルドしたら、以下のようなコマンドでenclaveを起動することができ、起動したenclaveの情報が返ってきます。

```
	$ ./nitro-cli run-enclave --cpu-count 4 --memory 2048 --eif-path command_executer.eif
	Start allocating memory...
	Running on instance CPUs [1, 5, 2, 6]
	Started enclave with enclave-cid: 16, memory: 2048 MiB, cpu-ids: [1, 5, 2, 6]
	Sending image to cid: 16 port: 7000
	{
	  "EnclaveID": "i-08aa8a2f7bff2ff99_enc103923520469154997",
	  "EnclaveCID": 16,
	  "NumberOfCPUs": 4,
	  "CPUIDs": [
	    1,
	    5,
	    2,
	    6
	  ],
	  "MemoryMiB": 2048
	}
```

この`run-enclave`コマンドの内部では、以下のような処理が行われています。

1. 親プロセス：CLIコマンドでコミュニケーションするためのEnclave用デーモンの生成
2. デーモン：コマンドを受け取るためのevent loopを回す
3. 親プロセス：socketに`run-enclave`コマンドを送信
4. デーモン：nitroデバイスファイルへVM生成コマンド送信
5. デーモン：メモリの割り当てとEIFをメモリに書き込み
6. デーモン：vCPUの割り当て
7. デーモン：nitroデバイスファイルを介しenclaveをスタート
8. 親プロセス：enclaveからのレスポンスを処理し表示

これらのステップ全ての実装を追っていくと量が多いので、特に参考になる実装、本質的な実装の部分を抽出して見ていきます。以下のコードのCLIコマンドをうけて、デーモンを生成する`enclave_proc_spawn`とコマンドをenclaveに送って初期化する`enclave_proc_command_send_single`の2つに分けて紹介します。

```rust
("run-enclave", Some(args)) => {
            ・・・
            let mut comm = enclave_proc_spawn(&logger)
                .map_err(|err| {
                    err.add_subaction("Failed to spawn enclave process".to_string())
                        .set_action(RUN_ENCLAVE_STR.to_string())
                })
                .ok_or_exit_with_errno(None);

            enclave_proc_command_send_single(
                EnclaveProcessCommandType::Run,
                Some(&run_args),
                &mut comm,
            )
            .map_err(|e| {
                e.add_subaction("Failed to send single command".to_string())
                    .set_action(RUN_ENCLAVE_STR.to_string())
            })
            .ok_or_exit_with_errno(None);
            ・・・
        }
```

### コマンド実行からデーモンの生成まで

#### enclave_proc_spawn

まず、`run-enclave`コマンドが実行されると、`enclave_proc_spawn`が呼ばれます。この関数では、enclaveとソケット通信するための準備、enclave用デーモン生成のためのフォーク、そしてその子プロセスで`enclave_process_run`関数が実行されています。

#### 

```rust
/// Spawn an enclave process and wait until it has detached and has
/// taken ownership of its communication socket.
pub fn enclave_proc_spawn(logger: &EnclaveProcLogWriter) -> NitroCliResult<UnixStream> {
    let (cli_socket, enclave_proc_socket) = UnixStream::pair().map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Could not create a socket pair: {:?}", e),
            NitroCliErrorEnum::SocketPairCreationFailure
        )
    })?;

    ・・・

    // Spawn an intermediate child process. This will fork again in order to
    // create the detached enclave process.
    let fork_status = fork();

    if let Ok(ForkResult::Child) = fork_status {
        // This is our intermediate child process.
        enclave_process_run(enclave_proc_socket, logger);
    } else {
        ・・・
    }

    ・・・
}
```

#### enclave_process_run

`enclave_process_run`を実行し、`create_enclave_process`でenclaveプロセス（デーモン）の生成処理を行い、そのデーモンでCLIなどからのコマンドを待ち受けるためのイベントループを`process_event_loop`を回しています。そして、この`process_event_loop`に先ほど生成した親インスタンスとのやりとりのためにソケットを渡しています。（ここでは、`process_event_loop`内の処理は本質ではないので省略します。）

```rust
/// Launch the enclave process.
///
/// * `comm_fd` - A descriptor used for initial communication with the parent Nitro CLI instance.
/// * `logger` - The current log writer, whose ID gets updated when an enclave is launched.
pub fn enclave_process_run(comm_stream: UnixStream, logger: &EnclaveProcLogWriter) {
    create_enclave_process(logger)
        .map_err(|e| e.set_action("Run Enclave".to_string()))
        .ok_or_exit_with_errno(None);
    let res = process_event_loop(comm_stream, logger);

    ・・・
}
```

#### create_enclave_process：Parentet instance側でデーモンの生成

続いて、実際にデーモンを生成する処理を見ていきます。ここではデーモンを生成するために大きく4つの処理が行われています。

1. SIGHUPのシグナルを無視するようハンドラ設定
2. 現在のプロセスをデーモン化
3. デタッチしたプロセスが確実にオーファンになるまで待つ
4. シグナルハンドラを元に戻す

以降、ここで作成したデーモンがCLIからのコマンドを受け取ってEnclaveとソケット通信をすることになります。

```rust
/// Create the enclave process.
fn create_enclave_process(logger: &EnclaveProcLogWriter) -> NitroCliResult<()> {
    // To get a detached process, we first:
    // (1) Temporarily ignore specific signals (SIGHUP).
    // (2) Daemonize the current process.
    // (3) Wait until the detached process is orphaned.
    // (4) Restore signal handlers.
    let signal_handler = SignalHandler::new(&[SIGHUP])
        .mask_all()
        .map_err(|e| e.add_subaction("Failed to mask signals".to_string()))?;
    let ppid = getpid();

    // Daemonize the current process. The working directory remains
    // unchanged and the standard descriptors are routed to '/dev/null'.
    daemon(true, false).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to daemonize enclave process: {:?}", e),
            NitroCliErrorEnum::DaemonizeProcessFailure
        )
    })?;

    // This is our detached process.
    logger
        .update_logger_id(format!("enc-xxxxxxx:{}", std::process::id()).as_str())
        .map_err(|e| e.add_subaction("Failed to update logger id".to_string()))?;
    info!("Enclave process PID: {}", process::id());

    // We must wait until we're 100% orphaned. That is, our parent must
    // no longer be the pre-fork process.
    while getppid() == ppid {
        thread::sleep(std::time::Duration::from_millis(10));
    }

    // Restore signal handlers.
    signal_handler
        .unmask_all()
        .map_err(|e| e.add_subaction("Failed to restore signal handlers".to_string()))?;

    Ok(())
}
```

### Enclaveの初期化まで

ここまでで、デーモンが立ち上がりましたので`run_enclave`コマンドを実際にenclaveへ送り、初期化する`enclave_proc_command_send_single`を見ていきます。`enclave_proc_command_send_single`内部では、親インスタンスがソケットにコマンドを書き込む処理をしています。それをデーモンがイベントループで受け取り、`run_enclaves`を呼び出します。

#### run_enclaves

ここでは、コマンドラインの引数として渡した、EIFのファイルパスやCPU数、メモリ数などを`EnclaveManager`に渡して、`run_enclave`を実行してenclaveを起動させています。

```rust
/// Launch an enclave with the specified arguments and provide the launch status through the given connection.
pub fn run_enclaves(
    args: &RunEnclavesArgs,
    connection: Option<&Connection>,
) -> NitroCliResult<EnclaveManager> {
    debug!("run_enclaves");

    let eif_file = File::open(&args.eif_path).map_err(|e| {
        new_nitro_cli_failure!(
            &format!("Failed to open the EIF file: {:?}", e),
            NitroCliErrorEnum::FileOperationFailure
        )
        .add_info(vec![&args.eif_path, "Open"])
    })?;

    let cpu_ids = CpuInfo::new()
        .map_err(|e| e.add_subaction("Failed to construct CPU information".to_string()))?
        .get_cpu_config(args)
        .map_err(|e| e.add_subaction("Failed to get CPU configuration".to_string()))?;
    let mut enclave_manager = EnclaveManager::new(
        args.enclave_cid,
        args.memory_mib,
        cpu_ids,
        eif_file,
        args.debug_mode.unwrap_or(false),
    )
    .map_err(|e| {
        e.add_subaction("Failed to construct EnclaveManager with given arguments".to_string())
    })?;
    enclave_manager
        .run_enclave(connection)
        .map_err(|e| e.add_subaction("Failed to run enclave".to_string()))?;
    enclave_manager
        .update_state(EnclaveState::Running)
        .map_err(|e| e.add_subaction("Failed to update enclave state".to_string()))?;

    Ok(enclave_manager)
}
```

`EnclaveManager`は以下のように定義されていて、実体は`EnclacveHandle`となっているので、その処理を見ていきます。

```rust
/// The structure which manages an enclave in a thread-safe manner.
#[derive(Clone, Default)]
pub struct EnclaveManager {
    /// The full ID of the managed enclave.
    pub enclave_id: String,
    /// A thread-safe handle to the enclave's resources.
    enclave_handle: Arc<Mutex<EnclaveHandle>>,
}
```

#### EnclaveHandle::new

メモリサイズやファイルディスクリプタのバリデーションは省略して、ハンドラ生成の本質的な部分はNitro Enclavesデバイスドライバのioctlを呼び出している部分です。コマンドにはNitro EnclavesのVMを作成する`NE_CREATE_VM`を渡しています。

```rust
/// Create a new enclave handle instance.
    fn new(
        enclave_cid: Option<u64>,
        memory_mib: u64,
        cpu_config: EnclaveCpuConfig,
        eif_file: File,
        debug_mode: bool,
    ) -> NitroCliResult<Self> {
        ・・・

        // Open the device file.
        let dev_file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(NE_DEV_FILEPATH)
            .map_err(|e| {
                ・・・
            })?;

        let mut slot_uid: u64 = 0;
        let enc_fd = EnclaveHandle::do_ioctl(dev_file.as_raw_fd(), NE_CREATE_VM, &mut slot_uid)
            .map_err(|e| e.add_subaction("Create VM ioctl failed".to_string()))?;
        ・・・
    }
```

ちなみに、Nitro Enclacvesのデバイスファイルパスとioctlのコマンドは以下のように定義されています。

```rust
/// Path corresponding to the Nitro Enclaves device file.
const NE_DEV_FILEPATH: &str = "/dev/nitro_enclaves";

/// IOCTL code for `NE_CREATE_VM`.
pub const NE_CREATE_VM: u64 = nix::request_code_read!(NE_MAGIC, 0x20, size_of::<u64>()) as _;

/// IOCTL code for `NE_ADD_VCPU`.
pub const NE_ADD_VCPU: u64 = nix::request_code_readwrite!(NE_MAGIC, 0x21, size_of::<u32>()) as _;

/// IOCTL code for `NE_GET_IMAGE_LOAD_INFO`.
pub const NE_GET_IMAGE_LOAD_INFO: u64 =
    nix::request_code_readwrite!(NE_MAGIC, 0x22, size_of::<ImageLoadInfo>()) as _;

/// IOCTL code for `NE_SET_USER_MEMORY_REGION`.
pub const NE_SET_USER_MEMORY_REGION: u64 =
    nix::request_code_write!(NE_MAGIC, 0x23, size_of::<MemoryRegion>()) as _;

/// IOCTL code for `NE_START_ENCLAVE`.
pub const NE_START_ENCLAVE: u64 =
    nix::request_code_readwrite!(NE_MAGIC, 0x24, size_of::<EnclaveStartInfo>()) as _;
```

### Enclave環境の初期化と実行

続いて、生成したenclaveでの初期化処理をしていきます。ここでもさまざまな処理を行っていますが、重要な部分はメモリの割り当て（EIFイメージの展開）とvCPUの割り当てとenclave起動のためのコマンドをioctlでNitro Systemに送っている部分です。

```rust
/// Initialize the enclave environment and start the enclave.
    fn create_enclave(&mut self, connection: Option<&Connection>) -> NitroCliResult<String> {
        self.init_memory(connection)
            .map_err(|e| e.add_subaction("Memory initialization issue".to_string()))?;
        self.init_cpus()
            .map_err(|e| e.add_subaction("vCPUs initialization issue".to_string()))?;

        ・・・

        let enclave_start = self
            .start(connection)
            .map_err(|e| e.add_subaction("Enclave start issue".to_string()))?;
        
        ・・・
    }
```

メモリの割り当てとeif_fileをローディングさせる実装は以下のようになっています。まず、EIFイメージを展開するenclaveでのメモリのオフセットを取得するために`NE_GET_IMAGE_LOAD_INFO`コマンドを渡してioctlを呼びます。続いて、`run-enclave`CLIコマンドの`--memory`オプションで指定したメモリサイズからenclaveメモリ領域（サイズと仮想アドレス）を生成し、先ほど取得したオフセットからEIFイメージをメモリに展開します。
その上でioctlを介して、`NE_SET_USER_MEMORY_REGION`コマンドとパラメタでメモリ領域を送り、そのメモリ領域の所有権をNitro Enclaveに切り替えます。

```rust
/// Allocate memory and provide it to the enclave.
    fn init_memory(&mut self, connection: Option<&Connection>) -> NitroCliResult<()> {
        // Allocate the memory regions needed by the enclave.
        safe_conn_eprintln(connection, "Start allocating memory...")?;

        let requested_mem_mib = self.resource_allocator.requested_mem >> 20;
        let regions = self
            .resource_allocator
            .allocate()
            .map_err(|e| e.add_subaction("Failed to allocate enclave memory".to_string()))?;

        self.allocated_memory_mib = regions.iter().fold(0, |mut acc, val| {
            acc += val.mem_size;
            acc
        }) >> 20;

        if self.allocated_memory_mib < requested_mem_mib {
            return Err(new_nitro_cli_failure!(
                &format!(
                    "Failed to allocate sufficient memory (requested {} MB, but got {} MB)",
                    requested_mem_mib, self.allocated_memory_mib
                ),
                NitroCliErrorEnum::InsufficientMemoryAvailable
            )
            .add_info(vec!["memory", &requested_mem_mib.to_string()]));
        }

        let eif_file = self.eif_file.as_mut().ok_or_else(|| {
            new_nitro_cli_failure!(
                "Failed to get mutable reference to EIF file",
                NitroCliErrorEnum::FileOperationFailure
            )
        })?;

        let mut image_load_info = ImageLoadInfo {
            flags: NE_EIF_IMAGE,
            memory_offset: 0,
        };
        EnclaveHandle::do_ioctl(self.enc_fd, NE_GET_IMAGE_LOAD_INFO, &mut image_load_info)
            .map_err(|e| e.add_subaction("Get image load info ioctl failed".to_string()))?;

        debug!("Memory load information: {:?}", image_load_info);
        write_eif_to_regions(eif_file, regions, image_load_info.memory_offset as usize)
            .map_err(|e| e.add_subaction("Write EIF to enclave memory regions".to_string()))?;

        // Provide the regions to the driver for ownership change.
        for region in regions {
            let mut user_mem_region: UserMemoryRegion = region.into();
            EnclaveHandle::do_ioctl(self.enc_fd, NE_SET_USER_MEMORY_REGION, &mut user_mem_region)
                .map_err(|e| e.add_subaction("Set user memory region ioctl failed".to_string()))?;
        }

        info!("Finished initializing memory.");

        Ok(())
    }
```

vCPUの割り当てに関しては、`run-enclave`コマンドの`--cpu-count`オプションで渡した回数だけ、enclaveにvCPUの割り当てるコマンド`NE_ADD_VCPU`を送ります。

```rust
    /// Provide CPUs from the parent instance to the enclave.
    fn init_cpus(&mut self) -> NitroCliResult<()> {
        let cpu_config = self.cpu_config.clone();

        match cpu_config {
            EnclaveCpuConfig::List(cpu_ids) => {
                for cpu_id in cpu_ids {
                    self.init_single_cpu(cpu_id.clone()).map_err(|e| {
                        e.add_subaction(format!("Failed to add CPU with ID {}", cpu_id))
                    })?;
                }
            }
            EnclaveCpuConfig::Count(cpu_count) => {
                for _ in 0..cpu_count {
                    self.init_single_cpu(0)?;
                }
            }
        }

        Ok(())
    }

    /// Initialize a single vCPU from a given ID.
    fn init_single_cpu(&mut self, mut cpu_id: u32) -> NitroCliResult<()> {
        EnclaveHandle::do_ioctl(self.enc_fd, NE_ADD_VCPU, &mut cpu_id)
            .map_err(|e| e.add_subaction("Add vCPU ioctl failed".to_string()))?;

        self.cpu_ids.push(cpu_id);
        debug!("Added CPU with ID {}.", cpu_id);

        Ok(())
    }
```

そして、メモリとvCPUの割り当てやEIFイメージの展開をした後にenclaveをスタートさせるコマンド`NE_START_ENCLAVE`を送ることで、Nitro Enclaveが起動します。

```rust
    /// Start an enclave after providing it with its necessary resources.
    fn start(&mut self, connection: Option<&Connection>) -> NitroCliResult<EnclaveStartInfo> {
        let mut start = EnclaveStartInfo::new(&self);

        EnclaveHandle::do_ioctl(self.enc_fd, NE_START_ENCLAVE, &mut start)
            .map_err(|e| e.add_subaction("Start enclave ioctl failed".to_string()))?;

        ・・・
    }
```


## まとめ

aws-nitro-enclaves-cliの`run-enclave`コマンドによりあらかじめビルドしておいたEIFを用いて、デバイスファイルを介してNitro Enclavesを起動させるまでの流れを追いました。また、今回生成したNitro Enclaveと親インスタンスがvsock経由で通信するサンプル実装は[Nitro Enclaves Samples](https://github.com/aws/aws-nitro-enclaves-samples)が参考になります。他にも、attestationによりNitro enclaveがブートプロセスから改ざんされていないか、イメージやenclaveのアプリケーションが改ざんされていないかといったことをPCRと呼ばれるハッシュ値や署名から検証することができます。（[NSM lib](https://github.com/aws/aws-nitro-enclaves-nsm-api)でRustのインターフェースが実装されています）

##　参考資料

* [Nitro Enclavesの実装](https://github.com/aws?q=nitro&type=&language=)
* [Deep Dive Into AWS Nitro Enclaves - AWS Online Tech Talks](https://www.youtube.com/watch?v=K5PRNHaEdOw)
* [AWS re:Invent 2019: Powering next-gen Amazon EC2: Deep dive into the Nitro system (CMP303-R2) ](https://www.youtube.com/watch?v=rUY-00yFlE4)