Skip to main content

lean_ctx/core/
a2a_transport.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5const MAX_ENVELOPE_BYTES: usize = 2_000_000;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct AgentIdentityV1 {
9    pub agent_id: String,
10    pub agent_type: String,
11    pub daemon_fingerprint: String,
12    pub capabilities: Vec<String>,
13}
14
15impl AgentIdentityV1 {
16    pub fn from_current(agent_id: &str, agent_type: &str) -> Self {
17        Self {
18            agent_id: agent_id.to_string(),
19            agent_type: agent_type.to_string(),
20            daemon_fingerprint: compute_daemon_fingerprint(),
21            capabilities: default_capabilities(),
22        }
23    }
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct TransportEnvelopeV1 {
28    pub format_version: u32,
29    pub sent_at: DateTime<Utc>,
30    pub sender: AgentIdentityV1,
31    pub recipient: Option<String>,
32    pub content_type: TransportContentType,
33    pub payload_json: String,
34    pub signature: Option<String>,
35    pub metadata: HashMap<String, String>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
39#[serde(rename_all = "snake_case")]
40pub enum TransportContentType {
41    HandoffBundle,
42    ContextPackage,
43    A2AMessage,
44    A2ATask,
45}
46
47impl TransportEnvelopeV1 {
48    pub fn new(
49        sender: AgentIdentityV1,
50        recipient: Option<&str>,
51        content_type: TransportContentType,
52        payload_json: String,
53    ) -> Self {
54        Self {
55            format_version: 1,
56            sent_at: Utc::now(),
57            sender,
58            recipient: recipient.map(std::string::ToString::to_string),
59            content_type,
60            payload_json,
61            signature: None,
62            metadata: HashMap::new(),
63        }
64    }
65
66    pub fn sign(&mut self, secret: &[u8]) {
67        let mac_bytes = self.compute_hmac(secret);
68        let mut sig = String::with_capacity(mac_bytes.len() * 2);
69        for b in &mac_bytes {
70            use std::fmt::Write;
71            let _ = write!(sig, "{b:02x}");
72        }
73        self.signature = Some(sig);
74    }
75
76    pub fn verify_signature(&self, secret: &[u8]) -> bool {
77        let Some(ref sig) = self.signature else {
78            return false;
79        };
80
81        let expected: Vec<u8> = (0..sig.len())
82            .step_by(2)
83            .filter_map(|i| {
84                sig.get(i..i + 2)
85                    .and_then(|h| u8::from_str_radix(h, 16).ok())
86            })
87            .collect();
88        if expected.len() != sig.len() / 2 {
89            return false;
90        }
91
92        let computed = self.compute_hmac(secret);
93        constant_time_eq(&computed, &expected)
94    }
95
96    fn compute_hmac(&self, secret: &[u8]) -> Vec<u8> {
97        use hmac::{Hmac, Mac};
98        use sha2::Sha256;
99
100        let recipient_str = self.recipient.as_deref().unwrap_or("");
101        let mut sorted_meta: Vec<(&str, &str)> = self
102            .metadata
103            .iter()
104            .map(|(k, v)| (k.as_str(), v.as_str()))
105            .collect();
106        sorted_meta.sort_by_key(|(k, _)| *k);
107        let meta_str: String = sorted_meta
108            .iter()
109            .map(|(k, v)| format!("{k}={v}"))
110            .collect::<Vec<_>>()
111            .join(",");
112
113        let header = format!(
114            "v2:{}:{}:{}:{}:{}:{}:{}",
115            self.format_version,
116            self.sender.agent_id,
117            recipient_str,
118            self.content_type_str(),
119            self.sent_at.to_rfc3339(),
120            meta_str,
121            self.payload_json.len()
122        );
123        let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("HMAC accepts any key length");
124        mac.update(header.as_bytes());
125        mac.update(b"\0");
126        mac.update(self.payload_json.as_bytes());
127        mac.finalize().into_bytes().to_vec()
128    }
129
130    fn content_type_str(&self) -> &str {
131        match self.content_type {
132            TransportContentType::HandoffBundle => "handoff_bundle",
133            TransportContentType::ContextPackage => "context_package",
134            TransportContentType::A2AMessage => "a2a_message",
135            TransportContentType::A2ATask => "a2a_task",
136        }
137    }
138}
139
140pub fn serialize_envelope(envelope: &TransportEnvelopeV1) -> Result<String, String> {
141    let json = serde_json::to_string_pretty(envelope).map_err(|e| e.to_string())?;
142    if json.len() > MAX_ENVELOPE_BYTES {
143        return Err(format!(
144            "envelope too large ({} bytes, max {})",
145            json.len(),
146            MAX_ENVELOPE_BYTES
147        ));
148    }
149    Ok(json)
150}
151
152pub fn parse_envelope(json: &str) -> Result<TransportEnvelopeV1, String> {
153    if json.len() > MAX_ENVELOPE_BYTES {
154        return Err(format!(
155            "envelope too large ({} bytes, max {})",
156            json.len(),
157            MAX_ENVELOPE_BYTES
158        ));
159    }
160    let env: TransportEnvelopeV1 = serde_json::from_str(json).map_err(|e| e.to_string())?;
161    if env.format_version != 1 {
162        return Err(format!(
163            "unsupported format_version {} (expected 1)",
164            env.format_version
165        ));
166    }
167    Ok(env)
168}
169
170fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
171    if a.len() != b.len() {
172        return false;
173    }
174    let mut diff = 0u8;
175    for (x, y) in a.iter().zip(b.iter()) {
176        diff |= x ^ y;
177    }
178    diff == 0
179}
180
181fn compute_daemon_fingerprint() -> String {
182    use sha2::{Digest, Sha256};
183    let mut hasher = Sha256::new();
184    hasher.update(env!("CARGO_PKG_VERSION").as_bytes());
185    if let Ok(exe) = std::env::current_exe() {
186        hasher.update(exe.to_string_lossy().as_bytes());
187    }
188    format!("{:x}", hasher.finalize())[..16].to_string()
189}
190
191fn default_capabilities() -> Vec<String> {
192    vec![
193        "context_compression".to_string(),
194        "knowledge_graph".to_string(),
195        "shared_sessions".to_string(),
196        "a2a_messaging".to_string(),
197    ]
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    fn test_sender() -> AgentIdentityV1 {
205        AgentIdentityV1 {
206            agent_id: "test-agent".to_string(),
207            agent_type: "cursor".to_string(),
208            daemon_fingerprint: "abcd1234".to_string(),
209            capabilities: vec!["context_compression".to_string()],
210        }
211    }
212
213    #[test]
214    fn envelope_roundtrip() {
215        let env = TransportEnvelopeV1::new(
216            test_sender(),
217            Some("target-agent"),
218            TransportContentType::A2AMessage,
219            r#"{"hello":"world"}"#.to_string(),
220        );
221        let json = serialize_envelope(&env).unwrap();
222        let parsed = parse_envelope(&json).unwrap();
223        assert_eq!(parsed.format_version, 1);
224        assert_eq!(parsed.sender.agent_id, "test-agent");
225        assert_eq!(parsed.recipient, Some("target-agent".to_string()));
226        assert_eq!(parsed.content_type, TransportContentType::A2AMessage);
227    }
228
229    #[test]
230    fn hmac_sign_verify() {
231        let secret = b"test-secret-key";
232        let mut env = TransportEnvelopeV1::new(
233            test_sender(),
234            None,
235            TransportContentType::HandoffBundle,
236            "payload".to_string(),
237        );
238        assert!(!env.verify_signature(secret));
239
240        env.sign(secret);
241        assert!(env.signature.is_some());
242        assert!(env.verify_signature(secret));
243        assert!(!env.verify_signature(b"wrong-key"));
244    }
245
246    #[test]
247    fn rejects_oversized_envelope() {
248        let big = "x".repeat(MAX_ENVELOPE_BYTES + 1);
249        assert!(parse_envelope(&big).is_err());
250    }
251
252    #[test]
253    fn rejects_wrong_version() {
254        let json = r#"{"format_version":99,"sent_at":"2026-01-01T00:00:00Z","sender":{"agent_id":"a","agent_type":"b","daemon_fingerprint":"c","capabilities":[]},"recipient":null,"content_type":"a2a_message","payload_json":"{}","signature":null,"metadata":{}}"#;
255        assert!(parse_envelope(json).is_err());
256    }
257}