Add YAML config support for transport configuration in build.rs

implants/Cargo.toml

@@ -114,6 +114,7 @@ russh-keys = "0.37.1"
 rust-embed = "8.5.0"
 serde = "1.0"
 serde_json = "1.0.87"
+serde_yaml = "0.9"
 sha1 = "0.10.5"
 sha2 = "0.10.7"
 sha256 = { version = "1.0.3", default-features = false }
@@ -134,6 +135,7 @@ tonic = { git = "https://github.com/hyperium/tonic.git", rev = "07e4ee1" }
 tonic-build = { git = "https://github.com/hyperium/tonic.git", rev = "c783652" } # Needed git for `.codec_path` in build.rs - previous version of codec setting is really gross. https://github.com/hyperium/tonic/blob/ea8cd3f384e953e177f20a62aa156a75676853f4/examples/build.rs#L44
 trait-variant = "0.1.1"
 uuid = "1.5.0"
+urlencoding = "2.1.3"
 static_vcruntime = "2.0"
 url = "2.5"
 which = "4.4.2"

implants/lib/eldritchv2/stdlib/tests/Cargo.toml

@@ -9,6 +9,6 @@ eldritch-core = { workspace = true, default-features = false }
 anyhow = { workspace = true }
 glob = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
-serde_yaml = "0.9"
+serde_yaml = {workspace = true}
 spin = { version = "0.10.0", features = ["rwlock"] }
 log = { workspace = true }

implants/lib/pb/Cargo.toml

@@ -40,3 +40,6 @@ which = { workspace = true }
 home = "=0.5.11"
 reqwest = { workspace = true, features = ["blocking", "json", "rustls-tls"] }
 serde_json = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_yaml = { workspace = true }
+urlencoding = { workspace = true }

implants/lib/pb/build.rs

@@ -1,8 +1,166 @@
+use serde::Deserialize;
 use std::collections::HashMap;
 use std::env;
 use std::path::PathBuf;
 use which::which;
 
