qos_core: add enclave timout test

also fix small nits

Author: Turnalek

src/init/init.rs

@@ -70,7 +70,7 @@ async fn main() {
 	);
 
 	const START_PORT: u32 = 3;
-	const INITIAL_POOL_SIZE: u32 = 1; // start at pool size 1, grow based on  manifest/args as necessary (see Reaper)
+	const INITIAL_POOL_SIZE: u8 = 1; // start at pool size 1, grow based on  manifest/args as necessary (see Reaper)
 	let core_pool = StreamPool::new(
 		SocketAddress::new_vsock(cid, START_PORT, VMADDR_NO_FLAGS),
 		INITIAL_POOL_SIZE,

src/integration/tests/enclave_app_client_socket_stress.rs

@@ -103,7 +103,7 @@ async fn enclave_app_client_socket_stress() {
 		StreamPool::single(SocketAddress::new_unix(ENCLAVE_SOCK)).unwrap();
 	let enclave_client = SocketClient::new(
 		enclave_client_pool.shared(),
-		INITIAL_CLIENT_TIMEOUT + Duration::from_secs(3), // needs to be bigger than the sdlow request below + some time for recovery
+		INITIAL_CLIENT_TIMEOUT + Duration::from_secs(3), // needs to be bigger than the slow request below + some time for recovery
 	);
 
 	let app_request =

src/integration/tests/reaper.rs

@@ -1,13 +1,17 @@
-use std::fs;
+use std::{fs, time::Duration};
 
 use integration::{
-	wait_for_usock, PIVOT_ABORT_PATH, PIVOT_OK_PATH, PIVOT_PANIC_PATH,
-	PIVOT_POOL_SIZE_PATH,
+	wait_for_usock, PivotSocketStressMsg, PIVOT_ABORT_PATH, PIVOT_OK_PATH,
+	PIVOT_PANIC_PATH, PIVOT_POOL_SIZE_PATH, PIVOT_SOCKET_STRESS_PATH,
 };
 use qos_core::{
+	client::SocketClient,
 	handles::Handles,
 	io::{SocketAddress, StreamPool},
-	protocol::services::boot::ManifestEnvelope,
+	protocol::{
+		msg::ProtocolMsg, services::boot::ManifestEnvelope, ProtocolError,
+		ProtocolPhase,
+	},
 	reaper::{Reaper, REAPER_EXIT_DELAY},
 };
 use qos_nsm::mock::MockNsm;
@@ -16,7 +20,6 @@ use qos_test_primitives::PathWrapper;
 #[tokio::test]
 async fn reaper_works() {
 	let secret_path: PathWrapper = "/tmp/reaper_works.secret".into();
-	// let eph_path = "reaper_works.eph.key";
 	let usock: PathWrapper = "/tmp/reaper_works.sock".into();
 	let manifest_path: PathWrapper = "/tmp/reaper_works.manifest".into();
 	let msg = "durp-a-durp";
@@ -36,16 +39,15 @@ async fn reaper_works() {
 	let mut manifest_envelope = ManifestEnvelope::default();
 	manifest_envelope.manifest.pivot.args =
 		vec!["--msg".to_string(), msg.to_string()];
-	manifest_envelope.manifest.client_timeout_ms = Some(2000); // check if this gets applied
 
 	handles.put_manifest_envelope(&manifest_envelope).unwrap();
 	assert!(handles.pivot_exists());
 
 	let enclave_pool =
-		StreamPool::new(SocketAddress::new_unix(&usock), 1).unwrap();
+		StreamPool::single(SocketAddress::new_unix(&usock)).unwrap();
 
 	let app_pool =
-		StreamPool::new(SocketAddress::new_unix("./never.sock"), 1).unwrap();
+		StreamPool::single(SocketAddress::new_unix("./never.sock")).unwrap();
 
 	let reaper_handle = tokio::spawn(async move {
 		Reaper::execute(
@@ -59,7 +61,7 @@ async fn reaper_works() {
 	});
 
 	// Give the enclave server time to bind to the socket
-	tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+	wait_for_usock(&usock).await;
 
 	// Check that the reaper is still running, presumably waiting for
 	// the secret.
@@ -76,6 +78,94 @@ async fn reaper_works() {
 	assert!(fs::remove_file(integration::PIVOT_OK_SUCCESS_FILE).is_ok());
 }
 
+#[tokio::test]
+async fn reaper_timeout_works() {
+	let secret_path: PathWrapper = "/tmp/reaper_timeout_works.secret".into();
+	let enclave_sock: PathWrapper = "/tmp/reaper_timeout_works.sock".into();
+	let app_sock: PathWrapper = "/tmp/reaper_timeout_works_app.sock".into();
+	let manifest_path: PathWrapper =
+		"/tmp/reaper_timeout_works.manifest".into();
+
+	// clean up old manifest if it's left from a panic
+	drop(std::fs::remove_file(&*manifest_path));
+
+	// For our sanity, ensure the secret does not yet exist
+	drop(fs::remove_file(&*secret_path));
+
+	let handles = Handles::new(
+		"eph_path".to_string(),
+		(*secret_path).to_string(),
+		(*manifest_path).to_string(),
+		PIVOT_SOCKET_STRESS_PATH.to_string(),
+	);
+
+	// Make sure we have written everything necessary to pivot, except the
+	// quorum key
+	let mut manifest_envelope = ManifestEnvelope::default();
+	// Tell pivot where to open up the server app socket
+	manifest_envelope.manifest.pivot.args = vec![app_sock.to_string()];
+
+	// we'll be checking if this is set by passing slow and fast requests
+	manifest_envelope.manifest.client_timeout_ms = Some(2000);
+
+	handles.put_manifest_envelope(&manifest_envelope).unwrap();
+	assert!(handles.pivot_exists());
+
+	let enclave_pool =
+		StreamPool::single(SocketAddress::new_unix(&enclave_sock)).unwrap();
+
+	let app_pool =
+		StreamPool::single(SocketAddress::new_unix(&app_sock)).unwrap();
+
+	let reaper_handle = tokio::spawn(async move {
+		Reaper::execute(
+			&handles,
+			Box::new(MockNsm),
+			enclave_pool,
+			app_pool,
+			Some(ProtocolPhase::QuorumKeyProvisioned),
+		)
+		.await;
+	});
+
+	// Give the enclave server time to bind to the socket
+	wait_for_usock(&enclave_sock).await;
+
+	// Check that the reaper is still running, presumably waiting for
+	// the secret.
+	assert!(!reaper_handle.is_finished());
+
+	// Create the file with the secret, which should cause the reaper
+	// to start executable.
+	fs::write(&*secret_path, b"super dank tank secret tech").unwrap();
+
+	// Give the app server time to bind to the socket
+	wait_for_usock(&app_sock).await;
+
+	// create a "slow" app request longer than client timeout from `Manifest`, but longer than 5s timeout on our local client.
+	let app_request =
+		borsh::to_vec(&PivotSocketStressMsg::SlowRequest(3000)).unwrap();
+	let request =
+		borsh::to_vec(&ProtocolMsg::ProxyRequest { data: app_request })
+			.unwrap();
+
+	// ensure our client to the enclave has longer timeout than the configured 2s and the slow request 3s
+	let client = SocketClient::single(
+		SocketAddress::new_unix(&enclave_sock),
+		Duration::from_millis(5000),
+	)
+	.unwrap();
+
+	let response: ProtocolMsg =
+		borsh::from_slice(&client.call(&request).await.unwrap()).unwrap();
+
+	// The response should be AppClientRecvTimeout which indicates the enclave short-circuited the timeout
+	assert_eq!(
+		response,
+		ProtocolMsg::ProtocolErrorResponse(ProtocolError::AppClientRecvTimeout)
+	);
+}
+
 #[tokio::test]
 async fn reaper_handles_non_zero_exits() {
 	let secret_path: PathWrapper =
@@ -117,7 +207,7 @@ async fn reaper_handles_non_zero_exits() {
 	});
 
 	// Give the enclave server time to bind to the socket
-	tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+	wait_for_usock(&usock).await;
 
 	// Check that the reaper is still running, presumably waiting for
 	// the secret.
@@ -174,7 +264,7 @@ async fn reaper_handles_panic() {
 	});
 
 	// Give the enclave server time to bind to the socket
-	tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+	wait_for_usock(&usock).await;
 
 	// Check that the reaper is still running, presumably waiting for
 	// the secret.
@@ -195,11 +285,10 @@ async fn reaper_handles_panic() {
 async fn reaper_handles_pool_size() {
 	let secret_path: PathWrapper =
 		"/tmp/reaper_handles_pool_size.secret".into();
-	// let eph_path = "reaper_works.eph.key";
 	let usock: PathWrapper = "/tmp/reaper_handles_pool_size.sock".into();
 	let manifest_path: PathWrapper =
 		"/tmp/reaper_handles_pool_size.manifest".into();
-	let msg = "5"; // must match pool-size in manifest bellow (test thing)
+	let msg = "5"; // must match pool-size in manifest below (test thing)
 
 	// For our sanity, ensure the secret does not yet exist
 	drop(fs::remove_file(&*secret_path));

src/integration/tests/simple_socket_stress.rs

@@ -22,8 +22,6 @@ async fn simple_socket_stress() {
 
 	wait_for_usock(SOCKET_STRESS_SOCK).await;
 
-	// needs to be long enough for process exit to register and not cause a timeout
-
 	let app_pool =
 		StreamPool::new(SocketAddress::new_unix(SOCKET_STRESS_SOCK), 1)
 			.unwrap();

src/qos_core/src/io/pool.rs

@@ -65,7 +65,7 @@ impl StreamPool {
 	/// Create a new `StreamPool` with given starting `SocketAddress`, timeout and number of addresses to populate.
 	pub fn new(
 		start_address: SocketAddress,
-		mut count: u32,
+		mut count: u8,
 	) -> Result<Self, IOError> {
 		eprintln!("StreamPool start address: {start_address}");
 

src/qos_net/src/cli.rs

@@ -37,7 +37,7 @@ impl ProxyOpts {
 	pub(crate) fn async_pool(
 		&self,
 	) -> Result<StreamPool, qos_core::io::IOError> {
-		let pool_size: u32 = self
+		let pool_size: u8 = self
 			.parsed
 			.single(POOL_SIZE)
 			.expect("invalid pool options")