Skip to content

Commit 4608d5b

Browse files
hyperfinitismhyperfinitism
committed
feat(enclave,proxy): adopt length prefix for vsock connection
1 parent fc960a1 commit 4608d5b

File tree

2 files changed

+212
-165
lines changed

2 files changed

+212
-165
lines changed

‎enclave/src/main.rs‎

Lines changed: 182 additions & 154 deletions
Original file line numberDiff line numberDiff line change
@@ -12,195 +12,223 @@ mod nsm;
1212
mod session;
1313
mod types;
1414

15-
const MAX_BUFFER_SIZE: usize = 8192;
15+
const LENGTH_PREFIX_SIZE: usize = 4;
1616

1717
#[derive(Debug, Parser)]
1818
#[command(name = "rafwne-enclave")]
1919
struct Args {
2020
/// Vsock port to listen on.
2121
#[arg(long, default_value_t = 5000)]
2222
vsock_port: u32,
23+
24+
/// Vsock buffer size in bytes.
25+
#[arg(long, default_value_t = 8192)]
26+
vsock_buffer_size: usize,
2327
}
2428

2529
fn main() -> Result<()> {
2630
let args = Args::parse();
2731
let addr = VsockAddr::new(VMADDR_CID_ANY, args.vsock_port);
2832
let listener = VsockListener::bind(&addr)?;
33+
let max_request_size = args.vsock_buffer_size - LENGTH_PREFIX_SIZE;
2934

3035
loop {
3136
let (mut stream, _) = listener.accept()?;
32-
let _ = handle_vsock_stream(&mut stream);
37+
let _ = handle_vsock_stream(&mut stream, max_request_size);
3338
}
3439
}
3540