+#[derive(Debug, Deserialize)]
+struct TransportConfig {
+    #[serde(rename = "URI")]
+    uri: String,
+    #[serde(rename = "type")]
+    transport_type: String,
+    extra: String,
+    #[serde(default)]
+    interval: Option<u64>,
+}
+
+#[derive(Debug, Deserialize)]
+struct YamlConfig {
+    transports: Vec<TransportConfig>,
+    #[serde(default)]
+    server_pubkey: Option<String>,
+}
+
+fn parse_yaml_config() -> Result<bool, Box<dyn std::error::Error>> {
+    // Check if IMIX_CONFIG is set
+    let config_yaml = match std::env::var("IMIX_CONFIG") {
+        Ok(yaml_content) => yaml_content,
+        Err(_) => return Ok(false), // No config set, return false
+    };
+
+    // Check that other configuration options are not set
+    let has_callback_uri = std::env::var("IMIX_CALLBACK_URI").is_ok();
+    let has_callback_interval = std::env::var("IMIX_CALLBACK_INTERVAL").is_ok();
+    let has_transport_extra = std::env::vars().any(|(k, _)| k.starts_with("IMIX_TRANSPORT_EXTRA_"));
+
+    if has_callback_uri || has_callback_interval || has_transport_extra {
+        let mut error_msg = String::from(
+            "Configuration error: Cannot use IMIX_CONFIG with other configuration options.\n",
+        );
+        error_msg.push_str(
+            "When IMIX_CONFIG is set, all configuration must be done through the YAML file.\n",
+        );
+        error_msg.push_str("Found one or more of:\n");
+
+        if has_callback_uri {
+            error_msg.push_str("  - IMIX_CALLBACK_URI\n");
+        }
+        if has_callback_interval {
+            error_msg.push_str("  - IMIX_CALLBACK_INTERVAL\n");
+        }
+        if has_transport_extra {
+            error_msg.push_str("  - IMIX_TRANSPORT_EXTRA_*\n");
+        }
+
+        error_msg.push_str(
+            "\nPlease use ONLY the YAML config file OR use environment variables, but not both.",
+        );
+
+        return Err(error_msg.into());
+    }
+
+    // Parse the YAML config
+    let config: YamlConfig = serde_yaml::from_str(&config_yaml)
+        .map_err(|e| format!("Failed to parse YAML config: {}", e))?;
+
+    // Validate that we have at least one transport
+    if config.transports.is_empty() {
+        return Err("YAML config must contain at least one transport".into());
+    }
+
+    // Build DSN string from transports
+    let mut dsn_parts = Vec::new();
+
+    for transport in &config.transports {
+        // Validate transport type
+        let transport_type_lower = transport.transport_type.to_lowercase();
+        if !["grpc", "http1", "dns"].contains(&transport_type_lower.as_str()) {
+            return Err(format!(
+                "Invalid transport type '{}'. Must be one of: GRPC, http1, DNS",
+                transport.transport_type
+            )
+            .into());
+        }
+
+        // Validate that extra is valid JSON
+        if !transport.extra.is_empty() {
+            serde_json::from_str::<serde_json::Value>(&transport.extra).map_err(|e| {
+                format!(
+                    "Invalid JSON in 'extra' field for transport '{}': {}",
+                    transport.uri, e
+                )
+            })?;
+        }
+
+        // Error if URI already contains query parameters
+        if transport.uri.contains('?') {
+            return Err(format!("URI '{}' already contains query parameters. Query parameters should not be present in the URI field.", transport.uri).into());
+        }
+
+        // Map transport type to appropriate schema
+        let schema = match transport_type_lower.as_str() {
+            "grpc" => "grpc",
+            "http1" => "http",
+            "dns" => "dns",
+            _ => unreachable!(), // Already validated above
+        };
+
+        // Strip any existing schema from the URI and replace with the correct one
+        let uri_without_schema = transport
+            .uri
+            .split_once("://")
+            .map(|(_, rest)| rest)
+            .unwrap_or(&transport.uri);
+
+        // Build DSN part with correct schema and query parameters
+        let mut dsn_part = format!("{}://{}", schema, uri_without_schema);
+
+        // Add query parameters
+        dsn_part.push('?');
+        let mut params = Vec::new();
+
+        // Add interval if present
+        if let Some(interval) = transport.interval {
+            params.push(format!("interval={}", interval));
+        }
+
+        // Add extra as query parameter if not empty
+        if !transport.extra.is_empty() {
+            let encoded_extra = urlencoding::encode(&transport.extra);
+            params.push(format!("extra={}", encoded_extra));
+        }
+
+        if !params.is_empty() {
+            dsn_part.push_str(&params.join("&"));
+        } else {
+            // Remove the trailing '?' if no params were added
+            dsn_part.pop();
+        }
+
+        dsn_parts.push(dsn_part);
+    }
+
+    // Join all DSN parts with semicolons
+    let dsn = dsn_parts.join(";");
+
+    // Emit the DSN configuration
+    println!("cargo:rustc-env=IMIX_CALLBACK_URI={}", dsn);
+
+    // Emit server_pubkey if present
+    if let Some(ref pubkey) = config.server_pubkey {
+        println!("cargo:rustc-env=IMIX_SERVER_PUBKEY={}", pubkey);
+        println!("cargo:warning=Using server_pubkey from YAML config");
+    }
+
+    println!(
+        "cargo:warning=Successfully parsed YAML config with {} transport(s)",
+        config.transports.len()
+    );
+
+    Ok(true)
+}
+
 fn get_pub_key() {
     // Check if IMIX_SERVER_PUBKEY is already set
     if std::env::var("IMIX_SERVER_PUBKEY").is_ok() {
@@ -95,6 +253,12 @@ fn build_extra_vars() -> Result<(), Box<dyn std::error::Error>> {
 }
 
 fn validate_dsn_config() -> Result<(), Box<dyn std::error::Error>> {
+    // Skip validation if YAML config is being used
+    // (parse_yaml_config already handles validation in that case)
+    if std::env::var("IMIX_CONFIG").is_ok() {
+        return Ok(());
+    }
+
     // Check if IMIX_CALLBACK_URI contains query parameters
     let callback_uri =
         std::env::var("IMIX_CALLBACK_URI").unwrap_or_else(|_| "http://127.0.0.1:8000".to_string());
@@ -126,7 +290,12 @@ fn validate_dsn_config() -> Result<(), Box<dyn std::error::Error>> {
 }
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // Parse YAML config if present (this will emit IMIX_CALLBACK_URI if successful)
+    parse_yaml_config()?;
+
+    // Validate DSN config (skips if YAML config was used)
     validate_dsn_config()?;
+
     get_pub_key();
     build_extra_vars()?;