agentmux_srv/backend/
container_credential.rs

1//! A scoped credential for container agents, in place of the instance auth key.
2//!
3//! A container agent is meant to see only its `/workspace` mount, but the
4//! instance key (`AGENTMUX_AUTH_KEY`) reaches every route srv serves: host
5//! shell creation, `stream-local-file`, editor reads, attachment ingest by
6//! path. Its `docker exec` env therefore carries a token minted here instead,
7//! which srv accepts only on [`container_route_allowed`] and attributes to the
8//! agent it was minted for.
9
10use sha2::{Digest, Sha256};
11use std::collections::HashMap;
12use std::sync::{Mutex, OnceLock};
13
14/// Prefix that marks a container token in logs and bug reports.
15pub const TOKEN_PREFIX: &str = "amxc_";
16
17/// What a container token was minted for.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct ContainerGrant {
20    pub block_id: String,
21    /// The agent's own `AGENTMUX_AGENT_TOKEN` when the token was last handed
22    /// out. `caller_middleware` attributes every request to it, ignoring any
23    /// `X-Agent-Token` header the container sends.
24    pub agent_token: Option<String>,
25}
26
27#[derive(Default)]
28struct Registry {
29    by_digest: HashMap<[u8; 32], ContainerGrant>,
30    by_block: HashMap<String, String>,
31}
32
33fn registry() -> &'static Mutex<Registry> {
34    static REGISTRY: OnceLock<Mutex<Registry>> = OnceLock::new();
35    REGISTRY.get_or_init(Default::default)
36}
37
38// Looked up by digest so a lookup's timing depends on the digest, not on how
39// much of a guessed token matches a real one.
40fn digest(token: &str) -> [u8; 32] {
41    Sha256::digest(token.as_bytes()).into()
42}
43
44fn mint() -> String {
45    format!(
46        "{TOKEN_PREFIX}{}{}",
47        uuid::Uuid::new_v4().simple(),
48        uuid::Uuid::new_v4().simple()
49    )
50}
51
52/// The block's container token: minted on first use, reused on later turns,
53/// with the grant's identity refreshed from this turn's env.
54pub fn token_for_block(block_id: &str, agent_token: Option<&str>) -> String {
55    let mut reg = registry().lock().unwrap();
56    let token = reg.by_block.get(block_id).cloned().unwrap_or_else(mint);
57    reg.by_block.insert(block_id.to_string(), token.clone());
58    reg.by_digest.insert(
59        digest(&token),
60        ContainerGrant {
61            block_id: block_id.to_string(),
62            agent_token: agent_token.map(str::trim).filter(|t| !t.is_empty()).map(str::to_string),
63        },
64    );
65    token
66}
67
68pub fn grant_for(token: &str) -> Option<ContainerGrant> {
69    if !token.starts_with(TOKEN_PREFIX) {
70        return None;
71    }
72    registry().lock().unwrap().by_digest.get(&digest(token)).cloned()
73}
74
75/// Invalidate the block's token; its container can no longer reach srv.
76pub fn revoke_block(block_id: &str) {
77    let mut reg = registry().lock().unwrap();
78    if let Some(token) = reg.by_block.remove(block_id) {
79        reg.by_digest.remove(&digest(&token));
80    }
81}
82
83/// The routes a container agent may call. Everything else is refused, notably
84/// host shell/pty creation, file streaming, editor and attachment paths, pane
85/// and agent opening, UI automation, cron, and the `/ws` and
86/// `/agentmux/service` RPC surfaces. Matched on the raw path, exactly as the
87/// router matches it.
88pub fn container_route_allowed(path: &str) -> bool {
89    const EXACT: &[&str] = &[
90        "/agentmux/reactive/inject",
91        "/agentmux/reactive/agent",
92        "/agentmux/reactive/agent-names",
93        "/agentmux/discovery",
94        "/agentmux/work",
95        "/agentmux/work/claim",
96        "/api/v1/self",
97        "/api/v1/agent/memory/list",
98        "/api/v1/agent/memory/read",
99        "/api/v1/agent/memory/write",
100        "/api/v1/agent/memory/history",
101        "/api/v1/agent/memory/diff",
102        "/api/v1/agent/memory/revert",
103        // Read-only: Global Memory is composed into every agent's context at
104        // launch, so writes from a sandbox would reach host agents.
105        "/api/v1/agent/globalmemory/list",
106        "/api/v1/agent/globalmemory/read",
107        "/api/v1/agent/globalmemory/history",
108        "/api/v1/agent/globalmemory/diff",
109        "/api/v1/agent/self/quit",
110        "/api/v1/agent/pane/close",
111    ];
112    if EXACT.contains(&path) {
113        return true;
114    }
115    let is_id = |s: &str| !s.is_empty() && s.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_');
116    if let Some(rest) = path.strip_prefix("/agentmux/work/") {
117        return match rest.split('/').collect::<Vec<_>>().as_slice() {
118            [id] => is_id(id),
119            [id, "heartbeat" | "complete" | "release"] => is_id(id),
120            _ => false,
121        };
122    }
123    if let Some(id) = path.strip_prefix("/api/v1/agent/shutdown/") {
124        return is_id(id);
125    }
126    false
127}
128
129/// Routes that act on the calling agent's own data, and so are refused to a
130/// container token that carries no agent identity rather than falling back to
131/// an agent named in the request body.
132pub fn route_needs_identity(path: &str) -> bool {
133    path.starts_with("/api/v1/agent/memory/")
134}
135
136/// A container turn's `docker exec` env: the container denylist applied, and
137/// the instance key replaced by this block's container token.
138pub fn container_exec_env(block_id: &str, env_vars: &HashMap<String, String>) -> Vec<(String, String)> {
139    let agent_token = env_vars.get("AGENTMUX_AGENT_TOKEN").map(String::as_str);
140    // The publish guard's hooks directory is a host path; the agent's other
141    // GIT_CONFIG_* settings still apply inside the container.
142    let env_vars = crate::backend::publish_guard::without_hooks_path(env_vars);
143    let build_info = crate::backend::pane_env::build_info();
144    let is_build_info = |k: &str| build_info.iter().any(|(b, _)| *b == k);
145    let mut out: Vec<(String, String)> = env_vars
146        .iter()
147        .filter(|(k, _)| !crate::backend::container::CONTAINER_ENV_DENYLIST.contains(&k.as_str()))
148        .filter(|(k, _)| !is_build_info(k))
149        .map(|(k, v)| {
150            if k == "AGENTMUX_AUTH_KEY" {
151                (k.clone(), token_for_block(block_id, agent_token))
152            } else {
153                (k.clone(), v.clone())
154            }
155        })
156        .collect();
157    // Which AgentMux this is, as every host pane is told: a container exec
158    // doesn't go through `pane_env`'s sanitizers, and srv's values win over
159    // any the agent's own env carried.
160    out.extend(build_info.iter().map(|(k, v)| (k.to_string(), v.clone())));
161    out
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    #[test]
169    fn tokens_are_reused_per_block_and_revocable() {
170        let a = token_for_block("cc-block-a", Some("tok-a"));
171        assert!(a.starts_with(TOKEN_PREFIX));
172        assert_eq!(token_for_block("cc-block-a", Some("tok-a2")), a);
173        assert_ne!(token_for_block("cc-block-b", None), a);
174        let grant = grant_for(&a).unwrap();
175        assert_eq!(grant.block_id, "cc-block-a");
176        assert_eq!(grant.agent_token.as_deref(), Some("tok-a2"));
177        revoke_block("cc-block-a");
178        assert_eq!(grant_for(&a), None);
179        assert_ne!(token_for_block("cc-block-a", None), a);
180    }
181
182    #[test]
183    fn unknown_and_unprefixed_tokens_have_no_grant() {
184        assert_eq!(grant_for("amxc_not-a-real-token"), None);
185        assert_eq!(grant_for("0b8e0c52-1b4e-4d3f-9a53-8a0c1d2e3f40"), None);
186        assert_eq!(grant_for(""), None);
187    }
188
189    #[test]
190    fn blank_agent_token_counts_as_none() {
191        let t = token_for_block("cc-block-blank", Some("  "));
192        assert_eq!(grant_for(&t).unwrap().agent_token, None);
193    }
194
195    #[test]
196    fn allowlist_admits_agent_routes_only() {
197        for ok in [
198            "/agentmux/reactive/inject",
199            "/agentmux/discovery",
200            "/agentmux/work",
201            "/agentmux/work/claim",
202            "/agentmux/work/7f3a/heartbeat",
203            "/agentmux/work/7f3a",
204            "/api/v1/self",
205            "/api/v1/agent/memory/write",
206            "/api/v1/agent/globalmemory/read",
207            "/api/v1/agent/pane/close",
208            "/api/v1/agent/shutdown/req-12",
209        ] {
210            assert!(container_route_allowed(ok), "{ok} should be allowed");
211        }
212        for denied in [
213            "/api/v1/shell/create",
214            "/api/v1/ptyshell/create",
215            "/agentmux/stream-local-file",
216            "/agentmux/service",
217            "/ws",
218            "/api/v1/pane/open",
219            "/api/v1/agent/open",
220            "/api/v1/attachments/upload",
221            "/api/v1/ui/browser/eval",
222            "/api/v1/ui/screenshot",
223            "/agentmux/cron",
224            "/api/v1/fleet/bulk-stop",
225            "/api/v1/agent/globalmemory/write",
226            "/api/v1/agent/globalmemory/remove",
227            "/agentmux/reactive/agents",
228            "/agentmux/work/7f3a/steal",
229            "/agentmux/work/../service",
230            "/api/v1/agent/shutdown/",
231            "/api/v1/self/",
232            "//api/v1/self",
233        ] {
234            assert!(!container_route_allowed(denied), "{denied} should be refused");
235        }
236    }
237
238    #[test]
239    fn exec_env_swaps_the_instance_key_for_the_container_token() {
240        let env: HashMap<String, String> = [
241            ("AGENTMUX_AUTH_KEY", "instance-key-value"),
242            ("AGENTMUX_AGENT_TOKEN", "agent-tok"),
243            ("AGENTMUX_LOCAL_URL", "http://host.docker.internal:1"),
244            ("HOME", "/host/home"),
245        ]
246        .into_iter()
247        .map(|(k, v)| (k.to_string(), v.to_string()))
248        .collect();
249        let out: HashMap<_, _> = container_exec_env("cc-block-env", &env).into_iter().collect();
250        let token = &out["AGENTMUX_AUTH_KEY"];
251        assert!(token.starts_with(TOKEN_PREFIX));
252        assert!(!out.values().any(|v| v == "instance-key-value"));
253        assert_eq!(grant_for(token).unwrap().agent_token.as_deref(), Some("agent-tok"));
254        assert_eq!(out["AGENTMUX_LOCAL_URL"], "http://host.docker.internal:1");
255        assert!(!out.contains_key("HOME"), "the denylist still applies");
256    }
257
258    #[test]
259    fn exec_env_without_a_key_gets_no_token() {
260        let env: HashMap<String, String> = [("FOO".to_string(), "bar".to_string())].into_iter().collect();
261        let out = container_exec_env("cc-block-nokey", &env);
262        assert!(out.iter().all(|(k, _)| k != "AGENTMUX_AUTH_KEY"));
263    }
264
265    #[test]
266    fn a_container_agent_is_told_which_build_it_runs_in() {
267        let env: HashMap<String, String> = [("AGENTMUX_VERSION".to_string(), "0.0.1-stale".to_string())].into_iter().collect();
268        let out = container_exec_env("cc-block-build", &env);
269        let versions: Vec<&str> = out.iter().filter(|(k, _)| k == "AGENTMUX_VERSION").map(|(_, v)| v.as_str()).collect();
270        assert_eq!(versions, [env!("CARGO_PKG_VERSION")], "srv's value, once, over the agent's own");
271        assert!(out.iter().any(|(k, _)| k == "AGENTMUX_BUILD"));
272    }
273}
agentmux_srv/backend/container_credential.rs

agentmux_srv/backend/
container_credential.rs
