a3s_box_runtime/tee/
kbs.rs

1//! KBS (Key Broker Service) client for TEE secret provisioning.
2//!
3//! Implements the IETF RATS (Remote ATtestation procedureS) challenge-response
4//! protocol for fetching secrets from a Key Broker Service:
5//!
6//! 1. Client sends attestation evidence (SNP report) to KBS
7//! 2. KBS verifies the report against its policy
8//! 3. If verified, KBS returns the requested secret(s)
9//!
10//! Supports both single-key fetch and batch resource retrieval.
11
12use std::collections::HashMap;
13
14use a3s_box_core::error::{BoxError, Result};
15use serde::{Deserialize, Serialize};
16
17/// KBS endpoint configuration.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct KbsConfig {
20    /// KBS server URL (e.g., "https://kbs.example.com")
21    pub url: String,
22    /// Optional API key for authentication
23    #[serde(default)]
24    pub api_key: Option<String>,
25    /// Request timeout in seconds (default: 30)
26    #[serde(default = "default_timeout")]
27    pub timeout_secs: u64,
28    /// Whether to accept self-signed TLS certificates (for testing)
29    #[serde(default)]
30    pub insecure_tls: bool,
31}
32
33fn default_timeout() -> u64 {
34    30
35}
36
37impl Default for KbsConfig {
38    fn default() -> Self {
39        Self {
40            url: String::new(),
41            api_key: None,
42            timeout_secs: 30,
43            insecure_tls: false,
44        }
45    }
46}
47
48/// A request to the KBS for secret provisioning.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct KbsRequest {
51    /// Resource path (e.g., "default/keys/my-secret")
52    pub resource_path: String,
53    /// TEE evidence: base64-encoded SNP attestation report
54    pub evidence: String,
55    /// TEE type identifier
56    #[serde(default = "default_tee_type")]
57    pub tee_type: String,
58    /// Additional claims to include in the request
59    #[serde(default)]
60    pub extra_claims: HashMap<String, String>,
61}
62
63fn default_tee_type() -> String {
64    "snp".to_string()
65}
66
67/// Response from the KBS.
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct KbsResponse {
70    /// Whether the attestation was verified successfully
71    pub verified: bool,
72    /// The secret payload (base64-encoded), if verification passed
73    pub payload: Option<String>,
74    /// Error message if verification failed
75    pub error: Option<String>,
76    /// Token for subsequent requests (session-based KBS)
77    #[serde(default)]
78    pub token: Option<String>,
79}
80
81/// Result of a KBS secret fetch operation.
82#[derive(Debug, Clone)]
83pub struct KbsSecret {
84    /// Resource path that was fetched
85    pub resource_path: String,
86    /// Decoded secret bytes
87    pub secret: Vec<u8>,
88    /// Optional token for session continuity
89    pub token: Option<String>,
90}
91
92/// KBS client for fetching secrets from a Key Broker Service.
93pub struct KbsClient {
94    config: KbsConfig,
95}
96
97impl KbsClient {
98    /// Create a new KBS client with the given configuration.
99    pub fn new(config: KbsConfig) -> Self {
100        Self { config }
101    }
102
103    /// Build the attestation request payload.
104    pub fn build_request(&self, resource_path: &str, evidence: &[u8]) -> KbsRequest {
105        use base64::Engine;
106        KbsRequest {
107            resource_path: resource_path.to_string(),
108            evidence: base64::engine::general_purpose::STANDARD.encode(evidence),
109            tee_type: "snp".to_string(),
110            extra_claims: HashMap::new(),
111        }
112    }
113
114    /// Parse a KBS response and extract the secret.
115    pub fn parse_response(&self, resource_path: &str, response: &KbsResponse) -> Result<KbsSecret> {
116        if !response.verified {
117            return Err(BoxError::AttestationError(format!(
118                "KBS attestation verification failed: {}",
119                response.error.as_deref().unwrap_or("unknown error")
120            )));
121        }
122
123        let payload = response.payload.as_ref().ok_or_else(|| {
124            BoxError::AttestationError("KBS response verified but no payload returned".to_string())
125        })?;
126
127        use base64::Engine;
128        let secret = base64::engine::general_purpose::STANDARD
129            .decode(payload)
130            .map_err(|e| {
131                BoxError::AttestationError(format!("Failed to decode KBS payload: {}", e))
132            })?;
133
134        Ok(KbsSecret {
135            resource_path: resource_path.to_string(),
136            secret,
137            token: response.token.clone(),
138        })
139    }
140
141    /// Get the KBS endpoint URL for a resource path.
142    pub fn resource_url(&self, resource_path: &str) -> String {
143        let base = self.config.url.trim_end_matches('/');
144        format!("{}/kbs/v0/resource/{}", base, resource_path)
145    }
146
147    /// Get the KBS attestation endpoint URL.
148    pub fn attest_url(&self) -> String {
149        let base = self.config.url.trim_end_matches('/');
150        format!("{}/kbs/v0/attest", base)
151    }
152
153    /// Get the configuration.
154    pub fn config(&self) -> &KbsConfig {
155        &self.config
156    }
157}
158
159/// Parse a KBS resource path into (repository, type, tag).
160///
161/// Format: `repository/type/tag` (e.g., `default/keys/my-secret`)
162pub fn parse_resource_path(path: &str) -> Result<(&str, &str, &str)> {
163    let parts: Vec<&str> = path.splitn(3, '/').collect();
164    if parts.len() != 3 {
165        return Err(BoxError::AttestationError(format!(
166            "Invalid KBS resource path '{}': expected 'repository/type/tag'",
167            path
168        )));
169    }
170    Ok((parts[0], parts[1], parts[2]))
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn test_kbs_config_default() {
179        let config = KbsConfig::default();
180        assert!(config.url.is_empty());
181        assert!(config.api_key.is_none());
182        assert_eq!(config.timeout_secs, 30);
183        assert!(!config.insecure_tls);
184    }
185
186    #[test]
187    fn test_kbs_config_serde_roundtrip() {
188        let config = KbsConfig {
189            url: "https://kbs.example.com".to_string(),
190            api_key: Some("secret-key".to_string()),
191            timeout_secs: 60,
192            insecure_tls: true,
193        };
194        let json = serde_json::to_string(&config).unwrap();
195        let parsed: KbsConfig = serde_json::from_str(&json).unwrap();
196        assert_eq!(parsed.url, "https://kbs.example.com");
197        assert_eq!(parsed.api_key, Some("secret-key".to_string()));
198        assert_eq!(parsed.timeout_secs, 60);
199        assert!(parsed.insecure_tls);
200    }
201
202    #[test]
203    fn test_kbs_client_build_request() {
204        let config = KbsConfig {
205            url: "https://kbs.example.com".to_string(),
206            ..Default::default()
207        };
208        let client = KbsClient::new(config);
209        let evidence = b"fake-snp-report";
210        let request = client.build_request("default/keys/my-key", evidence);
211
212        assert_eq!(request.resource_path, "default/keys/my-key");
213        assert_eq!(request.tee_type, "snp");
214        assert!(!request.evidence.is_empty());
215
216        // Verify base64 roundtrip
217        use base64::Engine;
218        let decoded = base64::engine::general_purpose::STANDARD
219            .decode(&request.evidence)
220            .unwrap();
221        assert_eq!(decoded, evidence);
222    }
223
224    #[test]
225    fn test_kbs_client_parse_response_success() {
226        let config = KbsConfig::default();
227        let client = KbsClient::new(config);
228
229        use base64::Engine;
230        let secret_data = b"my-secret-value";
231        let payload = base64::engine::general_purpose::STANDARD.encode(secret_data);
232
233        let response = KbsResponse {
234            verified: true,
235            payload: Some(payload),
236            error: None,
237            token: Some("session-token".to_string()),
238        };
239
240        let secret = client
241            .parse_response("default/keys/test", &response)
242            .unwrap();
243        assert_eq!(secret.resource_path, "default/keys/test");
244        assert_eq!(secret.secret, secret_data);
245        assert_eq!(secret.token, Some("session-token".to_string()));
246    }
247
248    #[test]
249    fn test_kbs_client_parse_response_verification_failed() {
250        let config = KbsConfig::default();
251        let client = KbsClient::new(config);
252
253        let response = KbsResponse {
254            verified: false,
255            payload: None,
256            error: Some("measurement mismatch".to_string()),
257            token: None,
258        };
259
260        let result = client.parse_response("default/keys/test", &response);
261        assert!(result.is_err());
262        assert!(result
263            .unwrap_err()
264            .to_string()
265            .contains("measurement mismatch"));
266    }
267
268    #[test]
269    fn test_kbs_client_parse_response_no_payload() {
270        let config = KbsConfig::default();
271        let client = KbsClient::new(config);
272
273        let response = KbsResponse {
274            verified: true,
275            payload: None,
276            error: None,
277            token: None,
278        };
279
280        let result = client.parse_response("default/keys/test", &response);
281        assert!(result.is_err());
282        assert!(result.unwrap_err().to_string().contains("no payload"));
283    }
284
285    #[test]
286    fn test_kbs_client_parse_response_invalid_base64() {
287        let config = KbsConfig::default();
288        let client = KbsClient::new(config);
289
290        let response = KbsResponse {
291            verified: true,
292            payload: Some("not-valid-base64!!!".to_string()),
293            error: None,
294            token: None,
295        };
296
297        let result = client.parse_response("default/keys/test", &response);
298        assert!(result.is_err());
299        assert!(result.unwrap_err().to_string().contains("decode"));
300    }
301
302    #[test]
303    fn test_kbs_client_resource_url() {
304        let config = KbsConfig {
305            url: "https://kbs.example.com".to_string(),
306            ..Default::default()
307        };
308        let client = KbsClient::new(config);
309        assert_eq!(
310            client.resource_url("default/keys/my-key"),
311            "https://kbs.example.com/kbs/v0/resource/default/keys/my-key"
312        );
313    }
314
315    #[test]
316    fn test_kbs_client_resource_url_trailing_slash() {
317        let config = KbsConfig {
318            url: "https://kbs.example.com/".to_string(),
319            ..Default::default()
320        };
321        let client = KbsClient::new(config);
322        assert_eq!(
323            client.resource_url("default/keys/my-key"),
324            "https://kbs.example.com/kbs/v0/resource/default/keys/my-key"
325        );
326    }
327
328    #[test]
329    fn test_kbs_client_attest_url() {
330        let config = KbsConfig {
331            url: "https://kbs.example.com".to_string(),
332            ..Default::default()
333        };
334        let client = KbsClient::new(config);
335        assert_eq!(client.attest_url(), "https://kbs.example.com/kbs/v0/attest");
336    }
337
338    #[test]
339    fn test_parse_resource_path_valid() {
340        let (repo, rtype, tag) = parse_resource_path("default/keys/my-secret").unwrap();
341        assert_eq!(repo, "default");
342        assert_eq!(rtype, "keys");
343        assert_eq!(tag, "my-secret");
344    }
345
346    #[test]
347    fn test_parse_resource_path_with_nested_tag() {
348        let (repo, rtype, tag) = parse_resource_path("myrepo/certs/tls/server.pem").unwrap();
349        assert_eq!(repo, "myrepo");
350        assert_eq!(rtype, "certs");
351        assert_eq!(tag, "tls/server.pem");
352    }
353
354    #[test]
355    fn test_parse_resource_path_invalid_too_few() {
356        assert!(parse_resource_path("default/keys").is_err());
357        assert!(parse_resource_path("just-one").is_err());
358        assert!(parse_resource_path("").is_err());
359    }
360
361    #[test]
362    fn test_kbs_request_serde_roundtrip() {
363        let request = KbsRequest {
364            resource_path: "default/keys/test".to_string(),
365            evidence: "base64evidence".to_string(),
366            tee_type: "snp".to_string(),
367            extra_claims: HashMap::from([("nonce".to_string(), "abc123".to_string())]),
368        };
369        let json = serde_json::to_string(&request).unwrap();
370        let parsed: KbsRequest = serde_json::from_str(&json).unwrap();
371        assert_eq!(parsed.resource_path, "default/keys/test");
372        assert_eq!(parsed.evidence, "base64evidence");
373        assert_eq!(parsed.tee_type, "snp");
374        assert_eq!(parsed.extra_claims.get("nonce").unwrap(), "abc123");
375    }
376
377    #[test]
378    fn test_kbs_response_serde_roundtrip() {
379        let response = KbsResponse {
380            verified: true,
381            payload: Some("c2VjcmV0".to_string()),
382            error: None,
383            token: Some("tok-123".to_string()),
384        };
385        let json = serde_json::to_string(&response).unwrap();
386        let parsed: KbsResponse = serde_json::from_str(&json).unwrap();
387        assert!(parsed.verified);
388        assert_eq!(parsed.payload, Some("c2VjcmV0".to_string()));
389        assert!(parsed.error.is_none());
390        assert_eq!(parsed.token, Some("tok-123".to_string()));
391    }
392}
a3s_box_runtime/tee/kbs.rs

a3s_box_runtime/tee/
kbs.rs
