diff --git a/CHANGELOG.md b/CHANGELOG.md index 32b069ed..ed285188 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,6 +66,8 @@ signals a backwards-compatible change. worker's seed and the aggregator's from one seed. The unseeded constructors draw from the thread generator. The seed is not serialized, and the ASAPv1 payload is unchanged. +- **`peek_kind_id(&[u8])`** returns an ASAPv1 envelope's `kind_id` after + checking its magic, version, and framing, without decoding the sketch. ### Changed diff --git a/docs/asapv1_wire_format.md b/docs/asapv1_wire_format.md index f85b81d8..3532dab5 100644 --- a/docs/asapv1_wire_format.md +++ b/docs/asapv1_wire_format.md @@ -249,7 +249,7 @@ This registry is the master list of algorithms still to design payloads for, and > Implementation note: the shared envelope module (`src/message_pack_format/envelope.rs`) owns rules 1-3 and the byte framing (`encode` / `split`); it is sketch-agnostic and does not know the registry. > Rule 4 (and metadata/kind_id validation) happens in each sketch's decoder, which checks the `kind_id` against the ones it owns. -> `peek_kind_id` checks only rules 1-3, then returns the `kind_id` bytes; otherwise it fails. +> `peek_kind_id` checks rules 2-3 and the length checks of rule 1, not that each block is filled exactly, then returns the `kind_id` bytes; otherwise it fails. --- diff --git a/src/lib.rs b/src/lib.rs index b062b8b9..5119bb3c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,6 +41,7 @@ pub mod __private { pub use common::*; pub use message_pack_format::MessagePackCodec; +pub use message_pack_format::envelope::peek_kind_id; pub use message_pack_format::portable::countminsketch::{CountMinSketch, CountMinSketchDelta}; pub use message_pack_format::portable::countminsketch_topk::{CmsHeapItem, CountMinSketchWithHeap}; pub use message_pack_format::portable::countsketch::{ diff --git a/src/message_pack_format/envelope.rs b/src/message_pack_format/envelope.rs index 4fb8c00d..ac1891bf 100644 --- a/src/message_pack_format/envelope.rs +++ b/src/message_pack_format/envelope.rs @@ -111,6 +111,13 @@ pub(crate) fn split(bytes: &[u8]) -> Result, String> { )) } +/// Returns the `kind_id` of an ASAPv1 envelope after checking its magic, version, +/// and framing; the metadata and payload blocks are not parsed. +pub fn peek_kind_id(bytes: &[u8]) -> Result<&[u8], RmpDecodeError> { + let (kind_id, _, _) = split(bytes).map_err(RmpDecodeError::Uncategorized)?; + Ok(kind_id) +} + /// Decodes a metadata or payload block as one msgpack value that fills the /// block exactly; a byte left after the value is an error. pub(crate) fn decode_block(block: &[u8]) -> Result { diff --git a/tests/asapv1_golden.rs b/tests/asapv1_golden.rs index 30b41651..933d112b 100644 --- a/tests/asapv1_golden.rs +++ b/tests/asapv1_golden.rs @@ -10,7 +10,7 @@ use asap_sketchlib::{ Classic, Count, CountMin, ErtlMLE, FastPath, HllSketch, HllVariant, HyperLogLogHIPP12, - HyperLogLogP12, KLL, MessagePackCodec, RegularPath, Vector2D, + HyperLogLogP12, KLL, MessagePackCodec, RegularPath, Vector2D, peek_kind_id, }; fn decode_hex(s: &str) -> Vec { @@ -287,3 +287,31 @@ fn kll_i64_k200_matches_golden() { assert_eq!(decoded.quantile(0.0), 1.0); assert_eq!(decoded.quantile(1.0), 50.0); } + +// --------------------------------------------------------------------------- +// peek_kind_id +// --------------------------------------------------------------------------- + +#[test] +fn peek_kind_id_checks_magic_version_and_framing() { + let bytes = decode_hex(GOLDEN_CMS_I64); + + let trailing = [&bytes[..], &[0xc0]].concat(); + assert_eq!( + peek_kind_id(&trailing).expect("trailing bytes ignored"), + [0x02, 0x00] + ); + + let mut bad_magic = bytes.clone(); + bad_magic[0] = b'X'; + assert!(peek_kind_id(&bad_magic).is_err(), "bad magic"); + + let mut bad_version = bytes.clone(); + bad_version[6] = 0x02; + assert!(peek_kind_id(&bad_version).is_err(), "unknown version"); + + // Cut inside the length fields, then inside the payload. + for len in [17, bytes.len() - 1] { + assert!(peek_kind_id(&bytes[..len]).is_err(), "truncated to {len}"); + } +}