a3s_common/transport/
tee.rs

1//! TEE protocol types for secure communication with TEE environments.
2
3use serde::{Deserialize, Serialize};
4
5/// Message types for TEE communication
6#[derive(Debug, Clone, Serialize, Deserialize)]
7#[serde(tag = "type", rename_all = "snake_case")]
8pub enum TeeMessage {
9    Request(TeeRequest),
10    Response(TeeResponse),
11    Heartbeat { timestamp: i64 },
12    Error { code: i32, message: String },
13}
14
15/// Request to TEE environment
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct TeeRequest {
18    pub id: String,
19    pub session_id: String,
20    pub request_type: TeeRequestType,
21    pub payload: Vec<u8>,
22    pub timestamp: i64,
23}
24
25impl TeeRequest {
26    pub fn new(session_id: String, request_type: TeeRequestType, payload: Vec<u8>) -> Self {
27        Self {
28            id: uuid::Uuid::new_v4().to_string(),
29            session_id,
30            request_type,
31            payload,
32            timestamp: chrono::Utc::now().timestamp_millis(),
33        }
34    }
35}
36
37/// Types of TEE requests
38#[derive(Debug, Clone, Serialize, Deserialize)]
39#[serde(rename_all = "snake_case")]
40pub enum TeeRequestType {
41    InitSession,
42    ProcessMessage,
43    ExecuteTool,
44    StoreSecret,
45    RetrieveSecret,
46    DeleteSecret,
47    GetSessionState,
48    TerminateSession,
49}
50
51/// Response from TEE environment
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct TeeResponse {
54    pub request_id: String,
55    pub session_id: String,
56    pub status: TeeResponseStatus,
57    pub payload: Vec<u8>,
58    pub timestamp: i64,
59}
60
61impl TeeResponse {
62    pub fn success(request_id: String, session_id: String, payload: Vec<u8>) -> Self {
63        Self {
64            request_id,
65            session_id,
66            status: TeeResponseStatus::Success,
67            payload,
68            timestamp: chrono::Utc::now().timestamp_millis(),
69        }
70    }
71
72    pub fn error(request_id: String, session_id: String, code: i32, message: String) -> Self {
73        Self {
74            request_id,
75            session_id,
76            status: TeeResponseStatus::Error { code, message },
77            payload: Vec::new(),
78            timestamp: chrono::Utc::now().timestamp_millis(),
79        }
80    }
81}
82
83/// Response status
84#[derive(Debug, Clone, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86pub enum TeeResponseStatus {
87    Success,
88    Error { code: i32, message: String },
89    Pending,
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_tee_request_creation() {
98        let request = TeeRequest::new(
99            "session-123".to_string(),
100            TeeRequestType::ProcessMessage,
101            vec![1, 2, 3],
102        );
103        assert!(!request.id.is_empty());
104        assert_eq!(request.session_id, "session-123");
105    }
106
107    #[test]
108    fn test_tee_message_serialization() {
109        let request = TeeRequest::new(
110            "session-123".to_string(),
111            TeeRequestType::InitSession,
112            vec![],
113        );
114        let message = TeeMessage::Request(request);
115        let json = serde_json::to_string(&message).unwrap();
116        let parsed: TeeMessage = serde_json::from_str(&json).unwrap();
117        assert!(matches!(parsed, TeeMessage::Request(_)));
118    }
119
120    #[test]
121    fn test_tee_response_success() {
122        let resp = TeeResponse::success("req-1".into(), "sess-1".into(), vec![42]);
123        assert!(matches!(resp.status, TeeResponseStatus::Success));
124        assert_eq!(resp.payload, vec![42]);
125    }
126
127    #[test]
128    fn test_tee_response_error() {
129        let resp = TeeResponse::error("req-1".into(), "sess-1".into(), 500, "fail".into());
130        assert!(matches!(resp.status, TeeResponseStatus::Error { .. }));
131        assert!(resp.payload.is_empty());
132    }
133}
a3s_common/transport/tee.rs

a3s_common/transport/
tee.rs
