Skip to main content

a3s_box_runtime/tee/
sealed.rs

1//! Sealed storage for TEE-bound encryption.
2//!
3//! Provides encryption/decryption of data bound to the TEE's identity
4//! (measurement + chip_id). Only the same TEE running the same firmware
5//! and guest image on the same physical chip can unseal the data.
6//!
7//! ## Key Derivation
8//!
9//! The sealing key is derived using HKDF-SHA256:
10//! - IKM (Input Key Material): `measurement || chip_id` from the SNP report
11//! - Salt: "a3s-sealed-storage-v1"
12//! - Info: caller-provided context (e.g., "session-keys", "model-weights")
13//!
14//! ## Encryption
15//!
16//! AES-256-GCM with a random 96-bit nonce per seal operation.
17//! The sealed blob format: `nonce (12 bytes) || ciphertext+tag`
18//!
19//! ## Sealing Policies
20//!
21//! - `MeasurementAndChip`: Binds to both measurement and chip (strictest)
22//! - `MeasurementOnly`: Binds to measurement only (portable across chips)
23//! - `ChipOnly`: Binds to chip only (survives firmware updates)
24
25use a3s_box_core::error::{BoxError, Result};
26use ring::aead::{self, Aad, BoundKey, Nonce, NonceSequence, NONCE_LEN};
27use ring::hkdf;
28use serde::{Deserialize, Serialize};
29
30/// Salt for HKDF key derivation.
31const HKDF_SALT: &[u8] = b"a3s-sealed-storage-v1";
32
33/// Sealing policy determines what TEE identity fields bind the key.
34#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
35pub enum SealingPolicy {
36    /// Bind to both measurement and chip_id (strictest).
37    /// Data can only be unsealed by the exact same guest image
38    /// on the exact same physical chip.
39    #[default]
40    MeasurementAndChip,
41
42    /// Bind to measurement only (portable across chips).
43    /// Data can be unsealed by the same guest image on any chip.
44    MeasurementOnly,
45
46    /// Bind to chip only (survives firmware/image updates).
47    /// Data can be unsealed by any guest image on the same chip.
48    ChipOnly,
49}
50
51/// Sealed data blob with metadata for unsealing.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct SealedData {
54    /// The sealing policy used.
55    pub policy: SealingPolicy,
56    /// Context string used for key derivation.
57    pub context: String,
58    /// Sealed blob: nonce (12 bytes) || ciphertext || tag (16 bytes).
59    #[serde(with = "base64_serde")]
60    pub blob: Vec<u8>,
61}
62
63// ============================================================================
64// Seal / Unseal operations
65// ============================================================================
66
67/// Seal (encrypt) data bound to the TEE identity.
68///
69/// # Arguments
70/// * `report` - Raw SNP report bytes (1184 bytes) containing measurement and chip_id
71/// * `plaintext` - Data to encrypt
72/// * `context` - Application-specific context for key derivation (e.g., "session-keys")
73/// * `policy` - Sealing policy determining which TEE fields bind the key
74///
75/// # Returns
76/// A `SealedData` blob that can only be unsealed with the same TEE identity.
77pub fn seal(
78    report: &[u8],
79    plaintext: &[u8],
80    context: &str,
81    policy: SealingPolicy,
82) -> Result<SealedData> {
83    let key = derive_sealing_key(report, context, policy)?;
84
85    // Generate random nonce
86    let rng = ring::rand::SystemRandom::new();
87    let mut nonce_bytes = [0u8; NONCE_LEN];
88    ring::rand::SecureRandom::fill(&rng, &mut nonce_bytes)
89        .map_err(|_| BoxError::AttestationError("Failed to generate random nonce".to_string()))?;
90
91    // Encrypt with AES-256-GCM
92    let mut in_out = plaintext.to_vec();
93    let unbound_key = aead::UnboundKey::new(&aead::AES_256_GCM, &key)
94        .map_err(|_| BoxError::AttestationError("Failed to create AES-256-GCM key".to_string()))?;
95
96    let nonce_seq = SingleNonce::new(nonce_bytes);
97    let mut sealing_key = aead::SealingKey::new(unbound_key, nonce_seq);
98
99    sealing_key
100        .seal_in_place_append_tag(Aad::from(context.as_bytes()), &mut in_out)
101        .map_err(|_| BoxError::AttestationError("AES-256-GCM seal failed".to_string()))?;
102
103    // Prepend nonce to ciphertext
104    let mut blob = Vec::with_capacity(NONCE_LEN + in_out.len());
105    blob.extend_from_slice(&nonce_bytes);
106    blob.extend_from_slice(&in_out);
107
108    Ok(SealedData {
109        policy,
110        context: context.to_string(),
111        blob,
112    })
113}
114
115/// Unseal (decrypt) data using the TEE identity.
116///
117/// # Arguments
118/// * `report` - Raw SNP report bytes (must match the TEE that sealed the data)
119/// * `sealed` - The sealed data blob
120///
121/// # Returns
122/// The original plaintext, or an error if the TEE identity doesn't match.
123pub fn unseal(report: &[u8], sealed: &SealedData) -> Result<Vec<u8>> {
124    if sealed.blob.len() < NONCE_LEN + aead::AES_256_GCM.tag_len() {
125        return Err(BoxError::AttestationError(
126            "Sealed blob too short".to_string(),
127        ));
128    }
129
130    let key = derive_sealing_key(report, &sealed.context, sealed.policy)?;
131
132    // Split nonce and ciphertext
133    let nonce_bytes: [u8; NONCE_LEN] = sealed.blob[..NONCE_LEN]
134        .try_into()
135        .map_err(|_| BoxError::AttestationError("Invalid nonce in sealed blob".to_string()))?;
136
137    let mut in_out = sealed.blob[NONCE_LEN..].to_vec();
138
139    let unbound_key = aead::UnboundKey::new(&aead::AES_256_GCM, &key)
140        .map_err(|_| BoxError::AttestationError("Failed to create AES-256-GCM key".to_string()))?;
141
142    let nonce_seq = SingleNonce::new(nonce_bytes);
143    let mut opening_key = aead::OpeningKey::new(unbound_key, nonce_seq);
144
145    let plaintext = opening_key
146        .open_in_place(Aad::from(sealed.context.as_bytes()), &mut in_out)
147        .map_err(|_| {
148            BoxError::AttestationError(
149                "Unseal failed: TEE identity mismatch or data corrupted".to_string(),
150            )
151        })?;
152
153    Ok(plaintext.to_vec())
154}
155
156// ============================================================================
157// Key derivation
158// ============================================================================
159
160/// Derive a 256-bit sealing key from the SNP report using HKDF-SHA256.
161fn derive_sealing_key(report: &[u8], context: &str, policy: SealingPolicy) -> Result<[u8; 32]> {
162    // Extract measurement (0x90, 48 bytes) and chip_id (0x1A0, 64 bytes)
163    if report.len() < 0x1E0 {
164        return Err(BoxError::AttestationError(
165            "Report too short to extract sealing identity".to_string(),
166        ));
167    }
168
169    let measurement = &report[0x90..0xC0]; // 48 bytes
170    let chip_id = &report[0x1A0..0x1E0]; // 64 bytes
171
172    // Build IKM based on policy
173    let ikm = match policy {
174        SealingPolicy::MeasurementAndChip => {
175            let mut v = Vec::with_capacity(112);
176            v.extend_from_slice(measurement);
177            v.extend_from_slice(chip_id);
178            v
179        }
180        SealingPolicy::MeasurementOnly => measurement.to_vec(),
181        SealingPolicy::ChipOnly => chip_id.to_vec(),
182    };
183
184    // HKDF extract + expand
185    let salt = hkdf::Salt::new(hkdf::HKDF_SHA256, HKDF_SALT);
186    let prk = salt.extract(&ikm);
187    let info = [context.as_bytes()];
188    let okm = prk
189        .expand(&info, HkdfLen(32))
190        .map_err(|_| BoxError::AttestationError("HKDF expand failed".to_string()))?;
191
192    let mut key = [0u8; 32];
193    okm.fill(&mut key)
194        .map_err(|_| BoxError::AttestationError("HKDF fill failed".to_string()))?;
195
196    Ok(key)
197}
198
199// ============================================================================
200// ring helper types
201// ============================================================================
202
203/// A NonceSequence that yields a single nonce then fails.
204struct SingleNonce {
205    nonce: Option<[u8; NONCE_LEN]>,
206}
207
208impl SingleNonce {
209    fn new(nonce: [u8; NONCE_LEN]) -> Self {
210        Self { nonce: Some(nonce) }
211    }
212}
213
214impl NonceSequence for SingleNonce {
215    fn advance(&mut self) -> std::result::Result<Nonce, ring::error::Unspecified> {
216        self.nonce
217            .take()
218            .map(Nonce::assume_unique_for_key)
219            .ok_or(ring::error::Unspecified)
220    }
221}
222
223/// HKDF output length wrapper for ring.
224struct HkdfLen(usize);
225
226impl hkdf::KeyType for HkdfLen {
227    fn len(&self) -> usize {
228        self.0
229    }
230}
231
232// ============================================================================
233// Base64 serde helper
234// ============================================================================
235
236mod base64_serde {
237    use serde::{Deserialize, Deserializer, Serializer};
238
239    pub fn serialize<S: Serializer>(bytes: &Vec<u8>, s: S) -> std::result::Result<S::Ok, S::Error> {
240        use base64::Engine;
241        s.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
242    }
243
244    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> std::result::Result<Vec<u8>, D::Error> {
245        use base64::Engine;
246        let s = String::deserialize(d)?;
247        base64::engine::general_purpose::STANDARD
248            .decode(&s)
249            .map_err(serde::de::Error::custom)
250    }
251}
252
253// ============================================================================
254// Tests
255// ============================================================================
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260
261    /// Build a fake 1184-byte report with known measurement and chip_id.
262    fn make_test_report() -> Vec<u8> {
263        let mut report = vec![0u8; 1184];
264        // measurement at 0x90 (48 bytes)
265        for i in 0..48 {
266            report[0x90 + i] = (i as u8).wrapping_mul(0xA3);
267        }
268        // chip_id at 0x1A0 (64 bytes)
269        for b in &mut report[0x1A0..0x1E0] {
270            *b = 0xA3;
271        }
272        report
273    }
274
275    #[test]
276    fn test_seal_unseal_roundtrip() {
277        let report = make_test_report();
278        let plaintext = b"secret data for TEE";
279        let sealed = seal(&report, plaintext, "test-context", SealingPolicy::default()).unwrap();
280        let unsealed = unseal(&report, &sealed).unwrap();
281        assert_eq!(unsealed, plaintext);
282    }
283
284    #[test]
285    fn test_seal_unseal_measurement_only() {
286        let report = make_test_report();
287        let plaintext = b"measurement-bound secret";
288        let sealed = seal(&report, plaintext, "ctx", SealingPolicy::MeasurementOnly).unwrap();
289        let unsealed = unseal(&report, &sealed).unwrap();
290        assert_eq!(unsealed, plaintext);
291    }
292
293    #[test]
294    fn test_seal_unseal_chip_only() {
295        let report = make_test_report();
296        let plaintext = b"chip-bound secret";
297        let sealed = seal(&report, plaintext, "ctx", SealingPolicy::ChipOnly).unwrap();
298        let unsealed = unseal(&report, &sealed).unwrap();
299        assert_eq!(unsealed, plaintext);
300    }
301
302    #[test]
303    fn test_unseal_wrong_measurement_fails() {
304        let report = make_test_report();
305        let plaintext = b"secret";
306        let sealed = seal(&report, plaintext, "ctx", SealingPolicy::MeasurementOnly).unwrap();
307
308        // Different measurement
309        let mut wrong_report = report.clone();
310        wrong_report[0x90] = 0xFF;
311        let result = unseal(&wrong_report, &sealed);
312        assert!(result.is_err());
313    }
314
315    #[test]
316    fn test_unseal_wrong_chip_fails() {
317        let report = make_test_report();
318        let plaintext = b"secret";
319        let sealed = seal(&report, plaintext, "ctx", SealingPolicy::ChipOnly).unwrap();
320
321        // Different chip_id
322        let mut wrong_report = report.clone();
323        wrong_report[0x1A0] = 0xFF;
324        let result = unseal(&wrong_report, &sealed);
325        assert!(result.is_err());
326    }
327
328    #[test]
329    fn test_unseal_wrong_context_fails() {
330        let report = make_test_report();
331        let plaintext = b"secret";
332        let sealed = seal(&report, plaintext, "context-a", SealingPolicy::default()).unwrap();
333
334        // Try to unseal with different context
335        let mut tampered = sealed.clone();
336        tampered.context = "context-b".to_string();
337        let result = unseal(&report, &tampered);
338        assert!(result.is_err());
339    }
340
341    #[test]
342    fn test_unseal_tampered_blob_fails() {
343        let report = make_test_report();
344        let plaintext = b"secret";
345        let sealed = seal(&report, plaintext, "ctx", SealingPolicy::default()).unwrap();
346
347        // Tamper with ciphertext
348        let mut tampered = sealed.clone();
349        if let Some(byte) = tampered.blob.get_mut(NONCE_LEN + 1) {
350            *byte ^= 0xFF;
351        }
352        let result = unseal(&report, &tampered);
353        assert!(result.is_err());
354    }
355
356    #[test]
357    fn test_seal_empty_plaintext() {
358        let report = make_test_report();
359        let sealed = seal(&report, b"", "ctx", SealingPolicy::default()).unwrap();
360        let unsealed = unseal(&report, &sealed).unwrap();
361        assert!(unsealed.is_empty());
362    }
363
364    #[test]
365    fn test_seal_large_plaintext() {
366        let report = make_test_report();
367        let plaintext = vec![0xAB; 1024 * 1024]; // 1 MiB
368        let sealed = seal(&report, &plaintext, "ctx", SealingPolicy::default()).unwrap();
369        let unsealed = unseal(&report, &sealed).unwrap();
370        assert_eq!(unsealed, plaintext);
371    }
372
373    #[test]
374    fn test_sealed_blob_size() {
375        let report = make_test_report();
376        let plaintext = b"hello";
377        let sealed = seal(&report, plaintext, "ctx", SealingPolicy::default()).unwrap();
378        // blob = nonce (12) + ciphertext (5) + tag (16) = 33
379        assert_eq!(
380            sealed.blob.len(),
381            NONCE_LEN + plaintext.len() + aead::AES_256_GCM.tag_len()
382        );
383    }
384
385    #[test]
386    fn test_report_too_short() {
387        let short_report = vec![0u8; 100];
388        let result = seal(&short_report, b"data", "ctx", SealingPolicy::default());
389        assert!(result.is_err());
390    }
391
392    #[test]
393    fn test_sealed_data_serialization() {
394        let report = make_test_report();
395        let sealed = seal(&report, b"secret", "ctx", SealingPolicy::default()).unwrap();
396        let json = serde_json::to_string(&sealed).unwrap();
397        let deserialized: SealedData = serde_json::from_str(&json).unwrap();
398        let unsealed = unseal(&report, &deserialized).unwrap();
399        assert_eq!(unsealed, b"secret");
400    }
401
402    #[test]
403    fn test_sealing_policy_default() {
404        assert_eq!(SealingPolicy::default(), SealingPolicy::MeasurementAndChip);
405    }
406
407    #[test]
408    fn test_different_nonces_per_seal() {
409        let report = make_test_report();
410        let s1 = seal(&report, b"same", "ctx", SealingPolicy::default()).unwrap();
411        let s2 = seal(&report, b"same", "ctx", SealingPolicy::default()).unwrap();
412        // Different nonces → different blobs
413        assert_ne!(s1.blob, s2.blob);
414        // But both unseal to the same plaintext
415        assert_eq!(unseal(&report, &s1).unwrap(), b"same");
416        assert_eq!(unseal(&report, &s2).unwrap(), b"same");
417    }
418
419    #[test]
420    fn test_chip_only_survives_measurement_change() {
421        let report = make_test_report();
422        let sealed = seal(&report, b"secret", "ctx", SealingPolicy::ChipOnly).unwrap();
423
424        // Change measurement but keep chip_id
425        let mut updated_report = report.clone();
426        updated_report[0x90] = 0xFF;
427        let unsealed = unseal(&updated_report, &sealed).unwrap();
428        assert_eq!(unsealed, b"secret");
429    }
430
431    #[test]
432    fn test_measurement_only_survives_chip_change() {
433        let report = make_test_report();
434        let sealed = seal(&report, b"secret", "ctx", SealingPolicy::MeasurementOnly).unwrap();
435
436        // Change chip_id but keep measurement
437        let mut other_chip = report.clone();
438        other_chip[0x1A0] = 0xFF;
439        let unsealed = unseal(&other_chip, &sealed).unwrap();
440        assert_eq!(unsealed, b"secret");
441    }
442}