fix(agent): receive broker capabilities over private startup input

Signed-off-by: Alex Rosenzweig <arosenzweig@squareup.com>

crates/buzz-agent/src/catalog.rs

@@ -143,7 +143,7 @@ pub async fn discover_databricks_models_with_cache_dir(
 ) -> Result<Vec<ModelEntry>, AgentError> {
     let token_source = if matches!(cfg.provider, Provider::Databricks | Provider::DatabricksV2)
         && cfg.api_key.is_empty()
-        && std::env::var_os(crate::sandbox_runtime::BROKER_ENV).is_none()
+        && !crate::sandbox_runtime::configured()?
     {
         crate::auth::PkceOAuthTokenSource::new(crate::llm::databricks_pkce_config(
             &cfg.base_url,

crates/buzz-agent/src/lib.rs

@@ -153,6 +153,7 @@ fn die(msg: String) -> ! {
 }
 
 pub fn run() -> Result<(), Box<dyn std::error::Error>> {
+    sandbox_runtime::initialize()?;
     let args: Vec<String> = std::env::args().collect();
     if matches!(args.get(1).map(String::as_str), Some("auth")) {
         return tokio::runtime::Builder::new_multi_thread()

crates/buzz-agent/src/sandbox_runtime.rs

@@ -14,14 +14,25 @@ use axum::{
 use serde::{Deserialize, Serialize};
 #[cfg(test)]
 use std::path::Path;
-use std::{sync::Arc, time::Duration};
+use std::{
+    io::Read,
+    sync::{Arc, OnceLock},
+    time::Duration,
+};
 
-/// Launcher-supplied JSON capability for a fixed loopback credential broker.
+/// Nonsecret startup marker. Set to `stdin-v1` and send a four-byte big-endian
+/// length followed by the JSON [`BrokerConfig`] before ACP input. Never put the
+/// capability in the environment or arguments.
+///
+/// The trusted launcher must isolate the agent from tool process inspection.
+/// On macOS this requires inherited kernel denials for other-process `process-info*` and
+/// `mach-priv-task-port` and `kern.procargs*` sysctl reads, including processes
+/// in the same sandbox. The startup channel does not replace OS isolation.
 pub const BROKER_ENV: &str = "BUZZ_SANDBOX_AUTH_BROKER";
 /// Path to the immutable certificate snapshot supplied by the launcher.
 pub const ROOTS_ENV: &str = "BUZZ_SANDBOX_TLS_ROOTS";
 
-/// Connection capability passed only to the protected agent, never its tools.
+/// Connection capability delivered privately before any tools are started.
 #[derive(Clone, Serialize, Deserialize)]
 pub struct BrokerConfig {
     /// Loopback TCP port selected by the trusted launcher.
@@ -45,6 +56,55 @@ struct BrokerState {
     source: Arc<dyn TokenSource>,
 }
 
+static BROKER: OnceLock<BrokerConfig> = OnceLock::new();
+const MAX_BOOTSTRAP_BYTES: usize = 16 * 1024;
+
+pub(crate) fn initialize() -> Result<(), AgentError> {
+    let Some(marker) = std::env::var_os(BROKER_ENV) else {
+        return Ok(());
+    };
+    if marker != "stdin-v1" {
+        return Err(AgentError::Llm(
+            "broker capability must be delivered through stdin-v1".into(),
+        ));
+    }
+    // Unlike clearing environ, this denies same-UID reads of the initial
+    // environment, memory, and descriptors through procfs/ptrace on Linux.
+    #[cfg(target_os = "linux")]
+    nix::sys::prctl::set_dumpable(false)
+        .map_err(|_| AgentError::Llm("cannot protect broker process from inspection".into()))?;
+    if !cfg!(any(target_os = "linux", target_os = "macos")) {
+        return Err(AgentError::Llm(
+            "credential broker unsupported on this platform".into(),
+        ));
+    }
+    let broker = read_bootstrap(&mut std::io::stdin().lock())?;
+    BROKER
+        .set(broker)
+        .map_err(|_| AgentError::Llm("broker already initialized".into()))
+}
+
+fn read_bootstrap(reader: &mut impl Read) -> Result<BrokerConfig, AgentError> {
+    let invalid = |_| AgentError::Llm("invalid broker startup frame".into());
+    let mut header = [0; 4];
+    reader.read_exact(&mut header).map_err(invalid)?;
+    let size = u32::from_be_bytes(header) as usize;
+    if size == 0 || size > MAX_BOOTSTRAP_BYTES {
+        return Err(AgentError::Llm("invalid broker startup frame size".into()));
+    }
+    let mut bytes = vec![0; size];
+    reader.read_exact(&mut bytes).map_err(invalid)?;
+    serde_json::from_slice(&bytes)
+        .map_err(|_| AgentError::Llm("invalid broker startup configuration".into()))
+}
+
+pub(crate) fn configured() -> Result<bool, AgentError> {
+    if BROKER.get().is_none() && std::env::var_os(BROKER_ENV).is_some() {
+        return Err(AgentError::Llm("broker startup was not initialized".into()));
+    }
+    Ok(BROKER.get().is_some())
+}
+
 /// Native trust is captured before confinement. Verification remains enabled,
 /// using rustls locally instead of macOS trustd (which can bypass network policy).
 pub fn http_builder() -> Result<reqwest::ClientBuilder, AgentError> {
@@ -60,7 +120,10 @@ fn builder_with_roots(
 }
 
 pub(crate) fn token_source(cfg: &Config) -> Result<Option<Arc<dyn TokenSource>>, AgentError> {
-    let Some(raw) = std::env::var_os(BROKER_ENV) else {
+    let Some(broker) = BROKER.get() else {
+        if std::env::var_os(BROKER_ENV).is_some() {
+            return Err(AgentError::Llm("broker startup was not initialized".into()));
+        }
         return Ok(None);
     };
     if !matches!(cfg.provider, Provider::Databricks | Provider::DatabricksV2)
@@ -68,8 +131,6 @@ pub(crate) fn token_source(cfg: &Config) -> Result<Option<Arc<dyn TokenSource>>,
     {
         return Ok(None);
     }
-    let broker: BrokerConfig = serde_json::from_str(&raw.to_string_lossy())
-        .map_err(|_| AgentError::Llm("invalid sandbox auth broker".into()))?;
     if broker.host != cfg.base_url || broker.port == 0 || broker.secret.len() < 32 {
         return Err(AgentError::Llm(
             "sandbox auth broker provider mismatch".into(),
@@ -80,7 +141,10 @@ pub(crate) fn token_source(cfg: &Config) -> Result<Option<Arc<dyn TokenSource>>,
         .timeout(Duration::from_secs(60))
         .build()
         .map_err(|_| AgentError::Llm("sandbox auth client unavailable".into()))?;
-    Ok(Some(Arc::new(BrokerClient { broker, client })))
+    Ok(Some(Arc::new(BrokerClient {
+        broker: broker.clone(),
+        client,
+    })))
 }
 struct BrokerClient {
     broker: BrokerConfig,
@@ -176,6 +240,39 @@ async fn token(
 #[cfg(test)]
 mod tests {
     use super::*;
+    #[test]
+    fn bootstrap_preserves_following_acp_input() {
+        let config = BrokerConfig {
+            port: 1234,
+            secret: "synthetic-capability".into(),
+            host: "https://model.example".into(),
+        };
+        let bytes = serde_json::to_vec(&config).unwrap();
+        let mut frame = (bytes.len() as u32).to_be_bytes().to_vec();
+        frame.extend(bytes);
+        frame.extend(b"{\"jsonrpc\":\"2.0\"}\n");
+        let mut reader = std::io::Cursor::new(frame);
+        assert_eq!(read_bootstrap(&mut reader).unwrap().secret, config.secret);
+        let mut acp = String::new();
+        reader.read_to_string(&mut acp).unwrap();
+        assert_eq!(acp, "{\"jsonrpc\":\"2.0\"}\n");
+    }
+
+    #[test]
+    fn bootstrap_rejects_invalid_frames_without_echoing_their_contents() {
+        for frame in [
+            vec![],
+            vec![0, 0],
+            0u32.to_be_bytes().to_vec(),
+            ((MAX_BOOTSTRAP_BYTES + 1) as u32).to_be_bytes().to_vec(),
+            [8u32.to_be_bytes().as_slice(), b"private"].concat(),
+            [7u32.to_be_bytes().as_slice(), b"private"].concat(),
+        ] {
+            let error = read_bootstrap(&mut frame.as_slice()).err().unwrap();
+            assert!(!error.to_string().contains("private"));
+        }
+    }
+
     #[tokio::test]
     async fn broker_rejects_unauthorized_callers_and_serves_fixed_source() {
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();

crates/buzz-agent/tests/broker_startup.rs

@@ -0,0 +1,98 @@
+//! The real executable consumes broker startup before ACP or provider creation.
+use serde_json::json;
+use std::{
+    io::Write,
+    process::{Command, Stdio},
+};
+
+fn agent() -> Command {
+    let mut cmd = Command::new(env!("CARGO_BIN_EXE_buzz-agent"));
+    cmd.env_clear()
+        .env("PATH", std::env::var_os("PATH").unwrap_or_default())
+        .env("HOME", std::env::var_os("HOME").unwrap_or_default())
+        .env("BUZZ_AGENT_PROVIDER", "openai-compat")
+        .env("OPENAI_COMPAT_BASE_URL", "http://127.0.0.1:1")
+        .env("OPENAI_COMPAT_API_KEY", "synthetic")
+        .env("BUZZ_AGENT_MODEL", "synthetic-model")
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped());
+    cmd
+}
+
+fn frame() -> Vec<u8> {
+    let data = serde_json::to_vec(&json!({"port":1234,
+        "secret":"synthetic-broker-capability-at-least-32-chars", "host":"https://model.example"}))
+    .unwrap();
+    let mut bytes = (data.len() as u32).to_be_bytes().to_vec();
+    bytes.extend(data);
+    bytes
+}
+
+#[test]
+fn private_bootstrap_preserves_real_acp_initialize() {
+    let mut cmd = agent();
+    cmd.env("BUZZ_SANDBOX_AUTH_BROKER", "stdin-v1");
+    let mut child = cmd.spawn().unwrap();
+    let mut input = frame();
+    input.extend(
+        serde_json::to_vec(&json!({"jsonrpc":"2.0", "id":1,
+        "method":"initialize", "params":{"protocolVersion":1,"clientCapabilities":{}}}))
+        .unwrap(),
+    );
+    input.push(b'\n');
+    child.stdin.take().unwrap().write_all(&input).unwrap();
+    let output = child.wait_with_output().unwrap();
+    assert!(
+        output.status.success(),
+        "{}",
+        String::from_utf8_lossy(&output.stderr)
+    );
+    let messages: Vec<serde_json::Value> = String::from_utf8(output.stdout)
+        .unwrap()
+        .lines()
+        .map(|line| serde_json::from_str(line).unwrap())
+        .collect();
+    assert!(
+        messages
+            .iter()
+            .any(|v| v["id"] == 1 && v.get("result").is_some()),
+        "{messages:?}"
+    );
+}
+
+#[test]
+fn legacy_environment_and_truncated_bootstrap_fail_closed() {
+    for (marker, input, expected) in [
+        (
+            r#"{"secret":"private-legacy-secret"}"#,
+            vec![],
+            "must be delivered through stdin-v1",
+        ),
+        ("stdin-v1", vec![0, 0], "invalid broker startup frame"),
+    ] {
+        let mut cmd = agent();
+        cmd.env("BUZZ_SANDBOX_AUTH_BROKER", marker);
+        let mut child = cmd.spawn().unwrap();
+        child.stdin.take().unwrap().write_all(&input).unwrap();
+        let output = child.wait_with_output().unwrap();
+        assert!(!output.status.success());
+        let error = String::from_utf8(output.stderr).unwrap();
+        assert!(error.contains(expected), "{error}");
+        assert!(!error.contains("private-legacy-secret"));
+        assert!(output.stdout.is_empty());
+    }
+}
+
+#[test]
+fn initialized_broker_still_rejects_wrong_provider_host() {
+    let mut cmd = agent();
+    cmd.env("BUZZ_SANDBOX_AUTH_BROKER", "stdin-v1")
+        .env("BUZZ_AGENT_PROVIDER", "databricks")
+        .env("DATABRICKS_HOST", "https://other.example");
+    let mut child = cmd.spawn().unwrap();
+    child.stdin.take().unwrap().write_all(&frame()).unwrap();
+    let output = child.wait_with_output().unwrap();
+    assert!(!output.status.success());
+    assert!(String::from_utf8_lossy(&output.stderr).contains("broker provider mismatch"));
+}