feat: add the common tee enclave utilities, and vlc in tee

And tee_vlc module verifiable logic clock is an implementation of Chronos's TEE backend.

Cargo.toml

@@ -11,9 +11,12 @@ members = [
     "crates/cops",
     "crates/vrf",
     "crates/crypto",
+    "crates/enclaves",
+    "crates/types",
     "demos/test_conflict",
-    "demos/coll-tx",
-    "demos/vlc-dag",
+    "demos/coll_tx",
+    "demos/vlc_dag",
+    "demos/tee_vlc",
 ]
 
 [profile.release]

crates/README.md

@@ -20,6 +20,11 @@ The crates folder of Chronos includes core functional code crates and utility li
 - The data store maintains a set of key-value pairs. 
 - It provides causal consistency to clients.
 
+## [enclaves](./enclaves/)
+
+- This module provides some common utilities of TEE (Trusted Execution Environment) Enclaves. 
+- For examples: AWS nitro enclave, Mircosoft Azure, Intel SGX, etc.
+
 ## [crypto](./crypto/)
 
 - Some common crypto utilities, signatures, verify, and hash functions for elliptic curve.

crates/enclaves/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "enclaves"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+nitro-enclaves = ["aws-nitro-enclaves-nsm-api", "aws-nitro-enclaves-attestation"]
+
+[dependencies]
+bincode = "1.3.3"
+tracing = "0.1.40"
+tracing-subscriber = "0.3.18"
+anyhow = { version = "1.0.79", features = ["backtrace"] }
+nix = { version = "0.28.0", features = ["socket", "sched", "resource"] }
+tokio = { version = "1.35.1", features = ["net", "time", "sync", "rt", "signal", "macros", "rt-multi-thread", "fs", "process", "io-util"] }
+aws-nitro-enclaves-nsm-api = { version = "0.4.0", optional = true }
+aws-nitro-enclaves-attestation = { git = "https://github.com/neatsys/aws-nitro-enclaves-attestation", version = "0.1.0", optional = true }

crates/enclaves/src/lib.rs

@@ -0,0 +1 @@
+pub mod nitro_secure;

crates/enclaves/src/nitro_secure.rs

@@ -0,0 +1,155 @@
+use std::sync::Arc;
+use std::{future::Future, pin::Pin};
+use tokio::sync::mpsc::UnboundedSender;
+use tracing::warn;
+
+/// HandleCallbackFn is running handler behind in vsock.
+/// params: input_buf, nsm (nitro secure module), pcrs, write_sender(reply sender)
+pub type HandleFn = Arc<
+    dyn Fn(
+            Vec<u8>,
+            Arc<NitroSecureModule>,
+            [Vec<u8>; 3],
+            UnboundedSender<Vec<u8>>,
+        ) -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send>>
+        + Send
+        + Sync,
+>;
+
+#[derive(Debug)]
+pub struct NitroSecureModule(pub i32);
+
+#[cfg(feature = "nitro-enclaves")]
+impl NitroSecureModule {
+    fn new() -> anyhow::Result<Self> {
+        let fd = aws_nitro_enclaves_nsm_api::driver::nsm_init();
+        anyhow::ensure!(fd >= 0);
+        Ok(Self(fd))
+    }
+
+    pub fn process_attestation(&self, user_data: Vec<u8>) -> anyhow::Result<Vec<u8>> {
+        use aws_nitro_enclaves_nsm_api::api::Request::Attestation;
+        // some silly code to avoid explicitly mention `serde_bytes::ByteBuf`
+        let mut request = Attestation {
+            user_data: Some(Default::default()),
+            nonce: None,
+            public_key: None,
+        };
+        let Attestation {
+            user_data: Some(buf),
+            ..
+        } = &mut request
+        else {
+            unreachable!()
+        };
+        buf.extend(user_data);
+        match aws_nitro_enclaves_nsm_api::driver::nsm_process_request(self.0, request) {
+            aws_nitro_enclaves_nsm_api::api::Response::Attestation { document } => Ok(document),
+            aws_nitro_enclaves_nsm_api::api::Response::Error(err) => anyhow::bail!("{err:?}"),
+            _ => anyhow::bail!("unimplemented"),
+        }
+    }
+
+    fn describe_pcr(&self, index: u16) -> anyhow::Result<Vec<u8>> {
+        use aws_nitro_enclaves_nsm_api::api::Request::DescribePCR;
+        match aws_nitro_enclaves_nsm_api::driver::nsm_process_request(self.0, DescribePCR { index })
+        {
+            aws_nitro_enclaves_nsm_api::api::Response::DescribePCR { lock: _, data } => Ok(data),
+            aws_nitro_enclaves_nsm_api::api::Response::Error(err) => anyhow::bail!("{err:?}"),
+            _ => anyhow::bail!("unimplemented"),
+        }
+    }
+
+    pub async fn run(port: u32, handler: HandleFn) -> anyhow::Result<()> {
+        use std::os::fd::AsRawFd;
+
+        use nix::sys::socket::{
+            bind, listen, socket, AddressFamily, Backlog, SockFlag, SockType, VsockAddr,
+        };
+        use tokio::{
+            io::{AsyncReadExt as _, AsyncWriteExt as _},
+            sync::mpsc::unbounded_channel,
+        };
+
+        let nsm = std::sync::Arc::new(Self::new()?);
+        let pcrs = [
+            nsm.describe_pcr(0)?,
+            nsm.describe_pcr(1)?,
+            nsm.describe_pcr(2)?,
+        ];
+
+        let socket_fd = socket(
+            AddressFamily::Vsock,
+            SockType::Stream,
+            SockFlag::empty(),
+            None,
+        )?;
+        bind(socket_fd.as_raw_fd(), &VsockAddr::new(0xFFFFFFFF, port))?;
+        // theoretically this is the earliest point to entering Tokio world, but i don't want to go
+        // unsafe with `FromRawFd`, and Tokio don't have a `From<OwnedFd>` yet
+        listen(&socket_fd, Backlog::new(64)?)?;
+        let socket = std::os::unix::net::UnixListener::from(socket_fd);
+        socket.set_nonblocking(true)?;
+        let socket = tokio::net::UnixListener::from_std(socket)?;
+
+        loop {
+            let (stream, _) = socket.accept().await?;
+            let (mut read_half, mut write_half) = stream.into_split();
+            let (write_sender, mut write_receiver) = unbounded_channel::<Vec<_>>();
+
+            let mut write_session = tokio::spawn(async move {
+                while let Some(buf) = write_receiver.recv().await {
+                    write_half.write_u64_le(buf.len() as _).await?;
+                    write_half.write_all(&buf).await?;
+                }
+                anyhow::Ok(())
+            });
+            let nsm = nsm.clone();
+            let pcrs = pcrs.clone();
+            let handler = handler.clone(); 
+            let mut read_session = tokio::spawn(async move {
+                loop {
+                    let task = async {
+                        let len = read_half.read_u64_le().await?;
+                        let mut buf = vec![0; len as _];
+                        read_half.read_exact(&mut buf).await?;
+                        anyhow::Ok(buf)
+                    };
+                    let buf = match task.await {
+                        Ok(buf) => buf,
+                        Err(err) => {
+                            warn!("{err}");
+                            return anyhow::Ok(());
+                        }
+                    };
+                    let nsm_clone = nsm.clone();
+                    let pcrs_clone = pcrs.clone();
+                    let write_sender = write_sender.clone();
+                    let handler = handler.clone();
+                    tokio::spawn(async move {
+                        if let Err(err) = handler(buf, nsm_clone, pcrs_clone, write_sender).await {
+                            eprintln!("Error: {:?}", err);
+                        }
+                    });
+                }
+            });
+            loop {
+                let result = tokio::select! {
+                    result = &mut read_session, if !read_session.is_finished() => result,
+                    result = &mut write_session, if !write_session.is_finished() => result,
+                    else => break,
+                };
+                if let Err(err) = result.map_err(Into::into).and_then(std::convert::identity) {
+                    warn!("{err}")
+                }
+            }
+        }
+    }
+}
+
+#[cfg(feature = "nitro-enclaves")]
+impl Drop for NitroSecureModule {
+    fn drop(&mut self) {
+        aws_nitro_enclaves_nsm_api::driver::nsm_exit(self.0)
+    }
+}

