implement async qos_net::AsyncProxy

this includes:
 - Reaper::execute_async [qos_core]
 - AsyncStream [qos_core]
 - AsyncPool/AsyncStream [qos_core]
 - AsyncListener [qos_core]
 - AsyncRequestProcessor [qos_core]
 - AsyncProxyConnection [qos_net]
 - AsyncProxy [qos_net]
 - AsyncProxyStream [qos_net]

as well as fixing qos_next's qos_core "vm" feature dependency issue
by only requiring it if qos_net is requested with it

Author: Turnalek

src/Cargo.lock

@@ -21,7 +21,8 @@ aws-nitro-enclaves-nsm-api = { version = "0.4", default-features = false }
 serde_bytes = { version = "0.11", default-features = false }
 serde = { version = "1", features = ["derive"], default-features = false }
 
-tokio = { version = "1.38.0", features = ["io-util", "macros", "net", "rt-multi-thread", "time"], default-features = false }
+futures = { version = "0.3.31" }
+tokio = { version = "1.38.0", features = ["io-util", "macros", "net", "rt-multi-thread", "time"], default-features = false, optional = true}
 tokio-vsock = { version = "0.7.1", optional = true }
 
 [dev-dependencies]
@@ -32,7 +33,9 @@ rustls = { version = "0.23.5" }
 webpki-roots = { version = "0.26.1" }
 
 [features]
+default = ["async", "vm"]
+async = ["tokio", "tokio-vsock"]
 # Support for VSOCK
-vm = ["tokio-vsock"]
+vm = []
 # Never use in production - support for mock NSM
 mock = ["qos_nsm/mock"]

src/qos_core/Cargo.toml

@@ -0,0 +1,26 @@
+//! Streaming socket based client to connect with
+//! [`crate::server::SocketServer`].
+
+use crate::{client::ClientError, io::AsyncStreamPool};
+
+/// Client for communicating with the enclave [`crate::server::SocketServer`].
+pub struct AsyncClient {
+	pool: AsyncStreamPool,
+}
+
+impl AsyncClient {
+	/// Create a new client.
+	#[must_use]
+	pub fn new(pool: AsyncStreamPool) -> Self {
+		Self { pool }
+	}
+
+	/// Send raw bytes and wait for a response until the clients configured
+	/// timeout.
+	pub async fn call(&self, request: &[u8]) -> Result<Vec<u8>, ClientError> {
+		let mut stream = self.pool.get().await;
+		let resp = stream.call(request).await?;
+
+		Ok(resp)
+	}
+}

src/qos_core/src/async_client.rs

