1use elara_core::{ElaraError, ElaraResult};
6
7use crate::{Extensions, FixedHeader, FIXED_HEADER_SIZE};
8
9pub const AUTH_TAG_SIZE: usize = 16;
11
12pub const MAX_FRAME_SIZE: usize = 1400;
14
15pub const MIN_FRAME_SIZE: usize = FIXED_HEADER_SIZE + AUTH_TAG_SIZE;
17
18#[derive(Clone, Debug)]
20pub struct Frame {
21 pub header: FixedHeader,
23 pub extensions: Extensions,
25 pub payload: Vec<u8>,
27 pub auth_tag: [u8; AUTH_TAG_SIZE],
29}
30
31impl Frame {
32 pub fn new(header: FixedHeader) -> Self {
34 Frame {
35 header,
36 extensions: Extensions::new(),
37 payload: Vec::new(),
38 auth_tag: [0u8; AUTH_TAG_SIZE],
39 }
40 }
41
42 pub fn parse(buf: &[u8]) -> ElaraResult<Self> {
44 if buf.len() < MIN_FRAME_SIZE {
45 return Err(ElaraError::BufferTooShort {
46 expected: MIN_FRAME_SIZE,
47 actual: buf.len(),
48 });
49 }
50
51 let header = FixedHeader::parse(buf)?;
53
54 if header.header_len as usize > buf.len() - AUTH_TAG_SIZE {
55 return Err(ElaraError::InvalidWireFormat(
56 "Header length exceeds frame".into(),
57 ));
58 }
59
60 let extensions =
62 if header.flags.has_extension() && header.header_len as usize > FIXED_HEADER_SIZE {
63 let ext_buf = &buf[FIXED_HEADER_SIZE..header.header_len as usize];
64 let (ext, _) = Extensions::parse(ext_buf, ext_buf.len())?;
65 ext
66 } else {
67 Extensions::new()
68 };
69
70 let payload_start = header.header_len as usize;
72 let payload_end = buf.len() - AUTH_TAG_SIZE;
73 let payload = buf[payload_start..payload_end].to_vec();
74
75 let mut auth_tag = [0u8; AUTH_TAG_SIZE];
77 auth_tag.copy_from_slice(&buf[payload_end..]);
78
79 Ok(Frame {
80 header,
81 extensions,
82 payload,
83 auth_tag,
84 })
85 }
86
87 pub fn serialize(&self) -> ElaraResult<Vec<u8>> {
89 let ext_size = if self.extensions.is_empty() {
90 0
91 } else {
92 self.extensions.serialized_size()
93 };
94
95 let total_size = FIXED_HEADER_SIZE + ext_size + self.payload.len() + AUTH_TAG_SIZE;
96
97 if total_size > MAX_FRAME_SIZE {
98 return Err(ElaraError::InvalidWireFormat(format!(
99 "Frame too large: {} > {}",
100 total_size, MAX_FRAME_SIZE
101 )));
102 }
103
104 let mut buf = vec![0u8; total_size];
105
106 let mut header = self.header.clone();
108 header.header_len = (FIXED_HEADER_SIZE + ext_size) as u16;
109 if !self.extensions.is_empty() {
110 header.flags.set_extension(true);
111 }
112 header.serialize(&mut buf)?;
113
114 if !self.extensions.is_empty() {
116 self.extensions
117 .serialize(&mut buf[FIXED_HEADER_SIZE..FIXED_HEADER_SIZE + ext_size])?;
118 }
119
120 let payload_start = FIXED_HEADER_SIZE + ext_size;
122 buf[payload_start..payload_start + self.payload.len()].copy_from_slice(&self.payload);
123
124 buf[total_size - AUTH_TAG_SIZE..].copy_from_slice(&self.auth_tag);
126
127 Ok(buf)
128 }
129
130 pub fn associated_data(&self) -> Vec<u8> {
132 let ext_size = if self.extensions.is_empty() {
133 0
134 } else {
135 self.extensions.serialized_size()
136 };
137
138 let mut aad = vec![0u8; FIXED_HEADER_SIZE + ext_size];
139
140 let mut header = self.header.clone();
141 header.header_len = (FIXED_HEADER_SIZE + ext_size) as u16;
142 header.serialize(&mut aad).unwrap();
143
144 if !self.extensions.is_empty() {
145 self.extensions
146 .serialize(&mut aad[FIXED_HEADER_SIZE..])
147 .unwrap();
148 }
149
150 aad
151 }
152
153 pub fn size(&self) -> usize {
155 let ext_size = if self.extensions.is_empty() {
156 0
157 } else {
158 self.extensions.serialized_size()
159 };
160 FIXED_HEADER_SIZE + ext_size + self.payload.len() + AUTH_TAG_SIZE
161 }
162
163 pub fn fits_mtu(&self) -> bool {
165 self.size() <= MAX_FRAME_SIZE
166 }
167}
168
169pub struct FrameBuilder {
171 frame: Frame,
172}
173
174impl FrameBuilder {
175 pub fn new(header: FixedHeader) -> Self {
176 FrameBuilder {
177 frame: Frame::new(header),
178 }
179 }
180
181 pub fn extensions(mut self, ext: Extensions) -> Self {
182 self.frame.extensions = ext;
183 self
184 }
185
186 pub fn payload(mut self, payload: Vec<u8>) -> Self {
187 self.frame.payload = payload;
188 self
189 }
190
191 pub fn auth_tag(mut self, tag: [u8; AUTH_TAG_SIZE]) -> Self {
192 self.frame.auth_tag = tag;
193 self
194 }
195
196 pub fn build(self) -> Frame {
197 self.frame
198 }
199}
200
201pub struct FrameSlice<'a> {
203 pub header: &'a [u8],
204 pub extensions: &'a [u8],
205 pub payload: &'a [u8],
206 pub auth_tag: &'a [u8; AUTH_TAG_SIZE],
207}
208
209impl<'a> FrameSlice<'a> {
210 pub fn from_bytes(buf: &'a [u8]) -> ElaraResult<Self> {
212 if buf.len() < MIN_FRAME_SIZE {
213 return Err(ElaraError::BufferTooShort {
214 expected: MIN_FRAME_SIZE,
215 actual: buf.len(),
216 });
217 }
218
219 let header_len = u16::from_le_bytes([buf[2], buf[3]]) as usize;
221
222 if header_len > buf.len() - AUTH_TAG_SIZE {
223 return Err(ElaraError::InvalidWireFormat(
224 "Header length exceeds frame".into(),
225 ));
226 }
227
228 let header = &buf[0..FIXED_HEADER_SIZE];
229 let extensions = &buf[FIXED_HEADER_SIZE..header_len];
230 let payload = &buf[header_len..buf.len() - AUTH_TAG_SIZE];
231 let auth_tag: &[u8; AUTH_TAG_SIZE] = buf[buf.len() - AUTH_TAG_SIZE..]
232 .try_into()
233 .map_err(|_| ElaraError::InvalidWireFormat("Invalid auth tag".into()))?;
234
235 Ok(FrameSlice {
236 header,
237 extensions,
238 payload,
239 auth_tag,
240 })
241 }
242
243 pub fn parse_header(&self) -> ElaraResult<FixedHeader> {
245 FixedHeader::parse(self.header)
246 }
247
248 pub fn parse_extensions(&self) -> ElaraResult<Extensions> {
250 if self.extensions.is_empty() {
251 Ok(Extensions::new())
252 } else {
253 let (ext, _) = Extensions::parse(self.extensions, self.extensions.len())?;
254 Ok(ext)
255 }
256 }
257}
258
259#[cfg(test)]
260mod tests {
261 use super::*;
262 use elara_core::{NodeId, PacketClass, SessionId};
263
264 #[test]
265 fn test_frame_roundtrip() {
266 let header = FixedHeader {
267 session_id: SessionId::new(12345),
268 node_id: NodeId::new(67890),
269 class: PacketClass::Perceptual,
270 time_hint: 100,
271 ..Default::default()
272 };
273
274 let mut ext = Extensions::new();
275 ext.ratchet_id = Some(42);
276
277 let frame = FrameBuilder::new(header)
278 .extensions(ext)
279 .payload(vec![1, 2, 3, 4, 5])
280 .auth_tag([0xAA; AUTH_TAG_SIZE])
281 .build();
282
283 let bytes = frame.serialize().unwrap();
284 let parsed = Frame::parse(&bytes).unwrap();
285
286 assert_eq!(parsed.header.session_id, frame.header.session_id);
287 assert_eq!(parsed.header.node_id, frame.header.node_id);
288 assert_eq!(parsed.header.class, frame.header.class);
289 assert_eq!(parsed.extensions.ratchet_id, Some(42));
290 assert_eq!(parsed.payload, vec![1, 2, 3, 4, 5]);
291 assert_eq!(parsed.auth_tag, [0xAA; AUTH_TAG_SIZE]);
292 }
293
294 #[test]
295 fn test_frame_slice_zero_copy() {
296 let header = FixedHeader::new(SessionId::new(1), NodeId::new(2));
297 let frame = FrameBuilder::new(header)
298 .payload(vec![10, 20, 30])
299 .auth_tag([0xBB; AUTH_TAG_SIZE])
300 .build();
301
302 let bytes = frame.serialize().unwrap();
303 let slice = FrameSlice::from_bytes(&bytes).unwrap();
304
305 assert_eq!(slice.payload, &[10, 20, 30]);
306 assert_eq!(slice.auth_tag, &[0xBB; AUTH_TAG_SIZE]);
307
308 let parsed_header = slice.parse_header().unwrap();
309 assert_eq!(parsed_header.session_id, SessionId::new(1));
310 }
311
312 #[test]
313 fn test_frame_size_limits() {
314 let header = FixedHeader::default();
315 let frame = FrameBuilder::new(header)
316 .payload(vec![0u8; MAX_FRAME_SIZE]) .build();
318
319 assert!(!frame.fits_mtu());
320 assert!(frame.serialize().is_err());
321 }
322
323 #[test]
324 fn test_associated_data() {
325 let header = FixedHeader::new(SessionId::new(100), NodeId::new(200));
326 let mut ext = Extensions::new();
327 ext.key_epoch = Some(5);
328
329 let frame = FrameBuilder::new(header)
330 .extensions(ext)
331 .payload(vec![1, 2, 3])
332 .build();
333
334 let aad = frame.associated_data();
335
336 assert!(aad.len() > FIXED_HEADER_SIZE);
338 assert!(aad.len() < frame.size());
339 }
340}