1use std::io::Cursor;
4
5use serde::{Deserialize, Serialize};
6
7use aetheris_protocol::error::EncodeError;
8use aetheris_protocol::events::{ComponentUpdate, NetworkEvent, ReplicationEvent, WireEvent};
9use aetheris_protocol::traits::Encoder;
10use aetheris_protocol::types::{ClientId, ComponentKind, NetworkId};
11
12#[derive(Debug, Serialize, Deserialize)]
16struct PacketHeader {
17 network_id: NetworkId,
18 component_kind: ComponentKind,
19 tick: u64,
20}
21
22#[derive(Debug, Default)]
27pub struct SerdeEncoder;
28
29impl SerdeEncoder {
30 #[must_use]
32 pub fn new() -> Self {
33 Self
34 }
35
36 pub fn encode_event(
41 &self,
42 event: &aetheris_protocol::events::NetworkEvent,
43 ) -> Result<Vec<u8>, EncodeError> {
44 let wire_event = match event {
45 NetworkEvent::Ping { tick, .. } => WireEvent::Ping { tick: *tick },
46 NetworkEvent::Pong { tick } => WireEvent::Pong { tick: *tick },
47 NetworkEvent::Auth { session_token } => WireEvent::Auth {
48 session_token: session_token.clone(),
49 },
50 NetworkEvent::Fragment { fragment, .. } => WireEvent::Fragment(fragment.clone()),
51 NetworkEvent::StressTest { count, rotate, .. } => WireEvent::StressTest {
52 count: *count,
53 rotate: *rotate,
54 },
55 NetworkEvent::Spawn {
56 entity_type,
57 x,
58 y,
59 rot,
60 ..
61 } => WireEvent::Spawn {
62 entity_type: *entity_type,
63 x: *x,
64 y: *y,
65 rot: *rot,
66 },
67 NetworkEvent::ClearWorld { .. } => WireEvent::ClearWorld,
68 NetworkEvent::ClientConnected(_)
69 | NetworkEvent::ClientDisconnected(_)
70 | NetworkEvent::UnreliableMessage { .. }
71 | NetworkEvent::ReliableMessage { .. }
72 | NetworkEvent::SessionClosed(_)
73 | NetworkEvent::StreamReset(_) => {
74 return Err(EncodeError::Io(std::io::Error::other(format!(
75 "Cannot encode local-only variant as wire event: {event:?}"
76 ))));
77 }
78 };
79 rmp_serde::to_vec(&wire_event)
80 .map_err(|e| EncodeError::Io(std::io::Error::other(e.to_string())))
81 }
82
83 pub fn decode_event(
88 &self,
89 data: &[u8],
90 ) -> Result<aetheris_protocol::events::NetworkEvent, EncodeError> {
91 let wire_event: WireEvent = rmp_serde::from_slice(data).map_err(|e| {
92 EncodeError::MalformedPayload {
93 offset: 0, message: e.to_string(),
95 }
96 })?;
97
98 Ok(match wire_event {
99 WireEvent::Ping { tick } => NetworkEvent::Ping {
100 client_id: ClientId(0), tick,
102 },
103 WireEvent::Pong { tick } => NetworkEvent::Pong { tick },
104 WireEvent::Auth { session_token } => NetworkEvent::Auth { session_token },
105 WireEvent::Fragment(fragment) => NetworkEvent::Fragment {
106 client_id: ClientId(0),
107 fragment,
108 },
109 WireEvent::StressTest { count, rotate } => NetworkEvent::StressTest {
110 client_id: ClientId(0), count,
112 rotate,
113 },
114 WireEvent::Spawn {
115 entity_type,
116 x,
117 y,
118 rot,
119 } => NetworkEvent::Spawn {
120 client_id: ClientId(0),
121 entity_type,
122 x,
123 y,
124 rot,
125 },
126 WireEvent::ClearWorld => NetworkEvent::ClearWorld {
127 client_id: ClientId(0),
128 },
129 })
130 }
131}
132
133impl Encoder for SerdeEncoder {
134 fn encode_event(&self, event: &NetworkEvent) -> Result<Vec<u8>, EncodeError> {
135 self.encode_event(event)
136 }
137
138 fn decode_event(&self, data: &[u8]) -> Result<NetworkEvent, EncodeError> {
139 self.decode_event(data)
140 }
141
142 fn encode(&self, event: &ReplicationEvent, buffer: &mut [u8]) -> Result<usize, EncodeError> {
143 #[cfg(not(target_arch = "wasm32"))]
144 let start = std::time::Instant::now();
145 let header = PacketHeader {
146 network_id: event.network_id,
147 component_kind: event.component_kind,
148 tick: event.tick,
149 };
150
151 let mut cursor = Cursor::new(buffer);
152 let mut serializer = rmp_serde::Serializer::new(&mut cursor);
153
154 header.serialize(&mut serializer).map_err(|_e| {
155 metrics::counter!("aetheris_encoder_errors_total", "type" => "header_serialize_fail")
156 .increment(1);
157 EncodeError::BufferOverflow {
159 needed: 32, available: cursor.get_ref().len(),
161 }
162 })?;
163
164 let header_len = usize::try_from(cursor.position()).unwrap_or(usize::MAX);
165 let payload_len = event.payload.len();
166 let total_needed = header_len + payload_len;
167
168 if total_needed > cursor.get_ref().len() {
169 metrics::counter!("aetheris_encoder_errors_total", "type" => "buffer_overflow")
170 .increment(1);
171 return Err(EncodeError::BufferOverflow {
172 needed: total_needed,
173 available: cursor.get_ref().len(),
174 });
175 }
176
177 let slice = cursor.into_inner();
179 slice[header_len..total_needed].copy_from_slice(&event.payload);
180
181 #[allow(clippy::cast_precision_loss)]
182 metrics::histogram!(
183 "aetheris_encoder_payload_size_bytes",
184 "operation" => "encode"
185 )
186 .record(total_needed as f64);
187
188 #[cfg(not(target_arch = "wasm32"))]
189 metrics::histogram!(
190 "aetheris_encoder_encode_duration_seconds",
191 "kind" => event.component_kind.0.to_string()
192 )
193 .record(start.elapsed().as_secs_f64());
194
195 Ok(total_needed)
196 }
197
198 fn decode(&self, buffer: &[u8]) -> Result<ComponentUpdate, EncodeError> {
199 #[cfg(not(target_arch = "wasm32"))]
200 let start = std::time::Instant::now();
201 let mut cursor = Cursor::new(buffer);
202 let mut deserializer = rmp_serde::Deserializer::new(&mut cursor);
203
204 let header = PacketHeader::deserialize(&mut deserializer).map_err(|e| {
205 metrics::counter!("aetheris_encoder_errors_total", "type" => "malformed_payload")
206 .increment(1);
207 EncodeError::MalformedPayload {
208 offset: usize::try_from(cursor.position()).unwrap_or(usize::MAX),
209 message: e.to_string(),
210 }
211 })?;
212
213 let header_len = usize::try_from(cursor.position()).unwrap_or(usize::MAX);
214 let payload = buffer
215 .get(header_len..)
216 .ok_or(EncodeError::MalformedPayload {
217 offset: header_len,
218 message: "Payload slice out of bounds".to_string(),
219 })?
220 .to_vec();
221
222 #[allow(clippy::cast_precision_loss)]
223 metrics::histogram!(
224 "aetheris_encoder_payload_size_bytes",
225 "operation" => "decode"
226 )
227 .record(buffer.len() as f64);
228
229 #[cfg(not(target_arch = "wasm32"))]
230 metrics::histogram!(
231 "aetheris_encoder_decode_duration_seconds",
232 "kind" => header.component_kind.0.to_string()
233 )
234 .record(start.elapsed().as_secs_f64());
235
236 Ok(ComponentUpdate {
237 network_id: header.network_id,
238 component_kind: header.component_kind,
239 payload,
240 tick: header.tick,
241 })
242 }
243
244 fn max_encoded_size(&self) -> usize {
245 aetheris_protocol::MAX_SAFE_PAYLOAD_SIZE
246 }
247}
248
249#[cfg(test)]
250mod tests {
251 use super::*;
252 use proptest::prelude::*;
253
254 #[test]
255 fn test_roundtrip() {
256 let encoder = SerdeEncoder::new();
257 let event = ReplicationEvent {
258 network_id: NetworkId(42),
259 component_kind: ComponentKind(1),
260 payload: vec![1, 2, 3, 4],
261 tick: 100,
262 };
263
264 let mut buffer = [0u8; 1200];
265 let bytes_written = encoder.encode(&event, &mut buffer).unwrap();
266 assert!(bytes_written > 0);
267
268 let update = encoder.decode(&buffer[..bytes_written]).unwrap();
269 assert_eq!(update.network_id, event.network_id);
270 assert_eq!(update.component_kind, event.component_kind);
271 assert_eq!(update.tick, event.tick);
272 assert_eq!(update.payload, event.payload);
273 }
274 #[test]
275 fn test_fragment_roundtrip() {
276 let encoder = SerdeEncoder::new();
277 let fragment = aetheris_protocol::events::FragmentedEvent {
278 message_id: 123,
279 fragment_index: 1,
280 total_fragments: 5,
281 payload: vec![1, 2, 3],
282 };
283
284 let event = NetworkEvent::Fragment {
285 client_id: aetheris_protocol::types::ClientId(0),
286 fragment: fragment.clone(),
287 };
288
289 let output = encoder.encode_event(&event).unwrap();
290 let decoded = encoder.decode_event(&output).unwrap();
291
292 if let NetworkEvent::Fragment {
293 client_id: _,
294 fragment: decoded_fragment,
295 } = decoded
296 {
297 assert_eq!(decoded_fragment.message_id, fragment.message_id);
298 assert_eq!(decoded_fragment.fragment_index, fragment.fragment_index);
299 assert_eq!(decoded_fragment.total_fragments, fragment.total_fragments);
300 assert_eq!(decoded_fragment.payload, fragment.payload);
301 } else {
302 panic!("Decoded event is not a Fragment: {decoded:?}");
303 }
304 }
305
306 #[test]
307 fn test_buffer_overflow() {
308 let encoder = SerdeEncoder::new();
309 let event = ReplicationEvent {
310 network_id: NetworkId(42),
311 component_kind: ComponentKind(1),
312 payload: vec![1, 2, 3, 4],
313 tick: 100,
314 };
315
316 let mut small_buffer = [0u8; 1];
317 let result = encoder.encode(&event, &mut small_buffer);
318 assert!(matches!(result, Err(EncodeError::BufferOverflow { .. })));
319 }
320
321 #[test]
322 fn test_malformed_payload() {
323 let encoder = SerdeEncoder::new();
324 let garbage = [0xff, 0xff, 0xff, 0xff];
325 let result = encoder.decode(&garbage);
326 if let Err(EncodeError::MalformedPayload { message, .. }) = result {
327 assert!(!message.is_empty());
328 } else {
329 panic!("Expected MalformedPayload error, got {result:?}");
330 }
331 }
332
333 proptest! {
334 #[test]
335 fn test_fuzz_decode(ref bytes in any::<Vec<u8>>()) {
336 let encoder = SerdeEncoder::new();
337 let _ = encoder.decode(bytes);
339 }
340
341 #[test]
342 fn test_fuzz_roundtrip(
343 nid in any::<u64>(),
344 kind in any::<u16>(),
345 tick in any::<u64>(),
346 ref payload in any::<Vec<u8>>()
347 ) {
348 let encoder = SerdeEncoder::new();
349 let event = ReplicationEvent {
350 network_id: NetworkId(nid),
351 component_kind: ComponentKind(kind),
352 payload: payload.clone(),
353 tick,
354 };
355
356 let mut buffer = vec![0u8; 2048 + payload.len()];
357 if let Ok(written) = encoder.encode(&event, &mut buffer) {
358 let update = encoder.decode(&buffer[..written])
359 .expect("Round-trip decode failed during fuzzed test");
360 assert_eq!(update.network_id, event.network_id);
361 assert_eq!(update.component_kind, event.component_kind);
362 assert_eq!(update.tick, event.tick);
363 assert_eq!(update.payload, event.payload);
364 }
365 }
366 }
367}