crates/types/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "types"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+derive_more = "0.99.17"
+derive-where = "1.2.7"
+serde_json = "1.0.114"
+serde = { version = "1.0.195", features = ["derive"] }

crates/types/src/configuration.rs

@@ -0,0 +1,9 @@
+
+/// Networks consts
+/// Suggested buffer size
+pub const DEFAULT_MAX_DATAGRAM_SIZE: usize = 65507;
+
+/// Round number of a block.
+pub type Round = u64;
+
+pub const GENESIS_ROUND: Round = 0;

crates/types/src/lib.rs

@@ -0,0 +1,2 @@
+pub mod raw_wrapper;
+pub mod configuration;

crates/types/src/raw_wrapper.rs

@@ -0,0 +1,32 @@
+use std::{fmt::Debug, hash::Hash};
+use serde::{Deserialize, Serialize};
+
+/// Payload is the vec<u8> but derive many traits, like hash, debug, clone, etc.
+#[derive(
+    Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default, derive_more::Deref, Serialize, Deserialize,
+)]
+pub struct Payload(pub Vec<u8>);
+
+impl Debug for Payload {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if let Ok(s) = std::str::from_utf8(&self.0) {
+            write!(f, "Payload(\"{s}\")")
+        } else {
+            write!(
+                f,
+                "Payload({}{})",
+                self.0
+                    .iter()
+                    .map(|b| format!("{b:02x}"))
+                    .take(32)
+                    .collect::<Vec<_>>()
+                    .concat(),
+                if self.0.len() > 32 {
+                    format!(".. <len {}>", self.0.len())
+                } else {
+                    String::new()
+                }
+            )
+        }
+    }
+}

crates/vlc/Cargo.toml

@@ -6,4 +6,20 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+sha2 = "0.10.8"
+sha3 = "0.10.1"
+rand = "0.8.5"
+rand_distr = "0.4.3"
+bincode = "1.3.3"
+tracing = "0.1.40"
+futures = "0.3.30"
+num_cpus = "1.13.1"
+derive_more = "0.99.17"
+derive-where = "1.2.7"
 serde = { version = "1", features = ["derive"] }
+anyhow = { version = "1.0.79", features = ["backtrace"] }
+tracing-subscriber = "0.3.18"
+secp256k1 = { version = "0.29.0", features = ["rand-std", "serde", "recovery"] }
+tokio = { version = "1.35.1", features = ["net", "time", "sync", "rt", "signal", "macros", "rt-multi-thread", "fs", "process", "io-util"] }
+tokio-util = "0.7.10"
+crypto ={ path = "../crypto", version = "0.1.0"}

crates/vlc/src/lib.rs

@@ -2,8 +2,9 @@
 //!
 //! This crate implements a verifiable logical clock construct. The clock
 //! can be used in a peer-to-peer network to order events. Any node in the
-//! network can verify the correctness of the clock.
-
+//! network can verify the correctness of the clock. And HashMap as its core 
+//! data structure.
+pub mod ordinary_clock;
 use serde::{Deserialize, Serialize};
 use std::cmp;
 use std::collections::HashMap;