Skip to content

Commit 6f7a8dd

Browse files
jstuczynjstuczyn
committed
use more strict typing for KKTFrame fields
1 parent 20b0fd1 commit 6f7a8dd

File tree

8 files changed

+139
-121
lines changed

8 files changed

+139
-121
lines changed

‎common/nym-kkt/benches/benches.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ pub fn kkt_benchmark(c: &mut Criterion) {
160160
b.iter(|| {
161161
responder_process(
162162
&mut r_context,
163-
i_frame_r.session_id_ref(),
163+
i_frame_r.session_id(),
164164
responder_ed25519_keypair.private_key(),
165165
&responder_kem_public_key,
166166
)
@@ -170,7 +170,7 @@ pub fn kkt_benchmark(c: &mut Criterion) {
170170
);
171171
let r_frame = responder_process(
172172
&mut r_context,
173-
i_frame_r.session_id_ref(),
173+
i_frame_r.session_id(),
174174
responder_ed25519_keypair.private_key(),
175175
&responder_kem_public_key,
176176
)
@@ -301,7 +301,7 @@ pub fn kkt_benchmark(c: &mut Criterion) {
301301
b.iter(|| {
302302
responder_process(
303303
&mut r_context,
304-
i_frame_r.session_id_ref(),
304+
i_frame_r.session_id(),
305305
responder_ed25519_keypair.private_key(),
306306
&responder_kem_public_key,
307307
)
@@ -312,7 +312,7 @@ pub fn kkt_benchmark(c: &mut Criterion) {
312312

313313
let r_frame = responder_process(
314314
&mut r_context,
315-
i_frame_r.session_id_ref(),
315+
i_frame_r.session_id(),
316316
responder_ed25519_keypair.private_key(),
317317
&responder_kem_public_key,
318318
)
@@ -450,7 +450,7 @@ pub fn kkt_benchmark(c: &mut Criterion) {
450450
b.iter(|| {
451451
responder_process(
452452
&mut r_context,
453-
i_frame_r.session_id_ref(),
453+
i_frame_r.session_id(),
454454
responder_ed25519_keypair.private_key(),
455455
&responder_kem_public_key,
456456
)
@@ -461,7 +461,7 @@ pub fn kkt_benchmark(c: &mut Criterion) {
461461

462462
let r_frame = responder_process(
463463
&mut r_context,
464-
i_frame_r.session_id_ref(),
464+
i_frame_r.session_id(),
465465
responder_ed25519_keypair.private_key(),
466466
&responder_kem_public_key,
467467
)

‎common/nym-kkt/src/ciphersuite.rs‎

Lines changed: 37 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ impl Ciphersuite {
203203
},
204204
})
205205
}
206-
pub fn encode(&self) -> [u8; 4] {
206+
pub fn encode(&self) -> [u8; CIPHERSUITE_ENCODING_LEN] {
207207
// [kem, hash, hashlen, sig]
208208
[
209209
match self.kem {
@@ -227,55 +227,45 @@ impl Ciphersuite {
227227
},
228228
]
229229
}
230-
pub fn decode(encoding: &[u8]) -> Result<Self, KKTError> {
231-
if encoding.len() == 4 {
232-
let kem = match encoding[0] {
233-
0 => KEM::XWing,
234-
1 => KEM::MlKem768,
235-
2 => KEM::McEliece,
236-
255 => KEM::X25519,
237-
_ => {
238-
return Err(KKTError::CiphersuiteDecodingError {
239-
info: format!("Undefined KEM: {}", encoding[0]),
240-
});
241-
}
242-
};
243-
let hash_function = match encoding[1] {
244-
0 => HashFunction::Blake3,
245-
1 => HashFunction::SHAKE256,
246-
2 => HashFunction::SHAKE128,
247-
3 => HashFunction::SHA256,
248-
_ => {
249-
return Err(KKTError::CiphersuiteDecodingError {
250-
info: format!("Undefined Hash Function: {}", encoding[1]),
251-
});
252-
}
253-
};
230+
pub fn decode(encoding: [u8; CIPHERSUITE_ENCODING_LEN]) -> Result<Self, KKTError> {
231+
let kem = match encoding[0] {
232+
0 => KEM::XWing,
233+
1 => KEM::MlKem768,
234+
2 => KEM::McEliece,
235+
255 => KEM::X25519,
236+
_ => {
237+
return Err(KKTError::CiphersuiteDecodingError {
238+
info: format!("Undefined KEM: {}", encoding[0]),
239+
});
240+
}
241+
};
242+
let hash_function = match encoding[1] {
243+
0 => HashFunction::Blake3,
244+
1 => HashFunction::SHAKE256,
245+
2 => HashFunction::SHAKE128,
246+
3 => HashFunction::SHA256,
247+
_ => {
248+
return Err(KKTError::CiphersuiteDecodingError {
249+
info: format!("Undefined Hash Function: {}", encoding[1]),
250+
});
251+
}
252+
};
254253

255-
let custom_hash_length = match encoding[2] {
256-
0 => None,
257-
_ => Some(encoding[2]),
258-
};
254+
let custom_hash_length = match encoding[2] {
255+
0 => None,
256+
_ => Some(encoding[2]),
257+
};
259258

260-
let signature_scheme = match encoding[3] {
261-
0 => SignatureScheme::Ed25519,
262-
_ => {
263-
return Err(KKTError::CiphersuiteDecodingError {
264-
info: format!("Undefined Signature Scheme: {}", encoding[3]),
265-
});
266-
}
267-
};
259+
let signature_scheme = match encoding[3] {
260+
0 => SignatureScheme::Ed25519,
261+
_ => {
262+
return Err(KKTError::CiphersuiteDecodingError {
263+
info: format!("Undefined Signature Scheme: {}", encoding[3]),
264+
});
265+
}
266+
};
268267

269-
Self::resolve_ciphersuite(kem, hash_function, signature_scheme, custom_hash_length)
270-
} else {
271-
Err(KKTError::CiphersuiteDecodingError {
272-
info: format!(
273-
"Incorrect Encoding Length: actual: {} != expected: {}",
274-
encoding.len(),
275-
CIPHERSUITE_ENCODING_LEN
276-
),
277-
})
278-
}
268+
Self::resolve_ciphersuite(kem, hash_function, signature_scheme, custom_hash_length)
279269
}
280270
}
281271

‎common/nym-kkt/src/context.rs‎

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
22
// SPDX-License-Identifier: Apache-2.0
33

4+
use crate::ciphersuite::CIPHERSUITE_ENCODING_LEN;
45
use crate::{KKT_VERSION, ciphersuite::Ciphersuite, error::KKTError, frame::KKT_SESSION_ID_LEN};
56
use num_enum::{IntoPrimitive, TryFromPrimitive};
67
use std::fmt::Display;
@@ -139,6 +140,9 @@ impl KKTContext {
139140
}
140141

141142
pub const fn session_id_len(&self) -> usize {
143+
// note: if anyone decides to update this function and changes the constant value,
144+
// you will have to adjust encoding/decoding functions
145+
142146
// match self.role {
143147
// KKTRole::Initiator | KKTRole::Responder => SESSION_ID_LENGTH,
144148
// It doesn't make sense to send a session_id if we send messages in the clear
@@ -151,31 +155,27 @@ impl KKTContext {
151155
self.body_len() + self.signature_len() + self.header_len() + self.session_id_len()
152156
}
153157

154-
pub fn encode(&self) -> Result<Vec<u8>, KKTError> {
155-
let mut header_bytes: Vec<u8> = Vec::with_capacity(KKT_CONTEXT_LEN);
158+
pub fn encode(&self) -> Result<[u8; KKT_CONTEXT_LEN], KKTError> {
159+
let mut header_bytes = [0u8; KKT_CONTEXT_LEN];
156160
if self.message_sequence >= 1 << 4 {
157161
return Err(KKTError::MessageCountLimitReached);
158162
}
159163

160-
header_bytes.push((KKT_VERSION << 4) + self.message_sequence);
161-
header_bytes.push(u8::from(self.status) + u8::from(self.mode) + u8::from(self.role));
164+
let ciphersuite_bytes = self.ciphersuite.encode();
162165

163-
header_bytes.extend_from_slice(&self.ciphersuite.encode());
164-
header_bytes.push(0);
165-
Ok(header_bytes)
166-
}
166+
header_bytes[0] = (KKT_VERSION << 4) + self.message_sequence;
167+
header_bytes[1] = u8::from(self.status) + u8::from(self.mode) + u8::from(self.role);
167168

168-
pub fn try_decode(header_bytes: &[u8]) -> Result<Self, KKTError> {
169-
if header_bytes.len() != KKT_CONTEXT_LEN {
170-
return Err(KKTError::FrameDecodingError {
171-
info: format!(
172-
"Header - Invalid Header Length: actual: {} != expected: {}",
173-
header_bytes.len(),
174-
KKT_CONTEXT_LEN
175-
),
176-
});
169+
let mut i = 2;
170+
for b in ciphersuite_bytes.into_iter() {
171+
header_bytes[i] = b;
172+
i += 1;
177173
}
174+
header_bytes[i] = 0;
175+
Ok(header_bytes)
176+
}
178177

178+
pub fn try_decode(header_bytes: [u8; KKT_CONTEXT_LEN]) -> Result<Self, KKTError> {
179179
let kkt_version = (header_bytes[0] & 0b1111_0000) >> 4;
180180
let message_sequence_counter = header_bytes[0] & 0b0000_1111;
181181

@@ -202,12 +202,20 @@ impl KKTContext {
202202
info: format!("Header - Invalid KKT Mode: {raw_kkt_mode}"),
203203
})?;
204204

205+
let ciphersuite_bytes = header_bytes[2..6].try_into().map_err(|_| {
206+
KKTError::CiphersuiteDecodingError {
207+
info: format!(
208+
"Incorrect Encoding Length: actual: 4 != expected: {CIPHERSUITE_ENCODING_LEN}",
209+
),
210+
}
211+
})?;
212+
205213
Ok(KKTContext {
206214
version: kkt_version,
207215
status,
208216
mode,
209217
role,
210-
ciphersuite: Ciphersuite::decode(&header_bytes[2..6])?,
218+
ciphersuite: Ciphersuite::decode(ciphersuite_bytes)?,
211219
message_sequence: message_sequence_counter,
212220
})
213221
}
@@ -222,11 +230,11 @@ mod tests {
222230
let valid_context = KKTContext::new(
223231
KKTRole::Initiator,
224232
KKTMode::Mutual,
225-
Ciphersuite::decode(&[255, 1, 0, 0]).unwrap(),
233+
Ciphersuite::decode([255, 1, 0, 0]).unwrap(),
226234
)
227235
.unwrap();
228236
let encoded = valid_context.encode().unwrap();
229-
let decoded = KKTContext::try_decode(&encoded).unwrap();
237+
let decoded = KKTContext::try_decode(encoded).unwrap();
230238

231239
assert_eq!(decoded, valid_context);
232240
}

‎common/nym-kkt/src/encryption.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -231,12 +231,12 @@ mod test {
231231
let valid_context = KKTContext::new(
232232
KKTRole::Initiator,
233233
KKTMode::Mutual,
234-
Ciphersuite::decode(&[255, 1, 0, 0])?,
234+
Ciphersuite::decode([255, 1, 0, 0])?,
235235
)?;
236236
let dummy_frame = KKTFrame::new(
237-
&valid_context.encode()?,
237+
valid_context.encode()?,
238238
&[2u8; 32],
239-
&[3u8; KKT_SESSION_ID_LEN],
239+
[3u8; KKT_SESSION_ID_LEN],
240240
&[4u8; 64],
241241
);
242242

0 commit comments

Comments
 (0)