36-
pub(crate) fn handle_vsock_stream(stream: &mut VsockStream) -> Result<()> {
37-
let mut buffer = [0; MAX_BUFFER_SIZE];
38-
41+
pub(crate) fn handle_vsock_stream(stream: &mut VsockStream, max_request_size: usize) -> Result<()> {
3942
loop {
40-
match stream.read(&mut buffer) {
41-
Ok(0) => {
43+
// Read 4-byte big-endian length prefix
44+
let mut len_buf = [0u8; LENGTH_PREFIX_SIZE];
45+
match stream.read_exact(&mut len_buf) {
46+
Ok(()) => {}
47+
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
4248
break;
4349
}
44-
Ok(len) => {
45-
let request_str = String::from_utf8_lossy(&buffer[..len]);
46-
let request = match types::Request::from_json(&request_str) {
47-
Ok(r) => r,
48-
Err(e) => {
49-
let response = types::Response::Error {
50-
error: format!("invalid request: {e}"),
51-
};
52-
let response_str = response.to_json()?;
53-
stream.write_all(response_str.as_bytes())?;
54-
continue;
50+
Err(e) => return Err(e.into()),
51+
}
52+
let req_len = u32::from_be_bytes(len_buf) as usize;
53+
54+
if req_len > max_request_size {
55+
let response = types::Response::Error {
56+
error: format!("request too large: {} > {}", req_len, max_request_size),
57+
};
58+
let response_str = response.to_json()?;
59+
write_response(stream, response_str.as_bytes())?;
60+
continue;
61+
}
62+
63+
// Read the exact amount of request data
64+
let mut buffer = vec![0u8; req_len];
65+
stream.read_exact(&mut buffer)?;
66+
67+
let request_str = String::from_utf8_lossy(&buffer);
68+
let request = match types::Request::from_json(&request_str) {
69+
Ok(r) => r,
70+
Err(e) => {
71+
let response = types::Response::Error {
72+
error: format!("invalid request: {e}"),
73+
};
74+
let response_str = response.to_json()?;
75+
write_response(stream, response_str.as_bytes())?;
76+
continue;
77+
}
78+
};
79+
80+
match request {
81+
types::Request::Init => {
82+
let sess = session::new_session()?;
83+
let response = types::Response::Init {
84+
session_id: sess.session_id.clone(),
85+
enclave_pubkey_b64: crypto::b64_encode(&sess.enclave_pub),
86+
};
87+
session::insert_session(sess)?;
88+
let response_str = response.to_json()?;
89+
write_response(stream, response_str.as_bytes())?;
90+
}
91+
types::Request::KeyExchange {
92+
session_id,
93+
client_pubkey_b64,
94+
} => {
95+
let mut sess = session::take_session(&session_id)?;
96+
let client_pub = crypto::b64_decode(&client_pubkey_b64)?;
97+
98+
let enclave_priv = sess
99+
.enclave_priv
100+
.take()
101+
.ok_or_else(|| anyhow::anyhow!("session already completed key exchange"))?;
102+
103+
let shared = session::ecdh_shared_secret(enclave_priv, &client_pub)?;
104+
105+
let ck = crypto::derive_key(&shared, b"CK");
106+
let mk = crypto::derive_key(&shared, b"MK");
107+
let vk = crypto::derive_key(&shared, b"VK");
108+
109+
sess.client_pub = Some(client_pub);
110+
sess.ck = Some(ck);
111+
sess.mk = Some(mk);
112+
sess.vk = Some(vk);
113+
114+
let user_data = session::session_user_data(&sess)?;
115+
let nonce = {
116+
let rnd = nsm::get_random()?;
117+
if rnd.len() >= 64 {
118+
rnd[..64].to_vec()
119+
} else {
120+
rnd
55121
}
56122
};
57123

58-
match request {
59-
types::Request::Init => {
60-
let sess = session::new_session()?;
61-
let response = types::Response::Init {
62-
session_id: sess.session_id.clone(),
63-
enclave_pubkey_b64: crypto::b64_encode(&sess.enclave_pub),
64-
};
65-
session::insert_session(sess)?;
66-
let response_str = response.to_json()?;
67-
stream.write_all(response_str.as_bytes())?;
124+
let doc = nsm::get_attestation_document(user_data, nonce)?;
125+
let response = types::Response::KeyExchange {
126+
attestation_document_b64: crypto::b64_encode(&doc),
127+
};
128+
session::put_session(sess)?;
129+
let response_str = response.to_json()?;
130+
write_response(stream, response_str.as_bytes())?;
131+
}
132+
types::Request::Add { session_id, x, y } => {
133+
let sess = session::get_session(&session_id)?;
134+
let ck = sess
135+
.ck
136+
.as_ref()
137+
.ok_or_else(|| anyhow::anyhow!("CK not set"))?;
138+
let mk = sess
139+
.mk
140+
.as_ref()
141+
.ok_or_else(|| anyhow::anyhow!("MK not set"))?;
142+
143+
let x_nonce = crypto::b64_decode(&x.nonce_b64)?;
144+
let x_ct = crypto::b64_decode(&x.ciphertext_b64)?;
145+
let y_nonce = crypto::b64_decode(&y.nonce_b64)?;
146+
let y_ct = crypto::b64_decode(&y.ciphertext_b64)?;
147+
148+
let x_pt = crypto::aes128gcm_decrypt(&ck[..16], &x_nonce, &x_ct)?;
149+
let y_pt = crypto::aes128gcm_decrypt(&ck[..16], &y_nonce, &y_ct)?;
150+
let x_u32 = crypto::u32_from_le_bytes(&x_pt)?;
151+
let y_u32 = crypto::u32_from_le_bytes(&y_pt)?;
152+
153+
let sum = x_u32
154+
.checked_add(y_u32)
155+
.ok_or_else(|| anyhow::anyhow!("u32 overflow"))?;
156+
let sum_pt = crypto::u32_to_le_bytes(sum);
157+
158+
let sum_nonce = {
159+
let rnd = nsm::get_random()?;
160+
if rnd.len() < 12 {
161+
return Err(anyhow::anyhow!("NSM random too short for nonce"));
68162
}
69-
types::Request::KeyExchange {
70-
session_id,
71-
client_pubkey_b64,
72-
} => {
73-
let mut sess = session::take_session(&session_id)?;
74-
let client_pub = crypto::b64_decode(&client_pubkey_b64)?;
75-
76-
let enclave_priv = sess.enclave_priv.take().ok_or_else(|| {
77-
anyhow::anyhow!("session already completed key exchange")
78-
})?;
79-
80-
let shared = session::ecdh_shared_secret(enclave_priv, &client_pub)?;
81-
82-
let ck = crypto::derive_key(&shared, b"CK");
83-
let mk = crypto::derive_key(&shared, b"MK");
84-
let vk = crypto::derive_key(&shared, b"VK");
85-
86-
sess.client_pub = Some(client_pub);
87-
sess.ck = Some(ck);
88-
sess.mk = Some(mk);
89-
sess.vk = Some(vk);
90-
91-
let user_data = session::session_user_data(&sess)?;
92-
let nonce = {
93-
let rnd = nsm::get_random()?;
94-
if rnd.len() >= 64 {
95-
rnd[..64].to_vec()
96-
} else {
97-
rnd
98-
}
99-
};
163+
rnd[..12].to_vec()
164+
};
165+
let sum_ct = crypto::aes128gcm_encrypt(&mk[..16], &sum_nonce, &sum_pt)?;
100166

101-
let doc = nsm::get_attestation_document(user_data, nonce)?;
102-
let response = types::Response::KeyExchange {
103-
attestation_document_b64: crypto::b64_encode(&doc),
167+
let response = types::Response::Add {
168+
sum: types::EncryptedBlob {
169+
nonce_b64: crypto::b64_encode(&sum_nonce),
170+
ciphertext_b64: crypto::b64_encode(&sum_ct),
171+
},
172+
};
173+
let response_str = response.to_json()?;
174+
write_response(stream, response_str.as_bytes())?;
175+
}
176+
types::Request::Attest {
177+
session_id,
178+
user_data_b64,
179+
nonce_b64,
180+
} => {
181+
let (user_data, nonce) = match (user_data_b64, nonce_b64) {
182+
(Some(ud), Some(nc)) => (crypto::b64_decode(&ud)?, crypto::b64_decode(&nc)?),
183+
(ud_opt, nc_opt) => {
184+
let sid = session_id.ok_or_else(|| {
185+
anyhow::anyhow!(
186+
"session_id is required when user_data/nonce not fully provided"
187+
)
188+
})?;
189+
let sess = session::get_session(&sid)?;
190+
let ud = match ud_opt {
191+
Some(ud) => crypto::b64_decode(&ud)?,
192+
None => session::session_user_data(&sess)?,
104193
};
105-
session::put_session(sess)?;
106-
let response_str = response.to_json()?;
107-
stream.write_all(response_str.as_bytes())?;
108-
}
109-
types::Request::Add { session_id, x, y } => {
110-
let sess = session::get_session(&session_id)?;
111-
let ck = sess
112-
.ck
113-
.as_ref()
114-
.ok_or_else(|| anyhow::anyhow!("CK not set"))?;
115-
let mk = sess
116-
.mk
117-
.as_ref()
118-
.ok_or_else(|| anyhow::anyhow!("MK not set"))?;
119-
120-
let x_nonce = crypto::b64_decode(&x.nonce_b64)?;
121-
let x_ct = crypto::b64_decode(&x.ciphertext_b64)?;
122-
let y_nonce = crypto::b64_decode(&y.nonce_b64)?;
123-
let y_ct = crypto::b64_decode(&y.ciphertext_b64)?;
124-
125-
let x_pt = crypto::aes128gcm_decrypt(&ck[..16], &x_nonce, &x_ct)?;
126-
let y_pt = crypto::aes128gcm_decrypt(&ck[..16], &y_nonce, &y_ct)?;
127-
let x_u32 = crypto::u32_from_le_bytes(&x_pt)?;
128-
let y_u32 = crypto::u32_from_le_bytes(&y_pt)?;
129-
130-
let sum = x_u32
131-
.checked_add(y_u32)
132-
.ok_or_else(|| anyhow::anyhow!("u32 overflow"))?;
133-
let sum_pt = crypto::u32_to_le_bytes(sum);
134-
135-
let sum_nonce = {
136-
let rnd = nsm::get_random()?;
137-
if rnd.len() < 12 {
138-
return Err(anyhow::anyhow!("NSM random too short for nonce"));
194+
let nc = match nc_opt {
195+
Some(nc) => crypto::b64_decode(&nc)?,
196+
None => {
197+
let rnd = nsm::get_random()?;
198+
if rnd.len() >= 64 {
199+
rnd[..64].to_vec()
200+
} else {
201+
rnd
202+
}
139203
}
140-
rnd[..12].to_vec()
141204
};
142-
let sum_ct = crypto::aes128gcm_encrypt(&mk[..16], &sum_nonce, &sum_pt)?;
143-
144-
let response = types::Response::Add {
145-
sum: types::EncryptedBlob {
146-
nonce_b64: crypto::b64_encode(&sum_nonce),
147-
ciphertext_b64: crypto::b64_encode(&sum_ct),
148-
},
149-
};
150-
let response_str = response.to_json()?;
151-
stream.write_all(response_str.as_bytes())?;
205+
(ud, nc)
152206
}
153-
types::Request::Attest {
154-
session_id,
155-
user_data_b64,
156-
nonce_b64,
157-
} => {
158-
let (user_data, nonce) = match (user_data_b64, nonce_b64) {
159-
(Some(ud), Some(nc)) => {
160-
(crypto::b64_decode(&ud)?, crypto::b64_decode(&nc)?)
161-
}
162-
(ud_opt, nc_opt) => {
163-
let sid = session_id.ok_or_else(|| {
164-
anyhow::anyhow!("session_id is required when user_data/nonce not fully provided")
165-
})?;
166-
let sess = session::get_session(&sid)?;
167-
let ud = match ud_opt {
168-
Some(ud) => crypto::b64_decode(&ud)?,
169-
None => session::session_user_data(&sess)?,
170-
};
171-
let nc = match nc_opt {
172-
Some(nc) => crypto::b64_decode(&nc)?,
173-
None => {
174-
let rnd = nsm::get_random()?;
175-
if rnd.len() >= 64 {
176-
rnd[..64].to_vec()
177-
} else {
178-
rnd
179-
}
180-
}
181-
};
182-
(ud, nc)
183-
}
184-
};
207+
};
185208

186-
let doc = nsm::get_attestation_document(user_data, nonce)?;
187-
let response = types::Response::Attest {
188-
attestation_document_b64: crypto::b64_encode(&doc),
189-
};
190-
let response_str = response.to_json()?;
191-
stream.write_all(response_str.as_bytes())?;
192-
}
193-
types::Request::Close { session_id } => {
194-
session::delete_session(&session_id)?;
195-
let response = types::Response::CloseOk {};
196-
let response_str = response.to_json()?;
197-
stream.write_all(response_str.as_bytes())?;
198-
}
199-
}
209+
let doc = nsm::get_attestation_document(user_data, nonce)?;
210+
let response = types::Response::Attest {
211+
attestation_document_b64: crypto::b64_encode(&doc),
212+
};
213+
let response_str = response.to_json()?;
214+
write_response(stream, response_str.as_bytes())?;
215+
}
216+
types::Request::Close { session_id } => {
217+
session::delete_session(&session_id)?;
218+
let response = types::Response::CloseOk {};
219+
let response_str = response.to_json()?;
220+
write_response(stream, response_str.as_bytes())?;
200221
}
201-
Err(e) => return Err(e.into()),
202222
}
203223
}
204224

205225
Ok(())
206226
}
227+
228+
/// Write response with 4-byte big-endian length prefix.
229+
fn write_response(stream: &mut VsockStream, data: &[u8]) -> Result<()> {
230+
let len = data.len() as u32;
231+
stream.write_all(&len.to_be_bytes())?;
232+
stream.write_all(data)?;
233+
Ok(())
234+
}

0 commit comments

Comments
 (0)