@@ -0,0 +1,89 @@
+//! Streaming socket based server for use in an enclave. Listens for connections
+//! from [`crate::client::Client`].
+
+use std::marker::PhantomData;
+
+use crate::{
+	io::{AsyncListener, AsyncStreamPool},
+	server::SocketServerError,
+};
+
+/// Something that can process requests in an async way.
+pub trait AsyncRequestProcessor: Send {
+	/// Process an incoming request and return a response in async.
+	///
+	/// The request and response are raw bytes. Likely this should be encoded
+	/// data and logic inside of this function should take care of decoding the
+	/// request and encoding a response.
+	fn process(
+		&self,
+		request: Vec<u8>,
+	) -> impl std::future::Future<Output = Vec<u8>> + Send;
+}
+
+/// A bare bones, socket based server.
+pub struct AsyncSocketServer<R: AsyncRequestProcessor> {
+	_phantom: PhantomData<R>,
+}
+
+impl<R> AsyncSocketServer<R>
+where
+	R: AsyncRequestProcessor + 'static + Clone,
+{
+	/// Listen and respond to incoming requests on all the pool's addresses with the given `processor`.
+	/// *NOTE*: the POOL_SIZE must match on both sides, since we expect ALL sockets to be connected
+	/// to right away (e.g. not on first use). The client side connect (above) will always connect them all.
+	pub async fn listen_all(
+		pool: AsyncStreamPool,
+		processor: R,
+	) -> Result<(), SocketServerError> {
+		println!("`AsyncSocketServer` listening on pool"); // TODO: add the addresses herei
+
+		let listeners = pool.listen().await?;
+
+		let mut tasks = Vec::new();
+		for listener in listeners {
+			let p = processor.clone();
+			let task =
+				tokio::spawn(async move { accept_loop(listener, p).await });
+
+			tasks.push(task);
+		}
+
+		// wait for ALL pool connections
+		let joined = futures::future::join_all(tasks).await;
+		for outer_result in joined {
+			match outer_result {
+				Err(_join_err) => {
+					return Err(SocketServerError::IOError(
+						crate::io::IOError::UnknownError, // TODO: add a join error translation to IOError
+					));
+				} // this really shouldn't happen
+				Ok(result) => result?,
+			}
+		}
+
+		Ok(())
+	}
+}
+
+async fn accept_loop<P>(
+	listener: AsyncListener,
+	processor: P,
+) -> Result<(), SocketServerError>
+where
+	P: AsyncRequestProcessor + Clone,
+{
+	loop {
+		let mut stream = listener.accept().await?;
+		loop {
+			match stream.recv().await {
+				Ok(payload) => {
+					let response = processor.process(payload).await;
+					let _ = stream.send(&response).await?;
+				}
+				Err(err) => return Err(SocketServerError::IOError(err)),
+			}
+		}
+	}
+}

src/qos_core/src/async_server.rs

@@ -9,9 +9,13 @@ use crate::{
 	io::SocketAddress,
 	parser::{GetParserForOptions, OptionsParser, Parser, Token},
 	reaper::Reaper,
-	EPHEMERAL_KEY_FILE, MANIFEST_FILE, PIVOT_FILE, QUORUM_FILE, SEC_APP_SOCK,
+	DEFAULT_POOL_SIZE, EPHEMERAL_KEY_FILE, MANIFEST_FILE, PIVOT_FILE,
+	QUORUM_FILE, SEC_APP_SOCK,
 };
 
+#[cfg(feature = "async")]
+use crate::io::AsyncStreamPool;
+
 /// "cid"
 pub const CID: &str = "cid";
 /// "port"
@@ -28,6 +32,8 @@ pub const EPHEMERAL_FILE_OPT: &str = "ephemeral-file";
 /// Name for the option to specify the manifest file.
 pub const MANIFEST_FILE_OPT: &str = "manifest-file";
 const APP_USOCK: &str = "app-usock";
+/// Name for the option to specify the maximum AsyncPool size.
+pub const POOL_SIZE: &str = "pool-size";
 
 /// CLI options for starting up the enclave server.
 #[derive(Default, Clone, Debug, PartialEq)]
@@ -44,11 +50,52 @@ impl EnclaveOpts {
 		Self { parsed }
 	}
 
+	/// Create a new [AsyncPool] of [AsyncStream] using the list of [SocketAddress] for the enclave server and
+	/// return the new [AsyncPool]. Analogous to [Self::addr] and [Self::app_addr] depending on the [app] parameter.
+	#[cfg(feature = "async")]
+	fn async_pool(&self, app: bool) -> AsyncStreamPool {
+		let pool_size: u32 = self
+			.parsed
+			.single(POOL_SIZE)
+			.expect("invalid pool options")
+			.parse()
+			.expect("invalid pool_size specified");
+		let usock_param = if app { APP_USOCK } else { USOCK };
+
+		match (
+			self.parsed.single(CID),
+			self.parsed.single(PORT),
+			self.parsed.single(usock_param),
+		) {
+			#[cfg(feature = "vm")]
+			(Some(c), Some(p), None) => {
+				let c = c.parse::<u32>().unwrap();
+				let start_port = p.parse::<u32>().unwrap();
+
+				let addresses = (start_port..start_port + pool_size).map(|p| {
+					SocketAddress::new_vsock(c, p, crate::io::VMADDR_NO_FLAGS)
+				});
+
+				AsyncStreamPool::new(addresses)
+			}
+			(None, None, Some(u)) => {
+				let addresses = (0..pool_size).map(|i| {
+					let u = format!("{}_{}", u, i); // add _X suffix for pooling
+					SocketAddress::new_unix(&u)
+				});
+
+				AsyncStreamPool::new(addresses)
+			}
+			_ => panic!("Invalid socket opts"),
+		}
+	}
+
 	/// Get the `SocketAddress` for the enclave server.
 	///
 	/// # Panics
 	///
 	/// Panics if the opts are not valid for exactly one of unix or vsock.
