1use sha2::{Digest, Sha256};
40use std::time::{SystemTime, UNIX_EPOCH};
41
42#[derive(Debug, Clone, Copy, PartialEq)]
44pub enum TeePlatform {
45 IntelSgx,
47 AmdSev,
49 ArmTrustZone,
51 Simulated,
53}
54
55#[derive(Debug, Clone, Copy, PartialEq)]
57pub enum EnclaveStatus {
58 Uninitialized,
60 Ready,
62 Secure,
64 Compromised,
66 Destroyed,
68}
69
70#[derive(Debug, Clone)]
72pub struct AttestationReport {
73 pub platform: TeePlatform,
75 pub measurement: [u8; 32],
77 pub signer: [u8; 32],
79 pub product_id: u16,
81 pub svn: u16,
83 pub report_data: [u8; 64],
85 pub timestamp: u64,
87 pub signature: [u8; 64],
89 pub is_valid: bool,
91}
92
93#[derive(Debug, Clone)]
95pub struct SealedData {
96 pub ciphertext: Vec<u8>,
98 pub aad: Vec<u8>,
100 pub nonce: [u8; 12],
102 pub tag: [u8; 16],
104 pub policy: SealingPolicy,
106 pub key_id: [u8; 32],
108}
109
110#[derive(Debug, Clone, Copy, PartialEq)]
112pub enum SealingPolicy {
113 EnclaveIdentity,
115 SignerIdentity,
117}
118
119#[derive(Debug, Clone)]
121pub struct EnclaveConfig {
122 pub platform: TeePlatform,
124 pub expected_measurement: Option<[u8; 32]>,
126 pub min_svn: u16,
128 pub allow_debug: bool,
130 pub heap_size: usize,
132 pub stack_size: usize,
134 pub num_tcs: u32,
136}
137
138impl Default for EnclaveConfig {
139 fn default() -> Self {
140 Self {
141 platform: TeePlatform::Simulated,
142 expected_measurement: None,
143 min_svn: 1,
144 allow_debug: false,
145 heap_size: 64 * 1024 * 1024, stack_size: 1024 * 1024, num_tcs: 4,
148 }
149 }
150}
151
152#[derive(Debug)]
154pub struct SecureEnclave {
155 config: EnclaveConfig,
157 status: EnclaveStatus,
159 enclave_id: u64,
161 sealed_master_key: Option<SealedData>,
163 attestation_cache: Option<AttestationReport>,
165 stats: EnclaveStats,
167}
168
169#[derive(Debug, Default)]
171pub struct EnclaveStats {
172 pub ecalls: u64,
174 pub ocalls: u64,
176 pub attestations: u64,
178 pub seals: u64,
180 pub unseals: u64,
182 pub failures: u64,
184}
185
186impl Default for SecureEnclave {
187 fn default() -> Self {
188 Self::new(EnclaveConfig::default())
189 }
190}
191
192impl SecureEnclave {
193 pub fn new(config: EnclaveConfig) -> Self {
195 Self {
196 config,
197 status: EnclaveStatus::Uninitialized,
198 enclave_id: 0,
199 sealed_master_key: None,
200 attestation_cache: None,
201 stats: EnclaveStats::default(),
202 }
203 }
204
205 pub fn with_sgx() -> Self {
207 Self::new(EnclaveConfig {
208 platform: TeePlatform::IntelSgx,
209 ..Default::default()
210 })
211 }
212
213 pub fn with_sev() -> Self {
215 Self::new(EnclaveConfig {
216 platform: TeePlatform::AmdSev,
217 ..Default::default()
218 })
219 }
220
221 pub fn initialize(&mut self) -> Result<(), EnclaveError> {
223 if self.status != EnclaveStatus::Uninitialized {
224 return Err(EnclaveError::AlreadyInitialized);
225 }
226
227 match self.config.platform {
229 TeePlatform::IntelSgx => {
230 self.enclave_id = self.generate_enclave_id();
232 }
233 TeePlatform::AmdSev => {
234 self.enclave_id = self.generate_enclave_id();
236 }
237 TeePlatform::ArmTrustZone => {
238 self.enclave_id = self.generate_enclave_id();
240 }
241 TeePlatform::Simulated => {
242 self.enclave_id = self.generate_enclave_id();
244 }
245 }
246
247 self.status = EnclaveStatus::Ready;
248 Ok(())
249 }
250
251 fn generate_enclave_id(&self) -> u64 {
253 let timestamp = SystemTime::now()
254 .duration_since(UNIX_EPOCH)
255 .unwrap_or_default()
256 .as_nanos() as u64;
257
258 let mut hasher = Sha256::new();
259 hasher.update(timestamp.to_le_bytes());
260 hasher.update(format!("{:?}", self.config.platform).as_bytes());
261 let hash = hasher.finalize();
262
263 u64::from_le_bytes(hash[0..8].try_into().unwrap())
264 }
265
266 pub fn enter_secure_mode(&mut self) -> Result<(), EnclaveError> {
268 if self.status != EnclaveStatus::Ready {
269 return Err(EnclaveError::NotReady);
270 }
271
272 self.status = EnclaveStatus::Secure;
273 self.stats.ecalls += 1;
274 Ok(())
275 }
276
277 pub fn exit_secure_mode(&mut self) -> Result<(), EnclaveError> {
279 if self.status != EnclaveStatus::Secure {
280 return Err(EnclaveError::NotInSecureMode);
281 }
282
283 self.status = EnclaveStatus::Ready;
284 self.stats.ocalls += 1;
285 Ok(())
286 }
287
288 pub fn attest(&mut self, challenge: &[u8; 32]) -> Result<AttestationReport, EnclaveError> {
290 if self.status == EnclaveStatus::Uninitialized {
291 return Err(EnclaveError::NotReady);
292 }
293
294 self.stats.attestations += 1;
295
296 let mut hasher = Sha256::new();
298 hasher.update(b"HOPE_GENOME_ENCLAVE_V1");
299 hasher.update(format!("{:?}", self.config.platform).as_bytes());
300 let measurement_hash = hasher.finalize();
301 let mut measurement = [0u8; 32];
302 measurement.copy_from_slice(&measurement_hash);
303
304 let mut hasher = Sha256::new();
306 hasher.update(b"HOPE_GENOME_SIGNER_V1");
307 let signer_hash = hasher.finalize();
308 let mut signer = [0u8; 32];
309 signer.copy_from_slice(&signer_hash);
310
311 let mut report_data = [0u8; 64];
313 report_data[0..32].copy_from_slice(challenge);
314
315 let mut hasher = Sha256::new();
317 hasher.update(measurement);
318 hasher.update(signer);
319 hasher.update(report_data);
320 let sig_hash = hasher.finalize();
321 let mut signature = [0u8; 64];
322 signature[0..32].copy_from_slice(&sig_hash);
323
324 let report = AttestationReport {
325 platform: self.config.platform,
326 measurement,
327 signer,
328 product_id: 1,
329 svn: self.config.min_svn,
330 report_data,
331 timestamp: SystemTime::now()
332 .duration_since(UNIX_EPOCH)
333 .unwrap_or_default()
334 .as_secs(),
335 signature,
336 is_valid: true,
337 };
338
339 self.attestation_cache = Some(report.clone());
340 Ok(report)
341 }
342
343 pub fn verify_attestation(&self, report: &AttestationReport) -> Result<bool, EnclaveError> {
345 if report.platform != self.config.platform {
347 return Ok(false);
348 }
349
350 if report.svn < self.config.min_svn {
352 return Ok(false);
353 }
354
355 if let Some(expected) = self.config.expected_measurement {
357 if report.measurement != expected {
358 return Ok(false);
359 }
360 }
361
362 let mut hasher = Sha256::new();
364 hasher.update(report.measurement);
365 hasher.update(report.signer);
366 hasher.update(report.report_data);
367 let expected_sig = hasher.finalize();
368
369 if report.signature[0..32] != expected_sig[..] {
370 return Ok(false);
371 }
372
373 Ok(report.is_valid)
374 }
375
376 pub fn seal(&mut self, data: &[u8], policy: SealingPolicy) -> Result<SealedData, EnclaveError> {
378 if self.status == EnclaveStatus::Uninitialized {
379 return Err(EnclaveError::NotReady);
380 }
381
382 self.stats.seals += 1;
383
384 let mut hasher = Sha256::new();
386 match policy {
387 SealingPolicy::EnclaveIdentity => hasher.update(b"MRENCLAVE"),
388 SealingPolicy::SignerIdentity => hasher.update(b"MRSIGNER"),
389 };
390 hasher.update(self.enclave_id.to_le_bytes());
391 let key_hash = hasher.finalize();
392 let mut key_id = [0u8; 32];
393 key_id.copy_from_slice(&key_hash);
394
395 let timestamp = SystemTime::now()
397 .duration_since(UNIX_EPOCH)
398 .unwrap_or_default()
399 .as_nanos() as u64;
400 let mut nonce = [0u8; 12];
401 nonce[0..8].copy_from_slice(×tamp.to_le_bytes());
402
403 let mut ciphertext = data.to_vec();
405 for (i, byte) in ciphertext.iter_mut().enumerate() {
406 *byte ^= key_id[i % 32];
407 }
408
409 let mut hasher = Sha256::new();
411 hasher.update(key_id);
412 hasher.update(nonce);
413 hasher.update(&ciphertext);
414 let tag_hash = hasher.finalize();
415 let mut tag = [0u8; 16];
416 tag.copy_from_slice(&tag_hash[0..16]);
417
418 Ok(SealedData {
419 ciphertext,
420 aad: Vec::new(),
421 nonce,
422 tag,
423 policy,
424 key_id,
425 })
426 }
427
428 pub fn unseal(&mut self, sealed: &SealedData) -> Result<Vec<u8>, EnclaveError> {
430 if self.status == EnclaveStatus::Uninitialized {
431 return Err(EnclaveError::NotReady);
432 }
433
434 self.stats.unseals += 1;
435
436 let mut hasher = Sha256::new();
438 hasher.update(sealed.key_id);
439 hasher.update(sealed.nonce);
440 hasher.update(&sealed.ciphertext);
441 let expected_tag = hasher.finalize();
442
443 if sealed.tag != expected_tag[0..16] {
444 self.stats.failures += 1;
445 return Err(EnclaveError::AuthenticationFailed);
446 }
447
448 let mut plaintext = sealed.ciphertext.clone();
450 for (i, byte) in plaintext.iter_mut().enumerate() {
451 *byte ^= sealed.key_id[i % 32];
452 }
453
454 Ok(plaintext)
455 }
456
457 pub fn execute_secure<F, T>(&mut self, operation: F) -> Result<T, EnclaveError>
459 where
460 F: FnOnce() -> T,
461 {
462 self.enter_secure_mode()?;
463 let result = operation();
464 self.exit_secure_mode()?;
465 Ok(result)
466 }
467
468 pub fn destroy(&mut self) -> Result<(), EnclaveError> {
470 self.sealed_master_key = None;
472 self.attestation_cache = None;
473 self.enclave_id = 0;
474
475 self.status = EnclaveStatus::Destroyed;
476 Ok(())
477 }
478
479 pub fn status(&self) -> EnclaveStatus {
481 self.status
482 }
483
484 pub fn stats(&self) -> &EnclaveStats {
486 &self.stats
487 }
488
489 pub fn enclave_id(&self) -> u64 {
491 self.enclave_id
492 }
493
494 pub fn platform(&self) -> TeePlatform {
496 self.config.platform
497 }
498}
499
500#[derive(Debug, Clone, PartialEq)]
502pub enum EnclaveError {
503 AlreadyInitialized,
505 NotReady,
507 NotInSecureMode,
509 AttestationFailed,
511 AuthenticationFailed,
513 PlatformNotSupported,
515 HardwareError(String),
517 MemoryError,
519}
520
521impl std::fmt::Display for EnclaveError {
522 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
523 match self {
524 EnclaveError::AlreadyInitialized => write!(f, "Enclave already initialized"),
525 EnclaveError::NotReady => write!(f, "Enclave not ready"),
526 EnclaveError::NotInSecureMode => write!(f, "Not in secure mode"),
527 EnclaveError::AttestationFailed => write!(f, "Attestation failed"),
528 EnclaveError::AuthenticationFailed => write!(f, "Authentication failed"),
529 EnclaveError::PlatformNotSupported => write!(f, "Platform not supported"),
530 EnclaveError::HardwareError(msg) => write!(f, "Hardware error: {}", msg),
531 EnclaveError::MemoryError => write!(f, "Memory allocation error"),
532 }
533 }
534}
535
536impl std::error::Error for EnclaveError {}
537
538#[derive(Debug)]
540pub struct ProtectedWatchdog {
541 enclave: SecureEnclave,
543 violation_count: u32,
545 max_violations: u32,
547 locked: bool,
549}
550
551impl ProtectedWatchdog {
552 pub fn new(platform: TeePlatform) -> Self {
554 let config = EnclaveConfig {
555 platform,
556 allow_debug: false,
557 ..Default::default()
558 };
559
560 Self {
561 enclave: SecureEnclave::new(config),
562 violation_count: 0,
563 max_violations: 10,
564 locked: false,
565 }
566 }
567
568 pub fn initialize(&mut self) -> Result<(), EnclaveError> {
570 self.enclave.initialize()
571 }
572
573 pub fn record_violation(&mut self) -> Result<u32, EnclaveError> {
575 self.enclave.execute_secure(|| {
576 self.violation_count += 1;
577 if self.violation_count >= self.max_violations {
578 self.locked = true;
579 }
580 self.violation_count
581 })
582 }
583
584 pub fn check_action(&mut self, _action: &str) -> Result<bool, EnclaveError> {
586 self.enclave.execute_secure(|| !self.locked)
587 }
588
589 pub fn get_attestation(
591 &mut self,
592 challenge: &[u8; 32],
593 ) -> Result<AttestationReport, EnclaveError> {
594 self.enclave.attest(challenge)
595 }
596
597 pub fn violation_count(&self) -> u32 {
599 self.violation_count
600 }
601
602 pub fn is_locked(&self) -> bool {
604 self.locked
605 }
606}
607
608#[cfg(test)]
609mod tests {
610 use super::*;
611
612 #[test]
613 fn test_enclave_lifecycle() {
614 let mut enclave = SecureEnclave::new(EnclaveConfig::default());
615 assert_eq!(enclave.status(), EnclaveStatus::Uninitialized);
616
617 enclave.initialize().unwrap();
618 assert_eq!(enclave.status(), EnclaveStatus::Ready);
619
620 enclave.enter_secure_mode().unwrap();
621 assert_eq!(enclave.status(), EnclaveStatus::Secure);
622
623 enclave.exit_secure_mode().unwrap();
624 assert_eq!(enclave.status(), EnclaveStatus::Ready);
625
626 enclave.destroy().unwrap();
627 assert_eq!(enclave.status(), EnclaveStatus::Destroyed);
628 }
629
630 #[test]
631 fn test_attestation() {
632 let mut enclave = SecureEnclave::with_sgx();
633 enclave.initialize().unwrap();
634
635 let challenge = [0u8; 32];
636 let report = enclave.attest(&challenge).unwrap();
637
638 assert_eq!(report.platform, TeePlatform::IntelSgx);
639 assert!(report.is_valid);
640
641 let verified = enclave.verify_attestation(&report).unwrap();
642 assert!(verified);
643 }
644
645 #[test]
646 fn test_seal_unseal() {
647 let mut enclave = SecureEnclave::new(EnclaveConfig::default());
648 enclave.initialize().unwrap();
649
650 let secret = b"Hope Genome Master Key";
651 let sealed = enclave
652 .seal(secret, SealingPolicy::EnclaveIdentity)
653 .unwrap();
654
655 assert_ne!(&sealed.ciphertext, secret);
656
657 let unsealed = enclave.unseal(&sealed).unwrap();
658 assert_eq!(&unsealed, secret);
659 }
660
661 #[test]
662 fn test_protected_watchdog() {
663 let mut watchdog = ProtectedWatchdog::new(TeePlatform::Simulated);
664 watchdog.initialize().unwrap();
665
666 assert!(!watchdog.is_locked());
667 assert_eq!(watchdog.violation_count(), 0);
668
669 for i in 1..=10 {
670 let count = watchdog.record_violation().unwrap();
671 assert_eq!(count, i);
672 }
673
674 assert!(watchdog.is_locked());
675
676 let allowed = watchdog.check_action("test").unwrap();
677 assert!(!allowed);
678 }
679
680 #[test]
681 fn test_execute_secure() {
682 let mut enclave = SecureEnclave::new(EnclaveConfig::default());
683 enclave.initialize().unwrap();
684
685 let result = enclave
686 .execute_secure(|| {
687 42 * 2
689 })
690 .unwrap();
691
692 assert_eq!(result, 84);
693 assert_eq!(enclave.stats().ecalls, 1);
694 assert_eq!(enclave.stats().ocalls, 1);
695 }
696}