1use std::collections::HashMap;
20
21use chia_traits::Streamable;
22
23use crate::envelope::{InteractionShape, MessageType};
24use crate::error::{MessageError, Result};
25
26pub const BAND_CORE: u32 = 0x0000_0000;
28pub const BAND_PEER_RPC: u32 = 0x0000_0100;
30pub const BAND_DIG_CHAT: u32 = 0x0000_0200;
32pub const BAND_DIG_EMAIL: u32 = 0x0000_0300;
34pub const BAND_DIG_VIDEO: u32 = 0x0000_0400;
36pub const BAND_PRESENCE: u32 = 0x0000_0500;
38pub const BAND_IPC: u32 = 0x0000_0600;
40pub const BAND_EXPERIMENTAL: u32 = 0x1000_0000;
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub enum MessageBand {
47 Core,
49 PeerRpc,
51 DigChat,
53 DigEmail,
55 DigVideo,
57 Presence,
59 Ipc,
61 Experimental,
63 Reserved,
65}
66
67impl MessageType {
68 pub const CORE_HANDSHAKE: MessageType = MessageType(BAND_CORE);
70 pub const CORE_ACK: MessageType = MessageType(BAND_CORE + 1);
72 pub const CORE_ERROR: MessageType = MessageType(BAND_CORE + 2);
74 pub const CORE_KEEPALIVE: MessageType = MessageType(BAND_CORE + 3);
76
77 #[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
95pub trait MessageKind {
102 const TYPE_ID: MessageType;
104
105 type Payload: Streamable;
107}
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub enum Dispatch {
113 Handled,
115 Dropped,
118}
119
120type Handler = Box<dyn Fn(&[u8]) -> Result<()> + Send + Sync>;
123
124#[derive(Default)]
132pub struct MessageRegistry {
133 handlers: HashMap<MessageType, Handler>,
134}
135
136impl MessageRegistry {
137 #[must_use]
139 pub fn new() -> Self {
140 Self::default()
141 }
142
143 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 #[must_use]
171 pub fn contains(&self, message_type: MessageType) -> bool {
172 self.handlers.contains_key(&message_type)
173 }
174
175 #[must_use]
177 pub fn len(&self) -> usize {
178 self.handlers.len()
179 }
180
181 #[must_use]
183 pub fn is_empty(&self) -> bool {
184 self.handlers.is_empty()
185 }
186
187 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 #[derive(Debug, Clone, PartialEq, Eq, Streamable)]
226 struct ChatText {
227 body: Vec<u8>,
228 }
229
230 struct ChatTextKind;
232 impl MessageKind for ChatTextKind {
233 const TYPE_ID: MessageType = MessageType(BAND_DIG_CHAT);
234 type Payload = ChatText;
235 }
236
237 #[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 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 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 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}