+	#[allow(unused)]
 	fn addr(&self) -> SocketAddress {
 		match (
 			self.parsed.single(CID),
@@ -66,6 +113,7 @@ impl EnclaveOpts {
 		}
 	}
 
+	#[allow(unused)]
 	fn app_addr(&self) -> SocketAddress {
 		SocketAddress::new_unix(
 			self.parsed
@@ -135,20 +183,48 @@ impl CLI {
 		} else if opts.parsed.help() {
 			println!("{}", opts.parsed.info());
 		} else {
-			Reaper::execute(
-				&Handles::new(
-					opts.ephemeral_file(),
-					opts.quorum_file(),
-					opts.manifest_file(),
-					opts.pivot_file(),
-				),
-				opts.nsm(),
-				opts.addr(),
-				opts.app_addr(),
-				None,
-			);
+			Self::run(opts)
 		}
 	}
+
+	#[cfg(not(feature = "async"))]
+	fn run(opts: EnclaveOpts) {
+		Reaper::execute(
+			&Handles::new(
+				opts.ephemeral_file(),
+				opts.quorum_file(),
+				opts.manifest_file(),
+				opts.pivot_file(),
+			),
+			opts.nsm(),
+			opts.addr(),
+			opts.app_addr(),
+			None,
+		);
+	}
+
+	#[cfg(feature = "async")]
+	fn run(opts: EnclaveOpts) {
+		tokio::runtime::Builder::new_current_thread()
+			.enable_all()
+			.build()
+			.expect("tokio main to run")
+			.block_on(async {
+				Reaper::async_execute(
+					&Handles::new(
+						opts.ephemeral_file(),
+						opts.quorum_file(),
+						opts.manifest_file(),
+						opts.pivot_file(),
+					),
+					opts.nsm(),
+					opts.async_pool(false),
+					opts.async_pool(true),
+					None,
+				)
+				.await;
+			});
+	}
 }
 
 /// Parser for enclave CLI
@@ -201,6 +277,11 @@ impl GetParserForOptions for EnclaveParser {
 					.takes_value(true)
 					.default_value(SEC_APP_SOCK)
 			)
+			.token(
+				Token::new(POOL_SIZE, "the pool size for use with all socket types")
+					.takes_value(true)
+					.default_value(DEFAULT_POOL_SIZE)
+			)
 	}
 }
 
@@ -281,6 +362,20 @@ mod test {
 		assert_eq!(opts.addr(), SocketAddress::new_unix("./test.sock"));
 	}
 
+	#[test]
+	#[cfg(feature = "async")]
+	fn parse_pool_size() {
+		let mut args: Vec<_> =
+			vec!["binary", "--usock", "./test.sock", "--pool-size", "7"]
+				.into_iter()
+				.map(String::from)
+				.collect();
+		let opts = EnclaveOpts::new(&mut args);
+
+		let pool = opts.async_pool(false);
+		assert_eq!(pool.len(), 7);
+	}
+
 	#[test]
 	fn parse_manifest_file() {
 		let mut args: Vec<_> = vec!["binary", "--usock", "./test.sock"]