Add acceptor config changes to CLI

crates/ev-cli/src/commands/enclave/init.rs

@@ -1,11 +1,13 @@
+use std::num::{NonZeroU32, NonZeroU64};
+
 use clap::{ArgGroup, Parser};
 use common::api::BasicAuth;
 use common::{api::AuthMode, CliError};
 use ev_enclave::api::enclave::{Enclave, EnclaveApi};
 use ev_enclave::cert::{create_new_cert, DesiredLifetime, DistinguishedName};
 use ev_enclave::config::{
-    default_dockerfile, EgressSettings, EnclaveConfig, HealthcheckConfig, ScalingSettings,
-    ServiceSettings, SigningInfo,
+    default_dockerfile, AcceptorConfig, EgressSettings, EnclaveConfig, HealthcheckConfig,
+    ScalingSettings, ServiceSettings, SigningInfo,
 };
 
 /// Initialize an Enclave.toml in the current directory
@@ -93,6 +95,18 @@ pub struct InitArgs {
     /// The port that all incoming traffic should be forwarded to within the Enclave.
     #[arg(long = "port")]
     pub port: Option<u16>,
+
+    /// Cap on the number of in-flight connections the Enclave will accept concurrently.
+    #[arg(long)]
+    pub max_concurrent_connections: Option<NonZeroU32>,
+
+    /// Cap on the number of in-flight TLS handshakes the Enclave will perform concurrently. Must not exceed --max-concurrent-connections.
+    #[arg(long)]
+    pub max_concurrent_handshakes: Option<NonZeroU32>,
+
+    /// Per-handshake timeout in milliseconds for incoming TLS handshakes.
+    #[arg(long)]
+    pub handshake_timeout: Option<NonZeroU64>,
 }
 
 impl std::convert::From<InitArgs> for EnclaveConfig {
@@ -115,6 +129,19 @@ impl std::convert::From<InitArgs> for EnclaveConfig {
             val.healthcheck.map(HealthcheckConfig::from)
         };
 
+        let acceptor = if val.max_concurrent_connections.is_none()
+            && val.max_concurrent_handshakes.is_none()
+            && val.handshake_timeout.is_none()
+        {
+            None
+        } else {
+            Some(AcceptorConfig {
+                max_concurrent_connections: val.max_concurrent_connections,
+                max_concurrent_handshakes: val.max_concurrent_handshakes,
+                handshake_timeout: val.handshake_timeout,
+            })
+        };
+
         EnclaveConfig {
             name: val.enclave_name,
             uuid: None,
@@ -138,6 +165,7 @@ impl std::convert::From<InitArgs> for EnclaveConfig {
             healthcheck,
             service: val.port.map(ServiceSettings::new),
             attestation_cors: None,
+            acceptor,
         }
     }
 }
@@ -171,6 +199,14 @@ async fn init_local_config(init_args: InitArgs, created_enclave: Enclave) -> exi
     let config_path = output_path.join("enclave.toml");
 
     let mut initial_config: EnclaveConfig = init_args.into();
+
+    if let Some(acceptor) = initial_config.acceptor.as_ref() {
+        if let Err(e) = acceptor.validate() {
+            log::error!("{e}");
+            return e.exitcode();
+        }
+    }
+
     initial_config.annotate(created_enclave);
 
     if initial_config.signing.is_none() {
@@ -237,6 +273,9 @@ mod init_tests {
             healthcheck: None,
             healthcheck_port: None,
             port: None,
+            max_concurrent_connections: None,
+            max_concurrent_handshakes: None,
+            handshake_timeout: None,
         }
     }
 
@@ -383,6 +422,59 @@ keyPath = "./key.pem"
         assert_eq!(config_content, expected_config_content);
     }
 
