diff --git a/Cargo.lock b/Cargo.lock index 99e0e295..37921f49 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2340,6 +2340,7 @@ dependencies = [ "tonic", "tracing", "url", + "uuid", ] [[package]] diff --git a/crates/forge_app/Cargo.toml b/crates/forge_app/Cargo.toml index 7b0452c9..e92b9527 100644 --- a/crates/forge_app/Cargo.toml +++ b/crates/forge_app/Cargo.toml @@ -37,6 +37,7 @@ backon.workspace = true sha2.workspace = true hex.workspace = true dashmap.workspace = true +uuid.workspace = true url.workspace = true reqwest.workspace = true diff --git a/crates/forge_app/src/lib.rs b/crates/forge_app/src/lib.rs index f98bfda6..5a40f627 100644 --- a/crates/forge_app/src/lib.rs +++ b/crates/forge_app/src/lib.rs @@ -20,6 +20,7 @@ mod operation; mod orch; #[cfg(test)] mod orch_spec; +pub mod redaction; mod retry; mod search_dedup; mod services; diff --git a/crates/forge_app/src/redaction.rs b/crates/forge_app/src/redaction.rs new file mode 100644 index 00000000..c408203b --- /dev/null +++ b/crates/forge_app/src/redaction.rs @@ -0,0 +1,258 @@ +//! Redaction layer: pipe sensitive tool data by *reference*, not by value. +//! +//! A tool result containing secrets is sealed into a salted hash commitment and +//! kept out-of-band (e.g. in the content-addressed [`crate::infra::KVStore`] +//! vault). The model only ever sees the [`Handle`] — a commitment plus a +//! non-revealing label — so the raw bytes never enter its context. A later tool +//! call references the handle and the executor opens it from the vault, so data +//! flows tool -> vault -> tool. The commitment is *binding* (a revealed value can +//! be checked against it) and *hiding* (the nonce keeps a low-entropy secret +//! unguessable) — the commit/reveal primitive underneath zero-knowledge +//! protocols, where a property of the data can be proven without revealing it. +//! +//! This module is the pure, deterministic core. Randomness for the nonce and the +//! vault storage/wiring into the tool executor live at the call site. + +use std::sync::Arc; + +use dashmap::DashMap; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +/// Token prefix the model sees in place of redacted data. +const PREFIX: &str = "zkref:sha256:"; +/// Hex length of a SHA-256 digest. +const DIGEST_HEX_LEN: usize = 64; + +/// A salted hash commitment to some bytes: `sha256(nonce || data)`, hex-encoded. +/// +/// Binding (no other value matches a given commitment) and hiding (the nonce +/// defeats guessing a low-entropy value), so it can be shown to the model +/// without revealing the data it commits to. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Commitment(String); + +impl Commitment { + /// Commit to `data` under `nonce`. The same `(data, nonce)` always yields the + /// same commitment; a different nonce yields a different commitment for the + /// same data (which is what makes it hiding). + pub fn new(data: &[u8], nonce: &[u8]) -> Self { + let mut hasher = Sha256::new(); + hasher.update(nonce); + hasher.update(data); + Commitment(hex::encode(hasher.finalize())) + } + + /// Verify that `data` opened under `nonce` matches this commitment (binding). + pub fn verify(&self, data: &[u8], nonce: &[u8]) -> bool { + Self::new(data, nonce) == *self + } + + /// The hex digest. + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// A reference the model carries in place of the raw data. It reveals only a +/// non-sensitive `label` (chosen by the caller, e.g. byte length or kind) and +/// the commitment — never the bytes. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Handle { + pub commitment: Commitment, + pub label: String, +} + +impl Handle { + /// Seal `data` (under `nonce`) into a handle carrying only a non-revealing + /// `label`. The caller is responsible for keeping the label non-sensitive. + pub fn seal(data: &[u8], nonce: &[u8], label: impl Into) -> Self { + Handle { + commitment: Commitment::new(data, nonce), + label: label.into(), + } + } + + /// The token the model sees, e.g. `zkref:sha256:ab12..#bytes=2048`. + pub fn token(&self) -> String { + format!("{}{}#{}", PREFIX, self.commitment.as_str(), self.label) + } +} + +/// Extract every handle commitment referenced in a blob of text (e.g. a tool +/// call's arguments), so the executor can resolve them from the vault before the +/// tool runs — piping the data by reference instead of through the model. +pub fn referenced_commitments(text: &str) -> Vec { + let mut out = Vec::new(); + let mut rest = text; + while let Some(idx) = rest.find(PREFIX) { + let after = &rest[idx + PREFIX.len()..]; + let digest: String = after.chars().take_while(|c| c.is_ascii_hexdigit()).collect(); + if digest.len() == DIGEST_HEX_LEN { + out.push(Commitment(digest.clone())); + } + // Always advance past the prefix we just consumed so the loop terminates. + rest = &after[digest.len()..]; + } + out +} + +struct VaultEntry { + nonce: Vec, + data: Vec, +} + +/// Session-scoped vault holding the plaintext for sealed values, keyed by their +/// commitment. The model carries only [`Handle`] tokens; the executor calls +/// [`RedactionVault::resolve_json`] / [`RedactionVault::resolve`] to swap the +/// tokens back to the real bytes right before a tool runs, so a sealed value +/// flows tool -> vault -> tool and never enters the model's context. Cloning +/// shares the same backing store (it is `Arc`-backed). +#[derive(Default, Clone)] +pub struct RedactionVault { + entries: Arc>, +} + +impl RedactionVault { + pub fn new() -> Self { + Self::default() + } + + /// Seal `data` behind a handle the model can carry. `label` must be chosen + /// non-sensitive (e.g. byte length); a fresh random nonce makes the + /// commitment hiding even for a low-entropy secret. + pub fn seal(&self, data: &[u8], label: impl Into) -> Handle { + let nonce = uuid::Uuid::new_v4().into_bytes(); + let handle = Handle::seal(data, &nonce, label); + self.entries.insert( + handle.commitment.as_str().to_string(), + VaultEntry { nonce: nonce.to_vec(), data: data.to_vec() }, + ); + handle + } + + /// Open a sealed value by its commitment, re-verifying the binding before + /// releasing the bytes. Returns `None` if the commitment is unknown. + pub fn open(&self, commitment: &Commitment) -> Option> { + let entry = self.entries.get(commitment.as_str())?; + (Commitment::new(&entry.data, &entry.nonce) == *commitment).then(|| entry.data.clone()) + } + + /// Replace every known handle token in `text` with its sealed plaintext. + /// Unknown handles and non-UTF-8 values are left verbatim. + pub fn resolve(&self, text: &str) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + while let Some(idx) = rest.find(PREFIX) { + out.push_str(&rest[..idx]); + let after = &rest[idx + PREFIX.len()..]; + let digest: String = after.chars().take_while(|c| c.is_ascii_hexdigit()).collect(); + let mut consumed = digest.len(); + if after[consumed..].starts_with('#') { + consumed += 1 + + after[consumed + 1..] + .chars() + .take_while(|c| { + c.is_ascii_alphanumeric() || matches!(c, '=' | '.' | '_' | '-') + }) + .count(); + } + let opened = (digest.len() == DIGEST_HEX_LEN) + .then(|| self.open(&Commitment(digest.clone()))) + .flatten(); + match opened.as_deref().and_then(|b| std::str::from_utf8(b).ok()) { + Some(plain) => out.push_str(plain), + None => { + out.push_str(PREFIX); + out.push_str(&after[..consumed]); + } + } + rest = &after[consumed..]; + } + out.push_str(rest); + out + } + + /// Resolve handles inside every string leaf of a JSON value, preserving + /// structure — safe to run on tool-call arguments. + pub fn resolve_json(&self, value: &Value) -> Value { + match value { + Value::String(s) => Value::String(self.resolve(s)), + Value::Array(items) => { + Value::Array(items.iter().map(|v| self.resolve_json(v)).collect()) + } + Value::Object(map) => Value::Object( + map.iter().map(|(k, v)| (k.clone(), self.resolve_json(v))).collect(), + ), + other => other.clone(), + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::{assert_eq, assert_ne}; + + use super::{Commitment, Handle, referenced_commitments}; + + const NONCE: &[u8] = b"0123456789abcdef"; + const SECRET: &[u8] = b"AKIA_SUPER_SECRET_TOKEN_value_4242"; + + #[test] + fn test_commitment_is_binding() { + let commitment = Commitment::new(SECRET, NONCE); + // The true opening verifies. + assert!(commitment.verify(SECRET, NONCE)); + // A tampered value does not. + assert!(!commitment.verify(b"AKIA_SUPER_SECRET_TOKEN_value_4243", NONCE)); + // The right value under the wrong nonce does not. + assert!(!commitment.verify(SECRET, b"different-nonce-")); + } + + #[test] + fn test_commitment_is_deterministic() { + let actual = Commitment::new(SECRET, NONCE); + let expected = Commitment::new(SECRET, NONCE); + assert_eq!(actual, expected); + } + + #[test] + fn test_commitment_is_hiding() { + // Different nonces hide the same data behind different commitments, yet + // each opens to the same secret. + let a = Commitment::new(SECRET, b"nonce-a-aaaaaaaa"); + let b = Commitment::new(SECRET, b"nonce-b-bbbbbbbb"); + assert_ne!(a, b); + assert!(a.verify(SECRET, b"nonce-a-aaaaaaaa")); + assert!(b.verify(SECRET, b"nonce-b-bbbbbbbb")); + } + + #[test] + fn test_handle_token_never_reveals_the_secret() { + let handle = Handle::seal(SECRET, NONCE, "bytes=34"); + let token = handle.token(); + let secret_text = std::str::from_utf8(SECRET).unwrap(); + // The model sees only the commitment + label, never the plaintext. + assert!(!token.contains(secret_text)); + assert!(!token.contains("SUPER_SECRET")); + assert!(token.starts_with("zkref:sha256:")); + assert!(token.ends_with("#bytes=34")); + } + + #[test] + fn test_referenced_commitments_round_trip() { + let handle = Handle::seal(SECRET, NONCE, "bytes=34"); + // A later tool call carries the handle token in its arguments. + let tool_args = format!("{{\"upload\": \"{}\", \"dest\": \"s3://bucket\"}}", handle.token()); + + let actual = referenced_commitments(&tool_args); + let expected = vec![handle.commitment.clone()]; + assert_eq!(actual, expected); + } + + #[test] + fn test_referenced_commitments_none_when_absent() { + let actual = referenced_commitments("a plain tool result with no references"); + assert!(actual.is_empty()); + } +} diff --git a/crates/forge_app/src/tool_registry.rs b/crates/forge_app/src/tool_registry.rs index 97eb70bc..1f5456c4 100644 --- a/crates/forge_app/src/tool_registry.rs +++ b/crates/forge_app/src/tool_registry.rs @@ -30,6 +30,11 @@ pub struct ToolRegistry { agent_executor: AgentExecutor, mcp_executor: McpExecutor, services: Arc, + /// Session vault for the zk-redaction layer: any sealed value referenced as a + /// handle in a tool call's arguments is resolved back to its real bytes here, + /// just before dispatch, so credentials/values pipe by reference instead of + /// through the model. + vault: crate::redaction::RedactionVault, } impl> ToolRegistry { @@ -39,9 +44,25 @@ impl> ToolReg tool_executor: ToolExecutor::new(services.clone()), agent_executor: AgentExecutor::new(services.clone()), mcp_executor: McpExecutor::new(services.clone()), + vault: crate::redaction::RedactionVault::new(), } } + /// Swap any zk-handle tokens in a tool call's arguments for their sealed + /// values from the vault, preserving JSON structure. Called after the + /// arguments are logged (so the log keeps handles, not secrets) and before + /// the tool is dispatched. + fn resolve_handles(&self, mut input: ToolCallFull) -> ToolCallFull { + use forge_domain::ToolCallArguments; + input.arguments = match input.arguments { + ToolCallArguments::Parsed(value) => { + ToolCallArguments::Parsed(self.vault.resolve_json(&value)) + } + ToolCallArguments::Unparsed(raw) => ToolCallArguments::Unparsed(self.vault.resolve(&raw)), + }; + input + } + /// Threads the trajectory backend down to the embedded `AgentExecutor` /// so child-agent dispatches can construct their own recorders. See /// `ForgeApp::with_trajectory_repo` for the full plumbing path. @@ -123,6 +144,11 @@ impl> ToolReg tracing::info!(tool_name = %input.name, arguments = %input.arguments.clone().into_string(), "Executing tool call"); let tool_name = input.name.clone(); + // Resolve any zk-handles to their real values now (after logging the + // handle-bearing args, before dispatch) so sealed data reaches the tool + // without ever entering the model's context. + let input = self.resolve_handles(input); + // First, try to call a Forge tool if ToolCatalog::contains(&input.name) { let tool_input: ToolCatalog = ToolCatalog::try_from(input)?; diff --git a/crates/forge_app/tests/redaction_pipe.rs b/crates/forge_app/tests/redaction_pipe.rs new file mode 100644 index 00000000..ef547a3f --- /dev/null +++ b/crates/forge_app/tests/redaction_pipe.rs @@ -0,0 +1,80 @@ +//! Integration test for the redaction (zk-pipe) primitive. +//! +//! Runs as a separate test binary linking only the `forge_app` lib, so it +//! executes even while the crate's in-module unit-test build is blocked by +//! unrelated pre-existing code. Proves the real public API by running it. + +use forge_app::redaction::{Commitment, Handle, RedactionVault, referenced_commitments}; +use serde_json::json; + +const NONCE: &[u8] = b"0123456789abcdef"; +const SECRET: &[u8] = b"AKIA_SUPER_SECRET_TOKEN_value_4242"; + +#[test] +fn commitment_binds_and_hides() { + let commitment = Commitment::new(SECRET, NONCE); + // Binding: only the true opening verifies. + assert!(commitment.verify(SECRET, NONCE)); + assert!(!commitment.verify(b"tampered-value", NONCE)); + // Hiding: a different nonce gives a different commitment for the same data. + assert_ne!(commitment, Commitment::new(SECRET, b"another-nonce-16")); +} + +#[test] +fn handle_pipes_data_by_reference_without_revealing_it() { + // A tool produced a secret; seal it into a handle the model can carry. + let handle = Handle::seal(SECRET, NONCE, "bytes=34"); + let token = handle.token(); + + // The model only ever sees the commitment + a non-revealing label. + let secret_text = std::str::from_utf8(SECRET).unwrap(); + assert!(!token.contains(secret_text)); + assert!(token.starts_with("zkref:sha256:")); + + // A later tool call references the handle; the executor recovers the + // commitment to resolve the data from the vault — piping by reference. + let tool_args = format!("{{\"upload\":\"{}\"}}", token); + let referenced = referenced_commitments(&tool_args); + assert_eq!(referenced, vec![handle.commitment.clone()]); + + // The recovered commitment still binds to the real bytes (vault open + verify). + assert!(referenced[0].verify(SECRET, NONCE)); +} + +#[test] +fn vault_reuses_a_credential_by_reference() { + let vault = RedactionVault::new(); + let credential = b"sk-live-CREDENTIAL-do-not-leak"; + + // A secrets source seals the credential; the model only ever sees the handle. + let handle = vault.seal(credential, "bytes=30"); + let token = handle.token(); + assert!(!token.contains("CREDENTIAL")); + + // The vault opens the handle back to the real bytes (binding re-verified). + assert_eq!(vault.open(&handle.commitment).unwrap(), credential); + + // A later tool call carries the handle; the executor resolves it to the real + // value before the tool runs — the credential is reused without ever being + // revealed to the model. + let args = format!("curl -H 'Authorization: Bearer {}' https://api", token); + let resolved = vault.resolve(&args); + assert!(resolved.contains("Bearer sk-live-CREDENTIAL-do-not-leak")); + assert!(!resolved.contains("zkref:")); +} + +#[test] +fn vault_resolve_json_preserves_structure_and_leaves_unknown_handles() { + let vault = RedactionVault::new(); + let handle = vault.seal(b"VALUE42", "bytes=7"); + + // Unknown handle (not in this vault) is left untouched. + let unknown = format!("zkref:sha256:{}#bytes=1", "a".repeat(64)); + assert_eq!(vault.resolve(&unknown), unknown); + + // JSON structure preserved; only the matching string leaf is resolved. + let args = json!({ "cmd": format!("echo {}", handle.token()), "count": 1 }); + let resolved = vault.resolve_json(&args); + assert_eq!(resolved["cmd"], json!("echo VALUE42")); + assert_eq!(resolved["count"], json!(1)); +}