ponk_protocol/
lib.rs

1#![doc = include_str!("../README.md")]
2#![forbid(unsafe_code)]
3
4use std::collections::HashMap;
5use std::fmt;
6use std::net::SocketAddr;
7use std::time::{Duration, Instant};
8
9pub const DEFAULT_PORT: u16 = 5583;
10pub const MULTICAST_ADDR: [u8; 4] = [239, 255, 10, 24];
11pub const HEADER_LEN: usize = 52;
12pub const MAX_DATAGRAM_LEN: usize = 65_507;
13pub const DEFAULT_MAX_ASSEMBLIES: usize = 64;
14pub const DEFAULT_MAX_FRAME_PAYLOAD_BYTES: usize = 1_048_576;
15pub const DEFAULT_MAX_BUFFERED_BYTES: usize = 8_388_608;
16pub const DEFAULT_MAX_PATHS: usize = 4_096;
17pub const DEFAULT_MAX_POINTS: usize = 65_535;
18pub const DEFAULT_ASSEMBLY_TIMEOUT: Duration = Duration::from_millis(300);
19
20const MAGIC: &[u8; 8] = b"PONK-UDP";
21const PROTOCOL_VERSION: u8 = 0;
22/// Byte width of the sender-name field in every PONK datagram header.
23///
24/// Encoders truncate at a UTF-8 character boundary. Callers that compose a
25/// display name can use this constant to keep the full name on the wire.
26pub const SENDER_NAME_LEN: usize = 32;
27const METADATA_KEY_LEN: usize = 8;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum DataFormat {
31    XyRgbU16,
32    XyF32RgbU8,
33}
34
35impl DataFormat {
36    fn from_wire(value: u8) -> Result<Self, PonkError> {
37        match value {
38            0 => Ok(Self::XyRgbU16),
39            1 => Ok(Self::XyF32RgbU8),
40            other => Err(PonkError::UnsupportedDataFormat(other)),
41        }
42    }
43
44    fn wire(self) -> u8 {
45        match self {
46            Self::XyRgbU16 => 0,
47            Self::XyF32RgbU8 => 1,
48        }
49    }
50
51    fn bytes_per_point(self) -> usize {
52        match self {
53            Self::XyRgbU16 => 10,
54            Self::XyF32RgbU8 => 11,
55        }
56    }
57}
58
59#[derive(Debug, Clone, PartialEq)]
60pub struct PonkPoint {
61    pub x: f32,
62    pub y: f32,
63    pub rgb: [u8; 3],
64}
65
66#[derive(Debug, Clone, PartialEq)]
67pub struct PonkMetadata {
68    pub key: String,
69    pub value: f32,
70}
71
72#[derive(Debug, Clone, PartialEq)]
73pub struct PonkPath {
74    pub metadata: Vec<PonkMetadata>,
75    pub points: Vec<PonkPoint>,
76}
77
78#[derive(Debug, Clone, PartialEq)]
79pub struct PonkFrame {
80    pub sender_id: u32,
81    pub sender_name: String,
82    pub frame_number: u8,
83    pub paths: Vec<PonkPath>,
84}
85
86/// Resource limits applied while decoding untrusted PONK datagrams.
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub struct PonkLimits {
89    pub max_assemblies: usize,
90    pub max_frame_payload_bytes: usize,
91    pub max_buffered_bytes: usize,
92    pub max_paths: usize,
93    pub max_points: usize,
94    pub assembly_timeout: Duration,
95}
96
97impl Default for PonkLimits {
98    fn default() -> Self {
99        Self {
100            max_assemblies: DEFAULT_MAX_ASSEMBLIES,
101            max_frame_payload_bytes: DEFAULT_MAX_FRAME_PAYLOAD_BYTES,
102            max_buffered_bytes: DEFAULT_MAX_BUFFERED_BYTES,
103            max_paths: DEFAULT_MAX_PATHS,
104            max_points: DEFAULT_MAX_POINTS,
105            assembly_timeout: DEFAULT_ASSEMBLY_TIMEOUT,
106        }
107    }
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub enum PonkError {
112    UnsupportedProtocolVersion(u8),
113    UnsupportedDataFormat(u8),
114    DatagramTooSmall,
115    DatagramTooLarge { max: usize, actual: usize },
116    MaxDatagramTooSmall { min: usize, actual: usize },
117    TooManyChunks(usize),
118    TooManyMetadata(usize),
119    TooManyPoints(usize),
120    TooManyPaths { max: usize },
121    TooManyTotalPoints { max: usize },
122    FramePayloadTooLarge { max: usize },
123    InvalidCoordinate,
124    InvalidChunkHeader,
125    MalformedPayload,
126}
127
128impl fmt::Display for PonkError {
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        match self {
131            Self::UnsupportedProtocolVersion(version) => {
132                write!(f, "unsupported PONK protocol version {version}")
133            }
134            Self::UnsupportedDataFormat(format) => {
135                write!(f, "Unsupported PONK data format {format}")
136            }
137            Self::DatagramTooSmall => write!(f, "PONK datagram is too small"),
138            Self::DatagramTooLarge { max, actual } => {
139                write!(f, "PONK datagram is {actual} bytes, max is {max}")
140            }
141            Self::MaxDatagramTooSmall { min, actual } => write!(
142                f,
143                "PONK max datagram size {actual} is too small, minimum is {min}"
144            ),
145            Self::TooManyChunks(count) => write!(f, "PONK frame needs too many chunks: {count}"),
146            Self::TooManyMetadata(count) => {
147                write!(f, "PONK path has too many metadata entries: {count}")
148            }
149            Self::TooManyPoints(count) => write!(f, "PONK path has too many points: {count}"),
150            Self::TooManyPaths { max } => write!(f, "PONK frame has more than {max} paths"),
151            Self::TooManyTotalPoints { max } => {
152                write!(f, "PONK frame has more than {max} total points")
153            }
154            Self::FramePayloadTooLarge { max } => {
155                write!(f, "PONK frame payload exceeds {max} bytes")
156            }
157            Self::InvalidCoordinate => write!(f, "PONK frame contains an invalid coordinate"),
158            Self::InvalidChunkHeader => write!(f, "invalid PONK chunk header"),
159            Self::MalformedPayload => write!(f, "malformed PONK payload"),
160        }
161    }
162}
163
164impl std::error::Error for PonkError {}
165
166#[derive(Debug, Clone)]
167struct Header {
168    sender_id: u32,
169    sender_name: String,
170    frame_number: u8,
171    chunk_count: u8,
172    chunk_number: u8,
173    data_crc: u32,
174}
175
176struct Assembly {
177    header: Header,
178    chunks: Vec<Option<Vec<u8>>>,
179    received_bytes: usize,
180    order: u64,
181    updated_at: Instant,
182}
183
184impl Assembly {
185    fn new(header: Header) -> Self {
186        Self {
187            chunks: vec![None; header.chunk_count as usize],
188            header,
189            received_bytes: 0,
190            order: 0,
191            updated_at: Instant::now(),
192        }
193    }
194
195    fn accepts(&self, header: &Header) -> bool {
196        self.header.sender_id == header.sender_id
197            && self.header.frame_number == header.frame_number
198            && self.header.chunk_count == header.chunk_count
199            && self.header.data_crc == header.data_crc
200    }
201
202    fn insert(&mut self, header: &Header, payload: &[u8]) -> Result<(), PonkError> {
203        let Some(slot) = self.chunks.get_mut(header.chunk_number as usize) else {
204            return Err(PonkError::InvalidChunkHeader);
205        };
206        if let Some(existing) = slot.as_ref() {
207            self.received_bytes = self.received_bytes.saturating_sub(existing.len());
208        }
209        self.received_bytes += payload.len();
210        *slot = Some(payload.to_vec());
211        self.updated_at = Instant::now();
212        Ok(())
213    }
214
215    fn complete(&self) -> bool {
216        self.chunks.iter().all(Option::is_some)
217    }
218
219    fn chunk_len(&self, chunk_number: u8) -> usize {
220        self.chunks
221            .get(chunk_number as usize)
222            .and_then(Option::as_ref)
223            .map_or(0, Vec::len)
224    }
225
226    fn into_frame(self, limits: &PonkLimits) -> Result<Option<PonkFrame>, PonkError> {
227        let mut payload = Vec::with_capacity(self.received_bytes);
228        for chunk in self.chunks {
229            let Some(chunk) = chunk else {
230                return Ok(None);
231            };
232            payload.extend(chunk);
233        }
234        decode_payload_checked(&self.header, &payload, limits)
235    }
236}
237
238pub struct PonkAssembler {
239    assemblies: HashMap<(SocketAddr, u32), Assembly>,
240    limits: PonkLimits,
241    next_order: u64,
242}
243
244impl Default for PonkAssembler {
245    fn default() -> Self {
246        Self::new()
247    }
248}
249
250impl PonkAssembler {
251    /// Creates an assembler with finite limits suitable for untrusted UDP.
252    pub fn new() -> Self {
253        Self::with_limits(PonkLimits::default())
254    }
255
256    /// Creates an assembler with the default byte, geometry, and timeout
257    /// limits and the requested concurrent assembly limit.
258    ///
259    /// A zero assembly limit ignores multipart frames without retaining data.
260    pub fn with_max_assemblies(max_assemblies: usize) -> Self {
261        Self::with_limits(PonkLimits {
262            max_assemblies,
263            ..PonkLimits::default()
264        })
265    }
266
267    /// Creates an assembler with caller-provided resource limits.
268    pub fn with_limits(limits: PonkLimits) -> Self {
269        Self {
270            assemblies: HashMap::new(),
271            limits,
272            next_order: 0,
273        }
274    }
275
276    pub fn assembly_count(&self) -> usize {
277        self.assemblies.len()
278    }
279
280    pub fn buffered_bytes(&self) -> usize {
281        self.assemblies
282            .values()
283            .map(|assembly| assembly.received_bytes)
284            .sum()
285    }
286
287    /// Drops incomplete frames that have exceeded `assembly_timeout`.
288    pub fn prune_stale(&mut self) {
289        let now = Instant::now();
290        self.assemblies.retain(|_, assembly| {
291            now.duration_since(assembly.updated_at) <= self.limits.assembly_timeout
292        });
293    }
294
295    fn evict_oldest(&mut self) -> bool {
296        let Some(oldest_key) = self
297            .assemblies
298            .iter()
299            .min_by_key(|(_, assembly)| assembly.order)
300            .map(|(key, _)| *key)
301        else {
302            return false;
303        };
304        self.assemblies.remove(&oldest_key);
305        true
306    }
307
308    pub fn push_datagram(
309        &mut self,
310        datagram: &[u8],
311        peer_addr: SocketAddr,
312    ) -> Result<Option<PonkFrame>, PonkError> {
313        self.prune_stale();
314        let Some((header, payload)) = split_datagram(datagram)? else {
315            return Ok(None);
316        };
317
318        if header.chunk_count == 1 {
319            return decode_payload_checked(&header, payload, &self.limits);
320        }
321        if self.limits.max_assemblies == 0 {
322            return Ok(None);
323        }
324        if payload.len() > self.limits.max_frame_payload_bytes {
325            return Err(PonkError::FramePayloadTooLarge {
326                max: self.limits.max_frame_payload_bytes,
327            });
328        }
329
330        let key = (peer_addr, header.sender_id);
331        if !self
332            .assemblies
333            .get(&key)
334            .is_some_and(|assembly| assembly.accepts(&header))
335        {
336            while self.assemblies.len() >= self.limits.max_assemblies {
337                if !self.evict_oldest() {
338                    return Ok(None);
339                }
340            }
341            self.assemblies.insert(key, Assembly::new(header.clone()));
342            if let Some(assembly) = self.assemblies.get_mut(&key) {
343                assembly.order = self.next_order;
344                self.next_order = self.next_order.wrapping_add(1);
345            }
346        }
347
348        let Some(assembly) = self.assemblies.get(&key) else {
349            return Ok(None);
350        };
351        let replaced_bytes = assembly.chunk_len(header.chunk_number);
352        let frame_bytes = assembly
353            .received_bytes
354            .saturating_sub(replaced_bytes)
355            .saturating_add(payload.len());
356        if frame_bytes > self.limits.max_frame_payload_bytes {
357            self.assemblies.remove(&key);
358            return Err(PonkError::FramePayloadTooLarge {
359                max: self.limits.max_frame_payload_bytes,
360            });
361        }
362        let additional_bytes = payload.len().saturating_sub(replaced_bytes);
363        while self.buffered_bytes().saturating_add(additional_bytes)
364            > self.limits.max_buffered_bytes
365        {
366            if !self.evict_oldest() {
367                return Err(PonkError::FramePayloadTooLarge {
368                    max: self.limits.max_buffered_bytes,
369                });
370            }
371        }
372        let Some(assembly) = self.assemblies.get_mut(&key) else {
373            return Ok(None);
374        };
375        assembly.insert(&header, payload)?;
376        if !assembly.complete() {
377            return Ok(None);
378        }
379
380        self.assemblies
381            .remove(&key)
382            .map(|assembly| assembly.into_frame(&self.limits))
383            .unwrap_or(Ok(None))
384    }
385}
386
387pub fn checksum(data: &[u8]) -> u32 {
388    data.iter()
389        .fold(0u32, |sum, byte| sum.wrapping_add(*byte as u32))
390}
391
392pub fn decode_datagram(datagram: &[u8]) -> Result<Option<PonkFrame>, PonkError> {
393    decode_datagram_with_limits(datagram, &PonkLimits::default())
394}
395
396pub fn decode_datagram_with_limits(
397    datagram: &[u8],
398    limits: &PonkLimits,
399) -> Result<Option<PonkFrame>, PonkError> {
400    let Some((header, payload)) = split_datagram(datagram)? else {
401        return Ok(None);
402    };
403    if header.chunk_count != 1 || header.chunk_number != 0 {
404        return Ok(None);
405    }
406    decode_payload_checked(&header, payload, limits)
407}
408
409pub fn encode_datagrams(
410    frame: &PonkFrame,
411    data_format: DataFormat,
412    max_datagram_len: usize,
413) -> Result<Vec<Vec<u8>>, PonkError> {
414    if max_datagram_len > MAX_DATAGRAM_LEN {
415        return Err(PonkError::DatagramTooLarge {
416            max: MAX_DATAGRAM_LEN,
417            actual: max_datagram_len,
418        });
419    }
420    if max_datagram_len <= HEADER_LEN {
421        return Err(PonkError::MaxDatagramTooSmall {
422            min: HEADER_LEN + 1,
423            actual: max_datagram_len,
424        });
425    }
426
427    let payload = encode_payload(&frame.paths, data_format)?;
428    let max_payload_len = max_datagram_len - HEADER_LEN;
429    let chunk_count = payload.len().max(1).div_ceil(max_payload_len);
430    if chunk_count > u8::MAX as usize {
431        return Err(PonkError::TooManyChunks(chunk_count));
432    }
433
434    let crc = checksum(&payload);
435    let mut datagrams = Vec::with_capacity(chunk_count);
436    if payload.is_empty() {
437        datagrams.push(encode_datagram_header(frame, 1, 0, crc));
438        return Ok(datagrams);
439    }
440
441    for (chunk_number, chunk) in payload.chunks(max_payload_len).enumerate() {
442        let mut datagram =
443            encode_datagram_header(frame, chunk_count as u8, chunk_number as u8, crc);
444        datagram.extend_from_slice(chunk);
445        datagrams.push(datagram);
446    }
447
448    Ok(datagrams)
449}
450
451fn encode_datagram_header(
452    frame: &PonkFrame,
453    chunk_count: u8,
454    chunk_number: u8,
455    data_crc: u32,
456) -> Vec<u8> {
457    let mut datagram = Vec::with_capacity(HEADER_LEN);
458    datagram.extend_from_slice(MAGIC);
459    datagram.push(PROTOCOL_VERSION);
460    datagram.extend_from_slice(&frame.sender_id.to_le_bytes());
461    datagram.extend_from_slice(&fixed_bytes::<SENDER_NAME_LEN>(&frame.sender_name));
462    datagram.push(frame.frame_number);
463    datagram.push(chunk_count);
464    datagram.push(chunk_number);
465    datagram.extend_from_slice(&data_crc.to_le_bytes());
466    datagram
467}
468
469fn encode_payload(paths: &[PonkPath], data_format: DataFormat) -> Result<Vec<u8>, PonkError> {
470    if paths.is_empty() {
471        return Ok(vec![data_format.wire(), 0, 0, 0]);
472    }
473    if paths.len() > DEFAULT_MAX_PATHS {
474        return Err(PonkError::TooManyPaths {
475            max: DEFAULT_MAX_PATHS,
476        });
477    }
478
479    let mut payload_len = 0usize;
480    let mut total_points = 0usize;
481    for path in paths {
482        if path.metadata.len() > u8::MAX as usize {
483            return Err(PonkError::TooManyMetadata(path.metadata.len()));
484        }
485        if path.points.len() > u16::MAX as usize {
486            return Err(PonkError::TooManyPoints(path.points.len()));
487        }
488        total_points =
489            total_points
490                .checked_add(path.points.len())
491                .ok_or(PonkError::TooManyTotalPoints {
492                    max: DEFAULT_MAX_POINTS,
493                })?;
494        if total_points > DEFAULT_MAX_POINTS {
495            return Err(PonkError::TooManyTotalPoints {
496                max: DEFAULT_MAX_POINTS,
497            });
498        }
499        if data_format == DataFormat::XyF32RgbU8
500            && path
501                .points
502                .iter()
503                .any(|point| !point.x.is_finite() || !point.y.is_finite())
504        {
505            return Err(PonkError::InvalidCoordinate);
506        }
507        let path_len = 4usize
508            .checked_add(path.metadata.len().checked_mul(12).ok_or(
509                PonkError::FramePayloadTooLarge {
510                    max: DEFAULT_MAX_FRAME_PAYLOAD_BYTES,
511                },
512            )?)
513            .and_then(|len| {
514                path.points
515                    .len()
516                    .checked_mul(data_format.bytes_per_point())
517                    .and_then(|point_bytes| len.checked_add(point_bytes))
518            })
519            .ok_or(PonkError::FramePayloadTooLarge {
520                max: DEFAULT_MAX_FRAME_PAYLOAD_BYTES,
521            })?;
522        payload_len = payload_len
523            .checked_add(path_len)
524            .ok_or(PonkError::FramePayloadTooLarge {
525                max: DEFAULT_MAX_FRAME_PAYLOAD_BYTES,
526            })?;
527        if payload_len > DEFAULT_MAX_FRAME_PAYLOAD_BYTES {
528            return Err(PonkError::FramePayloadTooLarge {
529                max: DEFAULT_MAX_FRAME_PAYLOAD_BYTES,
530            });
531        }
532    }
533
534    let mut payload = Vec::with_capacity(payload_len);
535    for path in paths {
536        payload.push(data_format.wire());
537        payload.push(path.metadata.len() as u8);
538        for metadata in &path.metadata {
539            payload.extend_from_slice(&fixed_bytes::<METADATA_KEY_LEN>(&metadata.key));
540            payload.extend_from_slice(&metadata.value.to_le_bytes());
541        }
542        payload.extend_from_slice(&(path.points.len() as u16).to_le_bytes());
543
544        for point in &path.points {
545            match data_format {
546                DataFormat::XyRgbU16 => {
547                    payload.extend_from_slice(&coord_to_u16(point.x).to_le_bytes());
548                    payload.extend_from_slice(&coord_to_u16(point.y).to_le_bytes());
549                    payload.extend_from_slice(&((point.rgb[0] as u16) * 257).to_le_bytes());
550                    payload.extend_from_slice(&((point.rgb[1] as u16) * 257).to_le_bytes());
551                    payload.extend_from_slice(&((point.rgb[2] as u16) * 257).to_le_bytes());
552                }
553                DataFormat::XyF32RgbU8 => {
554                    payload.extend_from_slice(&point.x.clamp(-1.0, 1.0).to_le_bytes());
555                    payload.extend_from_slice(&point.y.clamp(-1.0, 1.0).to_le_bytes());
556                    payload.extend_from_slice(&point.rgb);
557                }
558            }
559        }
560    }
561    Ok(payload)
562}
563
564fn split_datagram(datagram: &[u8]) -> Result<Option<(Header, &[u8])>, PonkError> {
565    if datagram.get(0..MAGIC.len()) != Some(MAGIC.as_slice()) {
566        return Ok(None);
567    }
568    if datagram.len() < HEADER_LEN {
569        return Err(PonkError::DatagramTooSmall);
570    }
571    if datagram.len() > MAX_DATAGRAM_LEN {
572        return Err(PonkError::DatagramTooLarge {
573            max: MAX_DATAGRAM_LEN,
574            actual: datagram.len(),
575        });
576    }
577
578    let mut offset = MAGIC.len();
579    let Some(protocol_version) = datagram.get(offset).copied() else {
580        return Ok(None);
581    };
582    offset += 1;
583    if protocol_version != PROTOCOL_VERSION {
584        return Err(PonkError::UnsupportedProtocolVersion(protocol_version));
585    }
586
587    let Some(sender_id) = read_u32_le(datagram, &mut offset) else {
588        return Ok(None);
589    };
590    let Some(sender_name_bytes) = datagram.get(offset..offset + SENDER_NAME_LEN) else {
591        return Ok(None);
592    };
593    let sender_name = parse_fixed_string(sender_name_bytes);
594    offset += SENDER_NAME_LEN;
595    let Some(frame_number) = datagram.get(offset).copied() else {
596        return Ok(None);
597    };
598    offset += 1;
599    let Some(chunk_count) = datagram.get(offset).copied() else {
600        return Ok(None);
601    };
602    offset += 1;
603    let Some(chunk_number) = datagram.get(offset).copied() else {
604        return Ok(None);
605    };
606    offset += 1;
607    let Some(data_crc) = read_u32_le(datagram, &mut offset) else {
608        return Ok(None);
609    };
610
611    if chunk_count == 0 || chunk_number >= chunk_count {
612        return Ok(None);
613    }
614
615    Ok(Some((
616        Header {
617            sender_id,
618            sender_name,
619            frame_number,
620            chunk_count,
621            chunk_number,
622            data_crc,
623        },
624        &datagram[HEADER_LEN..],
625    )))
626}
627
628fn decode_payload_checked(
629    header: &Header,
630    payload: &[u8],
631    limits: &PonkLimits,
632) -> Result<Option<PonkFrame>, PonkError> {
633    if payload.len() > limits.max_frame_payload_bytes {
634        return Err(PonkError::FramePayloadTooLarge {
635            max: limits.max_frame_payload_bytes,
636        });
637    }
638    if checksum(payload) != header.data_crc {
639        return Ok(None);
640    }
641    let paths = decode_payload(payload, limits)?;
642    Ok(Some(PonkFrame {
643        sender_id: header.sender_id,
644        sender_name: header.sender_name.clone(),
645        frame_number: header.frame_number,
646        paths,
647    }))
648}
649
650fn decode_payload(data: &[u8], limits: &PonkLimits) -> Result<Vec<PonkPath>, PonkError> {
651    let mut offset = 0;
652    let mut paths = Vec::new();
653    let mut total_points = 0usize;
654
655    while offset < data.len() {
656        if paths.len() >= limits.max_paths {
657            return Err(PonkError::TooManyPaths {
658                max: limits.max_paths,
659            });
660        }
661        let Some(data_format) = data.get(offset).copied() else {
662            return Err(PonkError::MalformedPayload);
663        };
664        let data_format = DataFormat::from_wire(data_format)?;
665        offset += 1;
666        let Some(metadata_count) = data.get(offset).copied() else {
667            return Err(PonkError::MalformedPayload);
668        };
669        offset += 1;
670
671        let mut metadata = Vec::with_capacity(metadata_count as usize);
672        for _ in 0..metadata_count {
673            let Some(key_bytes) = data.get(offset..offset + METADATA_KEY_LEN) else {
674                return Err(PonkError::MalformedPayload);
675            };
676            let key = parse_fixed_string(key_bytes);
677            offset += METADATA_KEY_LEN;
678            let Some(value) = read_f32_le(data, &mut offset) else {
679                return Err(PonkError::MalformedPayload);
680            };
681            metadata.push(PonkMetadata { key, value });
682        }
683
684        let Some(point_count) = read_u16_le(data, &mut offset) else {
685            return Err(PonkError::MalformedPayload);
686        };
687        let point_count = point_count as usize;
688        total_points =
689            total_points
690                .checked_add(point_count)
691                .ok_or(PonkError::TooManyTotalPoints {
692                    max: limits.max_points,
693                })?;
694        if total_points > limits.max_points {
695            return Err(PonkError::TooManyTotalPoints {
696                max: limits.max_points,
697            });
698        }
699        let Some(byte_count) = point_count.checked_mul(data_format.bytes_per_point()) else {
700            return Err(PonkError::MalformedPayload);
701        };
702        let Some(end_offset) = offset.checked_add(byte_count) else {
703            return Err(PonkError::MalformedPayload);
704        };
705        if data.len() < end_offset {
706            return Err(PonkError::MalformedPayload);
707        }
708
709        let mut points = Vec::with_capacity(point_count);
710        for _ in 0..point_count {
711            points.push(match data_format {
712                DataFormat::XyRgbU16 => {
713                    let Some(x) = read_u16_le(data, &mut offset) else {
714                        return Err(PonkError::MalformedPayload);
715                    };
716                    let Some(y) = read_u16_le(data, &mut offset) else {
717                        return Err(PonkError::MalformedPayload);
718                    };
719                    let Some(r) = read_u16_le(data, &mut offset) else {
720                        return Err(PonkError::MalformedPayload);
721                    };
722                    let Some(g) = read_u16_le(data, &mut offset) else {
723                        return Err(PonkError::MalformedPayload);
724                    };
725                    let Some(b) = read_u16_le(data, &mut offset) else {
726                        return Err(PonkError::MalformedPayload);
727                    };
728                    PonkPoint {
729                        x: u16_to_coord(x),
730                        y: u16_to_coord(y),
731                        rgb: [(r >> 8) as u8, (g >> 8) as u8, (b >> 8) as u8],
732                    }
733                }
734                DataFormat::XyF32RgbU8 => {
735                    let Some(x) = read_f32_le(data, &mut offset) else {
736                        return Err(PonkError::MalformedPayload);
737                    };
738                    let Some(y) = read_f32_le(data, &mut offset) else {
739                        return Err(PonkError::MalformedPayload);
740                    };
741                    let Some(rgb) = data.get(offset..offset + 3) else {
742                        return Err(PonkError::MalformedPayload);
743                    };
744                    offset += 3;
745                    if !x.is_finite() || !y.is_finite() {
746                        return Err(PonkError::InvalidCoordinate);
747                    }
748                    PonkPoint {
749                        x: x.clamp(-1.0, 1.0),
750                        y: y.clamp(-1.0, 1.0),
751                        rgb: [rgb[0], rgb[1], rgb[2]],
752                    }
753                }
754            });
755        }
756
757        paths.push(PonkPath { metadata, points });
758    }
759
760    Ok(paths)
761}
762
763fn fixed_bytes<const N: usize>(value: &str) -> [u8; N] {
764    let mut bytes = [0u8; N];
765    let mut len = value.len().min(N);
766    while !value.is_char_boundary(len) {
767        len -= 1;
768    }
769    bytes[..len].copy_from_slice(&value.as_bytes()[..len]);
770    bytes
771}
772
773fn parse_fixed_string(bytes: &[u8]) -> String {
774    let len = bytes.iter().position(|b| *b == 0).unwrap_or(bytes.len());
775    String::from_utf8_lossy(&bytes[..len]).to_string()
776}
777
778fn read_u16_le(data: &[u8], offset: &mut usize) -> Option<u16> {
779    let bytes = data.get(*offset..*offset + 2)?;
780    *offset += 2;
781    Some(u16::from_le_bytes([bytes[0], bytes[1]]))
782}
783
784fn read_u32_le(data: &[u8], offset: &mut usize) -> Option<u32> {
785    let bytes = data.get(*offset..*offset + 4)?;
786    *offset += 4;
787    Some(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
788}
789
790fn read_f32_le(data: &[u8], offset: &mut usize) -> Option<f32> {
791    read_u32_le(data, offset).map(f32::from_bits)
792}
793
794fn coord_to_u16(value: f32) -> u16 {
795    (((value.clamp(-1.0, 1.0) + 1.0) * 0.5) * u16::MAX as f32).round() as u16
796}
797
798fn u16_to_coord(value: u16) -> f32 {
799    -1.0 + 2.0 * (value as f32 / u16::MAX as f32)
800}
801
802#[cfg(test)]
803mod tests {
804    use std::net::{Ipv4Addr, SocketAddr};
805
806    use super::*;
807
808    // Literal vectors follow the canonical C++ sample's packed structs and
809    // little-endian push helpers at madmappersoftware/Ponk@2c166392.
810    const CANONICAL_F32_DATAGRAM: &[u8] = &[
811        0x50, 0x4f, 0x4e, 0x4b, 0x2d, 0x55, 0x44, 0x50, 0x00, 0x04, 0x03, 0x02, 0x01, 0x53, 0x00,
812        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
813        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
814        0x05, 0x01, 0x00, 0x3a, 0x05, 0x00, 0x00, 0x01, 0x01, 0x50, 0x41, 0x54, 0x48, 0x4e, 0x55,
815        0x4d, 0x42, 0x00, 0x00, 0x80, 0x3f, 0x01, 0x00, 0x00, 0x00, 0x00, 0x3f, 0x00, 0x00, 0x80,
816        0xbe, 0x12, 0x34, 0x56,
817    ];
818    const CANONICAL_U16_DATAGRAM: &[u8] = &[
819        0x50, 0x4f, 0x4e, 0x4b, 0x2d, 0x55, 0x44, 0x50, 0x00, 0x04, 0x03, 0x02, 0x01, 0x53, 0x00,
820        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
821        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
822        0x05, 0x01, 0x00, 0x7d, 0x05, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0xff, 0xff,
823        0x80, 0x80, 0x40, 0x40, 0xff, 0xff,
824    ];
825
826    fn sample_frame() -> PonkFrame {
827        PonkFrame {
828            sender_id: 42,
829            sender_name: "test sender".to_string(),
830            frame_number: 7,
831            paths: vec![PonkPath {
832                metadata: vec![PonkMetadata {
833                    key: "PATHNUMB".to_string(),
834                    value: 3.0,
835                }],
836                points: vec![
837                    PonkPoint {
838                        x: -0.5,
839                        y: 0.25,
840                        rgb: [255, 0, 64],
841                    },
842                    PonkPoint {
843                        x: 0.5,
844                        y: -0.25,
845                        rgb: [0, 255, 128],
846                    },
847                ],
848            }],
849        }
850    }
851
852    #[test]
853    fn canonical_f32_vector_decodes_and_encodes_exactly() {
854        let decoded = decode_datagram(CANONICAL_F32_DATAGRAM).unwrap().unwrap();
855        assert_eq!(decoded.sender_id, 0x0102_0304);
856        assert_eq!(decoded.sender_name, "S");
857        assert_eq!(decoded.frame_number, 5);
858        assert_eq!(decoded.paths[0].metadata[0].key, "PATHNUMB");
859        assert_eq!(decoded.paths[0].metadata[0].value, 1.0);
860        assert_eq!(decoded.paths[0].points[0].x, 0.5);
861        assert_eq!(decoded.paths[0].points[0].y, -0.25);
862        assert_eq!(decoded.paths[0].points[0].rgb, [0x12, 0x34, 0x56]);
863
864        assert_eq!(
865            encode_datagrams(&decoded, DataFormat::XyF32RgbU8, 1_472).unwrap(),
866            [CANONICAL_F32_DATAGRAM]
867        );
868    }
869
870    #[test]
871    fn canonical_u16_vector_decodes_and_encodes_exactly() {
872        let decoded = decode_datagram(CANONICAL_U16_DATAGRAM).unwrap().unwrap();
873        assert_eq!(decoded.sender_id, 0x0102_0304);
874        assert_eq!(decoded.sender_name, "S");
875        assert_eq!(decoded.frame_number, 5);
876        assert_eq!(decoded.paths[0].points[0].x, -1.0);
877        assert_eq!(decoded.paths[0].points[0].y, 1.0);
878        assert_eq!(decoded.paths[0].points[0].rgb, [0x80, 0x40, 0xff]);
879
880        assert_eq!(
881            encode_datagrams(&decoded, DataFormat::XyRgbU16, 1_472).unwrap(),
882            [CANONICAL_U16_DATAGRAM]
883        );
884    }
885
886    #[test]
887    fn empty_frame_encodes_as_explicit_zero_point_path() {
888        let frame = PonkFrame {
889            paths: Vec::new(),
890            ..sample_frame()
891        };
892
893        let datagram = encode_datagrams(&frame, DataFormat::XyF32RgbU8, 1_472)
894            .unwrap()
895            .remove(0);
896
897        assert!(datagram.len() > HEADER_LEN);
898        let decoded = decode_datagram(&datagram).unwrap().unwrap();
899        assert_eq!(decoded.paths.len(), 1);
900        assert!(decoded.paths[0].points.is_empty());
901    }
902
903    #[test]
904    fn encoded_single_datagram_decodes_to_frame() {
905        let datagrams = encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, 1200).unwrap();
906        assert_eq!(datagrams.len(), 1);
907
908        let decoded = decode_datagram(&datagrams[0]).unwrap().unwrap();
909
910        assert_eq!(decoded.sender_id, 42);
911        assert_eq!(decoded.sender_name, "test sender");
912        assert_eq!(decoded.frame_number, 7);
913        assert_eq!(decoded.paths.len(), 1);
914        assert_eq!(decoded.paths[0].metadata[0].key, "PATHNUMB");
915        assert_eq!(decoded.paths[0].metadata[0].value, 3.0);
916        assert_eq!(decoded.paths[0].points, sample_frame().paths[0].points);
917    }
918
919    #[test]
920    fn encoded_chunked_datagrams_reassemble_to_frame() {
921        let mut frame = sample_frame();
922        frame.paths[0].points = (0..50)
923            .map(|i| PonkPoint {
924                x: -1.0 + i as f32 * 0.04,
925                y: 0.25,
926                rgb: [i as u8, 255, 128],
927            })
928            .collect();
929
930        let datagrams = encode_datagrams(&frame, DataFormat::XyF32RgbU8, HEADER_LEN + 80).unwrap();
931        assert!(datagrams.len() > 1);
932
933        let peer = SocketAddr::from((Ipv4Addr::LOCALHOST, 5583));
934        let mut assembler = PonkAssembler::new();
935        let mut decoded = None;
936        for datagram in datagrams.iter().rev() {
937            decoded = assembler.push_datagram(datagram, peer).unwrap().or(decoded);
938        }
939
940        let decoded = decoded.expect("all chunks should reassemble");
941        assert_eq!(decoded.sender_id, frame.sender_id);
942        assert_eq!(decoded.sender_name, frame.sender_name);
943        assert_eq!(decoded.paths[0].points, frame.paths[0].points);
944    }
945
946    #[test]
947    fn malformed_datagram_classes_have_stable_public_results() {
948        let truncated_header = b"PONK-UDP\0";
949        assert_eq!(
950            decode_datagram(truncated_header),
951            Err(PonkError::DatagramTooSmall)
952        );
953
954        let mut unsupported_version = CANONICAL_F32_DATAGRAM.to_vec();
955        unsupported_version[MAGIC.len()] = 1;
956        assert_eq!(
957            decode_datagram(&unsupported_version),
958            Err(PonkError::UnsupportedProtocolVersion(1))
959        );
960
961        let mut invalid_chunk_numbering = CANONICAL_F32_DATAGRAM.to_vec();
962        let chunk_count_offset = HEADER_LEN - 6;
963        invalid_chunk_numbering[chunk_count_offset] = 0;
964        assert_eq!(decode_datagram(&invalid_chunk_numbering), Ok(None));
965
966        let mut unsupported_format = CANONICAL_F32_DATAGRAM.to_vec();
967        unsupported_format[HEADER_LEN] = 99;
968        let updated_checksum = checksum(&unsupported_format[HEADER_LEN..]);
969        unsupported_format[HEADER_LEN - 4..HEADER_LEN]
970            .copy_from_slice(&updated_checksum.to_le_bytes());
971        assert_eq!(
972            decode_datagram(&unsupported_format),
973            Err(PonkError::UnsupportedDataFormat(99))
974        );
975
976        let mut truncated_payload =
977            CANONICAL_F32_DATAGRAM[..CANONICAL_F32_DATAGRAM.len() - 1].to_vec();
978        let updated_checksum = checksum(&truncated_payload[HEADER_LEN..]);
979        truncated_payload[HEADER_LEN - 4..HEADER_LEN]
980            .copy_from_slice(&updated_checksum.to_le_bytes());
981        assert_eq!(
982            decode_datagram(&truncated_payload),
983            Err(PonkError::MalformedPayload)
984        );
985    }
986
987    #[test]
988    fn bad_checksum_is_ignored() {
989        let mut datagram = encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, 1200)
990            .unwrap()
991            .remove(0);
992        datagram[HEADER_LEN - 4] = datagram[HEADER_LEN - 4].wrapping_add(1);
993
994        assert!(decode_datagram(&datagram).unwrap().is_none());
995    }
996
997    #[test]
998    fn protocol_accepts_255_chunks_and_rejects_256() {
999        let point = PonkPoint {
1000            x: 0.0,
1001            y: 0.0,
1002            rgb: [0, 0, 0],
1003        };
1004        let metadata = PonkMetadata {
1005            key: "PATHNUMB".to_string(),
1006            value: 0.0,
1007        };
1008        let frame_255 = PonkFrame {
1009            paths: vec![PonkPath {
1010                metadata: vec![metadata.clone(); 9],
1011                points: vec![point.clone(); 13],
1012            }],
1013            ..sample_frame()
1014        };
1015        assert_eq!(
1016            encode_datagrams(&frame_255, DataFormat::XyF32RgbU8, HEADER_LEN + 1)
1017                .unwrap()
1018                .len(),
1019            255
1020        );
1021
1022        let frame_256 = PonkFrame {
1023            paths: vec![PonkPath {
1024                metadata: vec![metadata; 10],
1025                points: vec![point; 12],
1026            }],
1027            ..sample_frame()
1028        };
1029        assert_eq!(
1030            encode_datagrams(&frame_256, DataFormat::XyF32RgbU8, HEADER_LEN + 1),
1031            Err(PonkError::TooManyChunks(256))
1032        );
1033    }
1034
1035    #[test]
1036    fn zero_assembly_limit_never_retains_multipart_frames() {
1037        let datagrams =
1038            encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, HEADER_LEN + 8).unwrap();
1039        assert!(datagrams.len() > 1);
1040        let peer = SocketAddr::from((Ipv4Addr::LOCALHOST, 5583));
1041        let mut assembler = PonkAssembler::with_max_assemblies(0);
1042
1043        assert!(
1044            assembler
1045                .push_datagram(&datagrams[0], peer)
1046                .unwrap()
1047                .is_none()
1048        );
1049        assert_eq!(assembler.assembly_count(), 0);
1050        assert_eq!(assembler.buffered_bytes(), 0);
1051    }
1052
1053    #[test]
1054    fn assembly_payload_limit_rejects_frame_before_retaining_chunk() {
1055        let datagrams =
1056            encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, HEADER_LEN + 16).unwrap();
1057        let limits = PonkLimits {
1058            max_frame_payload_bytes: 8,
1059            ..PonkLimits::default()
1060        };
1061        let peer = SocketAddr::from((Ipv4Addr::LOCALHOST, 5583));
1062        let mut assembler = PonkAssembler::with_limits(limits);
1063
1064        assert_eq!(
1065            assembler.push_datagram(&datagrams[0], peer),
1066            Err(PonkError::FramePayloadTooLarge { max: 8 })
1067        );
1068        assert_eq!(assembler.assembly_count(), 0);
1069        assert_eq!(assembler.buffered_bytes(), 0);
1070    }
1071
1072    #[test]
1073    fn aggregate_buffer_limit_evicts_oldest_incomplete_frame() {
1074        let datagrams =
1075            encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, HEADER_LEN + 8).unwrap();
1076        let limits = PonkLimits {
1077            max_buffered_bytes: 12,
1078            ..PonkLimits::default()
1079        };
1080        let first_peer = SocketAddr::from((Ipv4Addr::LOCALHOST, 5583));
1081        let second_peer = SocketAddr::from((Ipv4Addr::LOCALHOST, 5584));
1082        let mut assembler = PonkAssembler::with_limits(limits);
1083
1084        assembler.push_datagram(&datagrams[0], first_peer).unwrap();
1085        assembler.push_datagram(&datagrams[0], second_peer).unwrap();
1086
1087        assert_eq!(assembler.assembly_count(), 1);
1088        assert!(assembler.buffered_bytes() <= limits.max_buffered_bytes);
1089    }
1090
1091    #[test]
1092    fn stale_incomplete_assemblies_are_pruned() {
1093        let datagrams =
1094            encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, HEADER_LEN + 8).unwrap();
1095        let limits = PonkLimits {
1096            assembly_timeout: std::time::Duration::ZERO,
1097            ..PonkLimits::default()
1098        };
1099        let peer = SocketAddr::from((Ipv4Addr::LOCALHOST, 5583));
1100        let mut assembler = PonkAssembler::with_limits(limits);
1101
1102        assembler.push_datagram(&datagrams[0], peer).unwrap();
1103        assert_eq!(assembler.assembly_count(), 1);
1104        assembler.prune_stale();
1105        assert_eq!(assembler.assembly_count(), 0);
1106        assert_eq!(assembler.buffered_bytes(), 0);
1107    }
1108
1109    #[test]
1110    fn oversized_decode_input_is_rejected() {
1111        let mut datagram = encode_datagrams(&sample_frame(), DataFormat::XyF32RgbU8, 1200)
1112            .unwrap()
1113            .remove(0);
1114        datagram.resize(MAX_DATAGRAM_LEN + 1, 0);
1115
1116        assert_eq!(
1117            decode_datagram(&datagram),
1118            Err(PonkError::DatagramTooLarge {
1119                max: MAX_DATAGRAM_LEN,
1120                actual: MAX_DATAGRAM_LEN + 1,
1121            })
1122        );
1123    }
1124
1125    #[test]
1126    fn checksum_wraps_instead_of_panicking() {
1127        let data = vec![u8::MAX; (u32::MAX as usize / u8::MAX as usize) + 2];
1128        let expected = data
1129            .iter()
1130            .fold(0u32, |sum, byte| sum.wrapping_add(*byte as u32));
1131
1132        assert_eq!(checksum(&data), expected);
1133    }
1134
1135    #[test]
1136    fn decode_limits_reject_path_amplification() {
1137        let frame = PonkFrame {
1138            paths: vec![
1139                PonkPath {
1140                    metadata: Vec::new(),
1141                    points: Vec::new(),
1142                },
1143                PonkPath {
1144                    metadata: Vec::new(),
1145                    points: Vec::new(),
1146                },
1147            ],
1148            ..sample_frame()
1149        };
1150        let datagram = encode_datagrams(&frame, DataFormat::XyF32RgbU8, 1200)
1151            .unwrap()
1152            .remove(0);
1153        let limits = PonkLimits {
1154            max_paths: 1,
1155            ..PonkLimits::default()
1156        };
1157
1158        assert_eq!(
1159            decode_datagram_with_limits(&datagram, &limits),
1160            Err(PonkError::TooManyPaths { max: 1 })
1161        );
1162    }
1163
1164    #[test]
1165    fn float_coordinates_are_finite_and_bounded_at_wire_boundaries() {
1166        let mut frame = sample_frame();
1167        frame.paths[0].points[0].x = f32::NAN;
1168        assert_eq!(
1169            encode_datagrams(&frame, DataFormat::XyF32RgbU8, 1200),
1170            Err(PonkError::InvalidCoordinate)
1171        );
1172
1173        let mut invalid_datagram = CANONICAL_F32_DATAGRAM.to_vec();
1174        let x_offset = HEADER_LEN + 16;
1175        invalid_datagram[x_offset..x_offset + 4].copy_from_slice(&f32::NAN.to_le_bytes());
1176        let crc = checksum(&invalid_datagram[HEADER_LEN..]);
1177        invalid_datagram[HEADER_LEN - 4..HEADER_LEN].copy_from_slice(&crc.to_le_bytes());
1178        assert_eq!(
1179            decode_datagram(&invalid_datagram),
1180            Err(PonkError::InvalidCoordinate)
1181        );
1182
1183        frame.paths[0].points[0].x = 2.0;
1184        let datagram = encode_datagrams(&frame, DataFormat::XyF32RgbU8, 1200)
1185            .unwrap()
1186            .remove(0);
1187        let decoded = decode_datagram(&datagram).unwrap().unwrap();
1188        assert_eq!(decoded.paths[0].points[0].x, 1.0);
1189    }
1190
1191    #[test]
1192    fn encoder_rejects_aggregate_path_amplification() {
1193        let frame = PonkFrame {
1194            paths: (0..=DEFAULT_MAX_PATHS)
1195                .map(|_| PonkPath {
1196                    metadata: Vec::new(),
1197                    points: Vec::new(),
1198                })
1199                .collect(),
1200            ..sample_frame()
1201        };
1202
1203        assert_eq!(
1204            encode_datagrams(&frame, DataFormat::XyF32RgbU8, 1_472),
1205            Err(PonkError::TooManyPaths {
1206                max: DEFAULT_MAX_PATHS
1207            })
1208        );
1209    }
1210
1211    #[test]
1212    fn fixed_utf8_fields_do_not_split_code_points() {
1213        let mut frame = sample_frame();
1214        frame.sender_name = format!("{}é", "a".repeat(SENDER_NAME_LEN - 1));
1215        frame.paths[0].metadata[0].key = "abcdefgé".to_string();
1216
1217        let datagram = encode_datagrams(&frame, DataFormat::XyF32RgbU8, 1200)
1218            .unwrap()
1219            .remove(0);
1220        let decoded = decode_datagram(&datagram).unwrap().unwrap();
1221
1222        assert_eq!(decoded.sender_name, "a".repeat(SENDER_NAME_LEN - 1));
1223        assert_eq!(decoded.paths[0].metadata[0].key, "abcdefg");
1224        assert!(!decoded.sender_name.contains('\u{fffd}'));
1225        assert!(!decoded.paths[0].metadata[0].key.contains('\u{fffd}'));
1226    }
1227}
ponk_protocol/lib.rs

ponk_protocol/
lib.rs