+    #[tokio::test]
+    async fn init_local_config_test_with_acceptor_config() {
+        let output_dir = TempDir::new().unwrap();
+        let sample_enclave = Enclave {
+            uuid: "1234".into(),
+            name: "hello-enclave".into(),
+            team_uuid: "1234".into(),
+            app_uuid: "1234".into(),
+            domain: "hello.com".into(),
+            state: EnclaveState::Pending,
+            created_at: "00:00:00".into(),
+            updated_at: "00:00:00".into(),
+        };
+        let mut init_args = default_init_args(&output_dir);
+        init_args.max_concurrent_connections = NonZeroU32::new(100);
+        init_args.max_concurrent_handshakes = NonZeroU32::new(10);
+        init_args.handshake_timeout = NonZeroU64::new(10000);
+        init_local_config(init_args, sample_enclave).await;
+        let config_path = output_dir.path().join("enclave.toml");
+        assert!(config_path.exists());
+        let config_content = String::from_utf8(read(config_path).unwrap()).unwrap();
+        let expected_config_content = r#"version = 1
+name = "hello"
+uuid = "1234"
+app_uuid = "1234"
+team_uuid = "1234"
+debug = false
+dockerfile = "Dockerfile"
+api_key_auth = true
+trx_logging = true
+tls_termination = true
+forward_proxy_protocol = false
+trusted_headers = ["X-Evervault-*"]
+
+[egress]
+enabled = true
+destinations = ["evervault.com"]
+
+[scaling]
+desired_replicas = 2
+
+[signing]
+certPath = "./cert.pem"
+keyPath = "./key.pem"
+
+[acceptor]
+max_concurrent_connections = 100
+max_concurrent_handshakes = 10
+handshake_timeout = 10000
+"#;
+        assert_eq!(config_content, expected_config_content);
+    }
+
     #[tokio::test]
     async fn init_local_config_test_with_service_port() {
         let output_dir = TempDir::new().unwrap();

crates/ev-enclave/src/build/mod.rs

@@ -255,7 +255,7 @@ async fn process_dockerfile<R: AsyncRead + std::marker::Unpin>(
         // catch for edge case where port is defined in the toml, but was never exposed in docker
         if let Some(false) = is_configured_port_missing_in_docker {
             log::warn!(
-              "Found service port in enclave.toml which is not exposed in the supplied Dockerfile. This may suggest a misconfiguration. The build will continue using the service port defined in the enclave.toml file.", 
+              "Found service port in enclave.toml which is not exposed in the supplied Dockerfile. This may suggest a misconfiguration. The build will continue using the service port defined in the enclave.toml file.",
             );
             log::warn!(
                 "Service port from enclave.toml: {}",
@@ -341,6 +341,29 @@ async fn process_dockerfile<R: AsyncRead + std::marker::Unpin>(
         dataplane_info["healthcheck_use_tls"] = json!(build_config.use_tls_healthcheck());
     }
 
+    if let Some(acceptor) = build_config.acceptor() {
+        let mut acc = serde_json::Map::new();
+        if let Some(c) = acceptor.max_concurrent_connections {
+            acc.insert("max_concurrent_connections".into(), json!(c));
+        }
+        if let Some(h) = acceptor.max_concurrent_handshakes {
+            acc.insert("max_concurrent_handshakes".into(), json!(h));
+        }
+        if let Some(ms) = acceptor.handshake_timeout {
+            let ms = ms.get();
+            acc.insert(
+                "handshake_timeout".into(),
+                json!({
+                    "secs": ms / 1000,
+                    "nanos": (ms % 1000) * 1_000_000,
+                }),
+            );
+        }
+        if !acc.is_empty() {
+            dataplane_info["acceptor"] = Value::Object(acc);
+        }
+    }
+
     let dataplane_env = format!(
         "echo {} > /etc/dataplane-config.json",
         dataplane_info.to_string().replace('"', "\\\"")
@@ -513,6 +536,8 @@ mod test {
     use crate::version::EnclaveRuntime;
     use common::enclave::types::AttestationCors;
     use std::iter::zip;
+    use std::num::NonZeroU32;
+    use std::num::NonZeroU64;
     use tempfile::TempDir;
 
     fn get_config(egress_enabled: bool) -> ValidatedEnclaveBuildConfig {
@@ -550,6 +575,7 @@ mod test {
             attestation_cors: Some(AttestationCors {
                 origin: "test.com".to_string(),
             }),
+            acceptor: None,
         }
     }
 
@@ -1366,6 +1392,76 @@ ENTRYPOINT ["/bootstrap", "1>&2"]
         }
     }
 
+    #[tokio::test]
+    async fn test_process_dockerfile_with_acceptor_config() {
+        let sample_dockerfile_contents = r#"FROM alpine
+RUN touch /hello-script;\
+    /bin/sh -c "echo -e '"'#!/bin/sh\nwhile true; do echo "hello"; sleep 2; done;\n'"' > /hello-script"
+ENTRYPOINT ["sh", "/hello-script"]"#;
+        let mut readable_contents = sample_dockerfile_contents.as_bytes();
+
+        let mut config = get_config(false);
+        config.acceptor = Some(crate::config::AcceptorConfig {
+            max_concurrent_connections: NonZeroU32::new(100),
+            max_concurrent_handshakes: NonZeroU32::new(10),
+            handshake_timeout: NonZeroU64::new(10000),
+        });
+
+        let enclave_runtime = EnclaveRuntime {
+            data_plane_version: "0.0.0".to_string(),
+            installer_version: "abcdef".to_string(),
+        };
+
+        let processed_file =
+            process_dockerfile(&config, &mut readable_contents, &enclave_runtime, false)
+                .await
+                .unwrap();
+
+        let dataplane_directive = processed_file
+            .iter()
+            .map(|directive| directive.to_string())
+            .find(|directive| directive.contains("/etc/dataplane-config.json"))
+            .expect("dataplane-config directive present");
+
+        let expected = r#"RUN echo {\"acceptor\":{\"handshake_timeout\":{\"nanos\":0,\"secs\":10},\"max_concurrent_connections\":100,\"max_concurrent_handshakes\":10},\"api_key_auth\":true,\"attestation_cors\":{\"origin\":\"test.com\"},\"forward_proxy_protocol\":false,\"trusted_headers\":[\"X-Evervault-*\"],\"trx_logging_enabled\":true} > /etc/dataplane-config.json"#;
+        assert_eq!(dataplane_directive, expected);
+    }
+
+    #[tokio::test]
+    async fn test_process_dockerfile_with_sub_second_handshake_timeout() {
+        let sample_dockerfile_contents = r#"FROM alpine
+RUN touch /hello-script;\
+    /bin/sh -c "echo -e '"'#!/bin/sh\nwhile true; do echo "hello"; sleep 2; done;\n'"' > /hello-script"
+ENTRYPOINT ["sh", "/hello-script"]"#;
+        let mut readable_contents = sample_dockerfile_contents.as_bytes();
+
+        let mut config = get_config(false);
+        config.acceptor = Some(crate::config::AcceptorConfig {
+            max_concurrent_connections: None,
+            max_concurrent_handshakes: None,
+            handshake_timeout: std::num::NonZeroU64::new(500),
+        });
+
+        let enclave_runtime = EnclaveRuntime {
+            data_plane_version: "0.0.0".to_string(),
+            installer_version: "abcdef".to_string(),
+        };
+
+        let processed_file =
+            process_dockerfile(&config, &mut readable_contents, &enclave_runtime, false)
+                .await
+                .unwrap();
+
+        let dataplane_directive = processed_file
+            .iter()
+            .map(|directive| directive.to_string())
+            .find(|directive| directive.contains("/etc/dataplane-config.json"))
+            .expect("dataplane-config directive present");
+
+        assert!(dataplane_directive
+            .contains(r#"\"acceptor\":{\"handshake_timeout\":{\"nanos\":500000000,\"secs\":0}}"#));
+    }
+
     #[tokio::test]
     #[serial_test::serial]
     async fn test_choose_output_dir() {