Skip to main content

dig_message/
registry.rs

1//! The extensible message-type registry (SPEC §4) — the runtime seam through which every downstream
2//! subsystem (dig-chat, dig-email, dig-video, peer-RPC, IPC) plugs its own message types into the one
3//! base protocol WITHOUT dig-message depending on any of them.
4//!
5//! Three pieces make the seam:
6//! - [`MessageBand`] + [`MessageType::band`] — the normative id allocation (each subsystem owns a
7//!   256-wide band), so a reader can classify any id it sees.
8//! - [`MessageKind`] — the compile-time contract a type declares: its reserved id + its typed,
9//!   byte-deterministic [`Streamable`] payload.
10//! - [`MessageRegistry`] — the runtime table that maps a [`MessageType`] to a decode-and-handle
11//!   closure, populated additively.
12//!
13//! Forward compatibility (SPEC §4, the §5.1 additive-only spirit) is the load-bearing property: a
14//! `message_type` the registry has never seen MUST fail CLEANLY — [`MessageError::UnsupportedType`] for
15//! a request/stream shape (so the caller replies with an error), a silent drop for a one-shot/response
16//! shape — and MUST NEVER panic. An old reader therefore keeps working when newer senders introduce
17//! new types.
18
19use std::collections::HashMap;
20
21use chia_traits::Streamable;
22
23use crate::envelope::{InteractionShape, MessageType};
24use crate::error::{MessageError, Result};
25
26/// The base of the core band (handshake / ack / error / keepalive) (SPEC §4).
27pub const BAND_CORE: u32 = 0x0000_0000;
28/// The base of the peer-RPC band (peer-to-peer request/response) (SPEC §4).
29pub const BAND_PEER_RPC: u32 = 0x0000_0100;
30/// The base of the dig-chat band (SPEC §4, #768).
31pub const BAND_DIG_CHAT: u32 = 0x0000_0200;
32/// The base of the dig-email band (SPEC §4, #794).
33pub const BAND_DIG_EMAIL: u32 = 0x0000_0300;
34/// The base of the dig-video-chat signaling band (SPEC §4, #795).
35pub const BAND_DIG_VIDEO: u32 = 0x0000_0400;
36/// The base of the presence / directed data-request band (SPEC §4).
37pub const BAND_PRESENCE: u32 = 0x0000_0500;
38/// The base of the dig-ipc-protocol band (authenticated local dig-app ↔ dig-node IPC) (SPEC §4).
39pub const BAND_IPC: u32 = 0x0000_0600;
40/// The base of the experimental / vendor band — never shipped as a canonical type (SPEC §4).
41pub const BAND_EXPERIMENTAL: u32 = 0x1000_0000;
42
43/// A subsystem's reserved id band (SPEC §4). Each named band is owned by one subsystem and allocated
44/// additively within it; ids that fall in no allocated subsystem band classify as [`MessageBand::Reserved`].
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub enum MessageBand {
47    /// Handshake, ack, error, keepalive (`0x0000_0000..=0x0000_00FF`).
48    Core,
49    /// Peer-to-peer request/response (`0x0000_0100..=0x0000_01FF`).
50    PeerRpc,
51    /// dig-chat (`0x0000_0200..=0x0000_02FF`).
52    DigChat,
53    /// dig-email (`0x0000_0300..=0x0000_03FF`).
54    DigEmail,
55    /// dig-video-chat signaling (`0x0000_0400..=0x0000_04FF`).
56    DigVideo,
57    /// Presence / directed data-request (`0x0000_0500..=0x0000_05FF`).
58    Presence,
59    /// dig-ipc-protocol (`0x0000_0600..=0x0000_06FF`).
60    Ipc,
61    /// Experimental / vendor, never canonical (`>= 0x1000_0000`).
62    Experimental,
63    /// A currently-unallocated id, reserved for a future subsystem band (SPEC §4).
64    Reserved,
65}
66
67impl MessageType {
68    /// Core: connection handshake (SPEC §4 core band).
69    pub const CORE_HANDSHAKE: MessageType = MessageType(BAND_CORE);
70    /// Core: generic acknowledgement (SPEC §4 core band).
71    pub const CORE_ACK: MessageType = MessageType(BAND_CORE + 1);
72    /// Core: protocol-level error report (SPEC §4 core band).
73    pub const CORE_ERROR: MessageType = MessageType(BAND_CORE + 2);
74    /// Core: liveness keepalive (SPEC §4 core band).
75    pub const CORE_KEEPALIVE: MessageType = MessageType(BAND_CORE + 3);
76
77    /// Classify this id into its reserved [`MessageBand`] (SPEC §4). Total — every `u32` maps to a
78    /// band, so a reader can always classify an id it does not otherwise recognize.
79    #[must_use]
80    pub fn band(self) -> MessageBand {
81        match self.0 {
82            0x0000_0000..=0x0000_00FF => MessageBand::Core,
83            0x0000_0100..=0x0000_01FF => MessageBand::PeerRpc,
84            0x0000_0200..=0x0000_02FF => MessageBand::DigChat,
85            0x0000_0300..=0x0000_03FF => MessageBand::DigEmail,
86            0x0000_0400..=0x0000_04FF => MessageBand::DigVideo,
87            0x0000_0500..=0x0000_05FF => MessageBand::Presence,
88            0x0000_0600..=0x0000_06FF => MessageBand::Ipc,
89            0x1000_0000..=0xFFFF_FFFF => MessageBand::Experimental,
90            _ => MessageBand::Reserved,
91        }
92    }
93}
94
95/// The compile-time contract a message type declares to plug into the registry (SPEC §4). A downstream
96/// crate implements it for each of its payload types; dig-message never depends on that crate.
97///
98/// The [`Payload`](MessageKind::Payload) is a Chia-[`Streamable`] type so its bytes are byte-
99/// deterministic across the Rust and wasm/JS targets (SPEC §1); [`TYPE_ID`](MessageKind::TYPE_ID) is
100/// the reserved id from the owning subsystem's band ([`MessageBand`]).
101pub trait MessageKind {
102    /// The reserved [`MessageType`] id for this kind (from the subsystem's band, SPEC §4).
103    const TYPE_ID: MessageType;
104
105    /// The typed, byte-deterministic payload carried by this message type (SPEC §1, §4).
106    type Payload: Streamable;
107}
108
109/// The outcome of dispatching a decoded payload (SPEC §4). Both variants are success — an unknown
110/// one-shot dropping is the intended forward-compat behavior, not an error.
111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub enum Dispatch {
113    /// A registered handler decoded and processed the payload.
114    Handled,
115    /// The type was unknown and the shape was one-shot/response, so the message was silently dropped
116    /// (SPEC §4). No handler ran and no error is surfaced.
117    Dropped,
118}
119
120/// A registered decode-and-handle closure: it decodes the payload bytes into the kind's Streamable
121/// payload and processes it, propagating any decode/handler failure as a [`MessageError`].
122type Handler = Box<dyn Fn(&[u8]) -> Result<()> + Send + Sync>;
123
124/// The runtime message-type table (SPEC §4). Maps a [`MessageType`] to its decode-and-handle closure,
125/// populated additively by each subsystem at startup.
126///
127/// It is deliberately additive-only: re-registering an already-present id is refused with
128/// [`MessageError::DuplicateType`] rather than silently overwriting, upholding the SPEC §4 rule that an
129/// id, once assigned, is never renumbered or repurposed. Registering a NEW id never disturbs the
130/// handlers already present.
131#[derive(Default)]
132pub struct MessageRegistry {
133    handlers: HashMap<MessageType, Handler>,
134}
135
136impl MessageRegistry {
137    /// An empty registry.
138    #[must_use]
139    pub fn new() -> Self {
140        Self::default()
141    }
142
143    /// Register a [`MessageKind`] with the closure that processes its decoded payload.
144    ///
145    /// The stored handler decodes the on-wire bytes into `K::Payload` (Streamable) before invoking
146    /// `handler`, so callers work with the typed payload, never raw bytes.
147    ///
148    /// # Errors
149    /// [`MessageError::DuplicateType`] if the kind's [`MessageKind::TYPE_ID`] is already registered
150    /// (SPEC §4 additive-only — never silently repurpose an id).
151    pub fn register<K, F>(&mut self, handler: F) -> Result<()>
152    where
153        K: MessageKind,
154        F: Fn(K::Payload) -> Result<()> + Send + Sync + 'static,
155    {
156        let type_id = K::TYPE_ID;
157        if self.handlers.contains_key(&type_id) {
158            return Err(MessageError::DuplicateType(type_id.0));
159        }
160        let handler: Handler = Box::new(move |bytes: &[u8]| {
161            let payload =
162                K::Payload::from_bytes(bytes).map_err(|e| MessageError::Codec(e.to_string()))?;
163            handler(payload)
164        });
165        self.handlers.insert(type_id, handler);
166        Ok(())
167    }
168
169    /// Whether a [`MessageType`] has a registered handler.
170    #[must_use]
171    pub fn contains(&self, message_type: MessageType) -> bool {
172        self.handlers.contains_key(&message_type)
173    }
174
175    /// The number of registered types.
176    #[must_use]
177    pub fn len(&self) -> usize {
178        self.handlers.len()
179    }
180
181    /// Whether no types are registered.
182    #[must_use]
183    pub fn is_empty(&self) -> bool {
184        self.handlers.is_empty()
185    }
186
187    /// Route a decoded payload to its registered handler, applying the SPEC §4 unknown-type rule.
188    ///
189    /// `payload` is the already-opened, decompressed type-payload bytes (the seal/decompression are
190    /// WU2/§1.1 concerns upstream of dispatch). `shape` decides how an unknown type fails.
191    ///
192    /// # Errors
193    /// - [`MessageError::UnsupportedType`] when the type is unregistered AND the shape is a request or
194    ///   stream frame (so the caller can reply UNSUPPORTED_TYPE).
195    /// - Whatever the registered handler returns (a decode failure or a handler-reported error).
196    ///
197    /// An unregistered type with a one-shot/response shape is NOT an error: it returns
198    /// [`Dispatch::Dropped`] (silent drop). Dispatch NEVER panics on an unknown type.
199    pub fn dispatch(
200        &self,
201        message_type: MessageType,
202        shape: InteractionShape,
203        payload: &[u8],
204    ) -> Result<Dispatch> {
205        match self.handlers.get(&message_type) {
206            Some(handler) => handler(payload).map(|()| Dispatch::Handled),
207            None => match shape {
208                InteractionShape::Request | InteractionShape::StreamFrame => {
209                    Err(MessageError::UnsupportedType(message_type.0))
210                }
211                InteractionShape::OneShot | InteractionShape::Response => Ok(Dispatch::Dropped),
212            },
213        }
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use chia_streamable_macro::Streamable;
221    use std::sync::atomic::{AtomicU32, Ordering};
222    use std::sync::Arc;
223
224    /// A minimal Streamable payload standing in for a real dig-chat text message in these tests.
225    #[derive(Debug, Clone, PartialEq, Eq, Streamable)]
226    struct ChatText {
227        body: Vec<u8>,
228    }
229
230    /// A dig-chat text message kind, reserved in the dig-chat band (SPEC §4).
231    struct ChatTextKind;
232    impl MessageKind for ChatTextKind {
233        const TYPE_ID: MessageType = MessageType(BAND_DIG_CHAT);
234        type Payload = ChatText;
235    }
236
237    /// A second, distinct kind (peer-RPC ping) used to prove additive registration is non-disturbing.
238    #[derive(Debug, Clone, PartialEq, Eq, Streamable)]
239    struct Ping {
240        nonce: u64,
241    }
242    struct PingKind;
243    impl MessageKind for PingKind {
244        const TYPE_ID: MessageType = MessageType(BAND_PEER_RPC);
245        type Payload = Ping;
246    }
247
248    #[test]
249    fn register_then_dispatch_decodes_and_routes_to_the_handler() {
250        let seen = Arc::new(AtomicU32::new(0));
251        let sink = Arc::clone(&seen);
252
253        let mut registry = MessageRegistry::new();
254        registry
255            .register::<ChatTextKind, _>(move |msg: ChatText| {
256                sink.store(msg.body.len() as u32, Ordering::SeqCst);
257                Ok(())
258            })
259            .unwrap();
260
261        let payload = ChatText {
262            body: vec![1, 2, 3, 4, 5],
263        }
264        .to_bytes()
265        .unwrap();
266
267        let outcome = registry
268            .dispatch(ChatTextKind::TYPE_ID, InteractionShape::Request, &payload)
269            .unwrap();
270
271        assert_eq!(outcome, Dispatch::Handled);
272        assert_eq!(
273            seen.load(Ordering::SeqCst),
274            5,
275            "handler saw the decoded payload"
276        );
277    }
278
279    #[test]
280    fn unknown_type_request_returns_unsupported_type() {
281        let registry = MessageRegistry::new();
282        let unknown = MessageType(BAND_DIG_EMAIL + 0x42);
283        assert_eq!(
284            registry
285                .dispatch(unknown, InteractionShape::Request, &[])
286                .unwrap_err(),
287            MessageError::UnsupportedType(unknown.0)
288        );
289    }
290
291    #[test]
292    fn unknown_type_stream_frame_returns_unsupported_type() {
293        let registry = MessageRegistry::new();
294        let unknown = MessageType(BAND_DIG_VIDEO + 7);
295        assert_eq!(
296            registry
297                .dispatch(unknown, InteractionShape::StreamFrame, &[])
298                .unwrap_err(),
299            MessageError::UnsupportedType(unknown.0)
300        );
301    }
302
303    #[test]
304    fn unknown_type_one_shot_is_silently_dropped_not_an_error() {
305        let registry = MessageRegistry::new();
306        let unknown = MessageType(BAND_EXPERIMENTAL);
307        assert_eq!(
308            registry
309                .dispatch(unknown, InteractionShape::OneShot, &[])
310                .unwrap(),
311            Dispatch::Dropped
312        );
313    }
314
315    #[test]
316    fn unknown_type_response_is_silently_dropped() {
317        let registry = MessageRegistry::new();
318        let unknown = MessageType(BAND_PRESENCE + 1);
319        assert_eq!(
320            registry
321                .dispatch(unknown, InteractionShape::Response, &[])
322                .unwrap(),
323            Dispatch::Dropped
324        );
325    }
326
327    #[test]
328    fn registering_a_new_type_never_disturbs_existing_registrations() {
329        let mut registry = MessageRegistry::new();
330        registry
331            .register::<ChatTextKind, _>(|_: ChatText| Ok(()))
332            .unwrap();
333        assert!(registry.contains(ChatTextKind::TYPE_ID));
334
335        // Additively add a second, unrelated type.
336        registry.register::<PingKind, _>(|_: Ping| Ok(())).unwrap();
337
338        assert_eq!(registry.len(), 2);
339        assert!(
340            registry.contains(ChatTextKind::TYPE_ID),
341            "the first type is undisturbed"
342        );
343        assert!(registry.contains(PingKind::TYPE_ID));
344    }
345
346    #[test]
347    fn re_registering_the_same_id_is_refused_additive_only() {
348        let mut registry = MessageRegistry::new();
349        registry
350            .register::<ChatTextKind, _>(|_: ChatText| Ok(()))
351            .unwrap();
352        assert_eq!(
353            registry
354                .register::<ChatTextKind, _>(|_: ChatText| Ok(()))
355                .unwrap_err(),
356            MessageError::DuplicateType(ChatTextKind::TYPE_ID.0)
357        );
358        assert_eq!(
359            registry.len(),
360            1,
361            "the failed re-registration left the table unchanged"
362        );
363    }
364
365    #[test]
366    fn a_handler_decode_failure_propagates_and_never_panics() {
367        let mut registry = MessageRegistry::new();
368        registry.register::<PingKind, _>(|_: Ping| Ok(())).unwrap();
369        // A `Ping` is a u64 (8 bytes); a 3-byte frame cannot decode.
370        let err = registry
371            .dispatch(PingKind::TYPE_ID, InteractionShape::Request, &[1, 2, 3])
372            .unwrap_err();
373        assert!(matches!(err, MessageError::Codec(_)));
374    }
375
376    #[test]
377    fn a_handler_reported_error_propagates() {
378        let mut registry = MessageRegistry::new();
379        registry
380            .register::<PingKind, _>(|_: Ping| Err(MessageError::UnsupportedType(0)))
381            .unwrap();
382        let payload = Ping { nonce: 9 }.to_bytes().unwrap();
383        assert!(registry
384            .dispatch(PingKind::TYPE_ID, InteractionShape::OneShot, &payload)
385            .is_err());
386    }
387
388    #[test]
389    fn every_reserved_band_classifies_at_its_boundaries() {
390        let cases = [
391            (BAND_CORE, MessageBand::Core),
392            (BAND_CORE + 0xFF, MessageBand::Core),
393            (BAND_PEER_RPC, MessageBand::PeerRpc),
394            (BAND_PEER_RPC + 0xFF, MessageBand::PeerRpc),
395            (BAND_DIG_CHAT, MessageBand::DigChat),
396            (BAND_DIG_CHAT + 0xFF, MessageBand::DigChat),
397            (BAND_DIG_EMAIL, MessageBand::DigEmail),
398            (BAND_DIG_VIDEO, MessageBand::DigVideo),
399            (BAND_PRESENCE, MessageBand::Presence),
400            (BAND_IPC, MessageBand::Ipc),
401            (BAND_IPC + 0xFF, MessageBand::Ipc),
402            (BAND_EXPERIMENTAL, MessageBand::Experimental),
403            (0xFFFF_FFFF, MessageBand::Experimental),
404        ];
405        for (id, expected) in cases {
406            assert_eq!(MessageType(id).band(), expected, "id {id:#010x}");
407        }
408    }
409
410    #[test]
411    fn unallocated_ids_classify_as_reserved() {
412        // Just past the last named subsystem band (IPC ends at 0x06FF) and below experimental.
413        for id in [0x0000_0700, 0x0000_1000, 0x00FF_FFFF, 0x0FFF_FFFF] {
414            assert_eq!(
415                MessageType(id).band(),
416                MessageBand::Reserved,
417                "id {id:#010x}"
418            );
419        }
420    }
421
422    #[test]
423    fn named_core_type_constants_land_in_the_core_band() {
424        for mt in [
425            MessageType::CORE_HANDSHAKE,
426            MessageType::CORE_ACK,
427            MessageType::CORE_ERROR,
428            MessageType::CORE_KEEPALIVE,
429        ] {
430            assert_eq!(mt.band(), MessageBand::Core);
431        }
432    }
433}