Skip to main content

unb_core/
route_control.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::{CoreError, Kind};
5
6pub const MAX_ROUTE_PATH: usize = 8;
7pub const MAX_ROUTES_PER_UPDATE: usize = 4096;
8pub const MAX_SUBJECT_LEN: usize = 256;
9pub const MAX_ROUTE_IDENTIFIER_LEN: usize = 256;
10
11pub fn validate_node_identifier(node: &str) -> Result<(), CoreError> {
12    if node.is_empty() || node.len() > MAX_ROUTE_IDENTIFIER_LEN {
13        return Err(CoreError::Establishment(
14            "node identifier is empty or exceeds its length limit".into(),
15        ));
16    }
17    Ok(())
18}
19
20#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
21pub struct NodeIdentity {
22    pub node_id: String,
23    pub instance_id: String,
24    pub epoch: u64,
25    #[serde(default, skip_serializing_if = "Value::is_null")]
26    pub proof: Value,
27}
28
29#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
30pub struct RouteAdvertisement {
31    pub subject: String,
32    pub owner: String,
33    pub owner_instance: String,
34    pub owner_epoch: u64,
35    pub owner_revision: u64,
36    pub distance: u32,
37    pub path: Vec<String>,
38}
39
40#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
41pub struct RouteWithdrawal {
42    pub subject: String,
43    pub owner: String,
44    pub owner_instance: String,
45    pub owner_epoch: u64,
46    pub owner_revision: u64,
47}
48
49#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
50pub struct RouteSnapshot {
51    pub generation: u64,
52    pub routes: Vec<RouteAdvertisement>,
53}
54
55impl RouteSnapshot {
56    pub fn canonical(generation: u64, mut routes: Vec<RouteAdvertisement>) -> RouteSnapshot {
57        routes.sort();
58        RouteSnapshot { generation, routes }
59    }
60}
61
62#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
63pub struct RouteDelta {
64    pub generation: u64,
65    pub upsert: Vec<RouteAdvertisement>,
66    pub withdraw: Vec<RouteWithdrawal>,
67}
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71pub enum RouteAckStatus {
72    Applied,
73    ResyncRequired,
74}
75
76#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
77pub struct RouteAck {
78    pub generation: u64,
79    pub status: RouteAckStatus,
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub enum SessionClass {
84    Unclassified,
85    NodeCandidate,
86    Client,
87}
88
89impl SessionClass {
90    pub fn classify_first(kind: Kind) -> SessionClass {
91        match kind {
92            Kind::Identify => SessionClass::NodeCandidate,
93            Kind::Request | Kind::Subscribe | Kind::Channel | Kind::Discover => {
94                SessionClass::Client
95            }
96            _ => SessionClass::Unclassified,
97        }
98    }
99}
100
101#[derive(Debug, Default, Clone, PartialEq)]
102pub(crate) struct Establishment {
103    sent_identify: bool,
104    remote_identity: Option<NodeIdentity>,
105    accepted_remote: bool,
106    remote_accepted_us: bool,
107    sent_snapshot: bool,
108    sent_snapshot_generation: Option<u64>,
109    applied_remote_snapshot: bool,
110    remote_acked_snapshot: bool,
111}
112
113impl Establishment {
114    pub(crate) fn identify_sent(&mut self) -> Result<(), CoreError> {
115        if self.sent_identify {
116            return Err(CoreError::Establishment("identify was already sent".into()));
117        }
118        self.sent_identify = true;
119        Ok(())
120    }
121
122    pub fn on_identify(&mut self, identity: NodeIdentity) -> Result<(), CoreError> {
123        if self.remote_identity.is_some() {
124            return Err(CoreError::Establishment(
125                "peer identity was already declared".into(),
126            ));
127        }
128        self.remote_identity = Some(identity);
129        Ok(())
130    }
131
132    pub(crate) fn local_accept(&mut self) -> Result<&NodeIdentity, CoreError> {
133        let Some(identity) = self.remote_identity.as_ref() else {
134            return Err(CoreError::Establishment(
135                "cannot accept an identity that was never declared".into(),
136            ));
137        };
138        self.accepted_remote = true;
139        Ok(identity)
140    }
141
142    pub(crate) fn on_identity_accepted(&mut self) -> Result<(), CoreError> {
143        if !self.sent_identify {
144            return Err(CoreError::Establishment(
145                "identity_accepted before identify was sent".into(),
146            ));
147        }
148        self.remote_accepted_us = true;
149        Ok(())
150    }
151
152    pub fn identities_accepted(&self) -> bool {
153        self.accepted_remote && self.remote_accepted_us
154    }
155
156    pub(crate) fn snapshot_sent(&mut self, generation: u64) -> Result<(), CoreError> {
157        if !self.identities_accepted() {
158            return Err(CoreError::Establishment(
159                "route snapshot before bilateral identity acceptance".into(),
160            ));
161        }
162        self.sent_snapshot = true;
163        self.sent_snapshot_generation = Some(generation);
164        Ok(())
165    }
166
167    pub(crate) fn on_snapshot_applied(&mut self) -> Result<(), CoreError> {
168        if !self.identities_accepted() {
169            return Err(CoreError::Establishment(
170                "route snapshot before bilateral identity acceptance".into(),
171            ));
172        }
173        self.applied_remote_snapshot = true;
174        Ok(())
175    }
176
177    pub(crate) fn on_route_ack(&mut self, ack: &RouteAck) -> Result<(), CoreError> {
178        if !self.sent_snapshot {
179            return Err(CoreError::Establishment(
180                "route_ack before the initial snapshot was sent".into(),
181            ));
182        }
183        if ack.status != RouteAckStatus::Applied
184            || self.sent_snapshot_generation != Some(ack.generation)
185        {
186            return Err(CoreError::Establishment(
187                "initial snapshot ack was not Applied for the exact generation".into(),
188            ));
189        }
190        self.remote_acked_snapshot = true;
191        Ok(())
192    }
193
194    pub fn remote_identity(&self) -> Option<&NodeIdentity> {
195        self.remote_identity.as_ref()
196    }
197
198    pub(crate) fn ready(&self) -> bool {
199        self.identities_accepted() && self.applied_remote_snapshot && self.remote_acked_snapshot
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn node_identifier_validation_uses_route_identity_limits() {
209        assert!(validate_node_identifier("node").is_ok());
210        assert!(validate_node_identifier("").is_err());
211        assert!(validate_node_identifier(&"n".repeat(MAX_ROUTE_IDENTIFIER_LEN + 1)).is_err());
212    }
213    use crate::{Envelope, DEFAULT_HOPS, PROTOCOL_VERSION};
214    use bytes::Bytes;
215
216    fn identity(node: &str) -> NodeIdentity {
217        NodeIdentity {
218            node_id: node.into(),
219            instance_id: format!("{node}-inst"),
220            epoch: 1,
221            proof: Value::Null,
222        }
223    }
224
225    fn advertisement(
226        subject: &str,
227        owner: &str,
228        distance: u32,
229        path: &[&str],
230    ) -> RouteAdvertisement {
231        RouteAdvertisement {
232            subject: subject.into(),
233            owner: owner.into(),
234            owner_instance: format!("{owner}-inst"),
235            owner_epoch: 1,
236            owner_revision: 0,
237            distance,
238            path: path.iter().map(|s| s.to_string()).collect(),
239        }
240    }
241
242    fn control_frame(kind: Kind, payload: Bytes) -> Envelope {
243        Envelope {
244            v: PROTOCOL_VERSION,
245            id: "f1".into(),
246            subject: String::new(),
247            kind,
248            corr: None,
249            seq: None,
250            hops: None,
251            body_token: None,
252            payload,
253            path: Vec::new(),
254            headers: Default::default(),
255        }
256    }
257
258    #[test]
259    fn identity_payload_round_trips_through_an_envelope() {
260        let sent = identity("node-a");
261        let frame = control_frame(
262            Kind::Identify,
263            Envelope::encode_payload(&serde_json::to_value(&sent).unwrap()),
264        );
265        let decoded = Envelope::decode(frame.encode()).unwrap();
266        assert_eq!(decoded.kind, Kind::Identify);
267        assert_eq!(decoded.parse_payload::<NodeIdentity>().unwrap(), sent);
268    }
269
270    #[test]
271    fn an_arbitrary_proof_round_trips_through_the_identify_frame() {
272        let mut sent = identity("node-proof");
273        sent.proof = serde_json::json!({ "region": "eu", "tags": [1, 2, 3], "ok": true });
274        let frame = control_frame(
275            Kind::Identify,
276            Envelope::encode_payload(&serde_json::to_value(&sent).unwrap()),
277        );
278        let decoded = Envelope::decode(frame.encode()).unwrap();
279        assert_eq!(
280            decoded.parse_payload::<NodeIdentity>().unwrap().proof,
281            sent.proof
282        );
283    }
284
285    #[test]
286    fn a_null_proof_is_omitted_from_the_identify_payload_and_defaults_back() {
287        let sent = identity("node-null");
288        let json = serde_json::to_value(&sent).unwrap();
289        assert!(
290            json.get("proof").is_none(),
291            "null proof is omitted on the wire"
292        );
293        let decoded: NodeIdentity = serde_json::from_value(json).unwrap();
294        assert_eq!(decoded.proof, Value::Null);
295    }
296
297    #[test]
298    fn snapshot_canonical_order_is_input_order_independent() {
299        let a = advertisement("chess", "leaf-a", 0, &["leaf-a"]);
300        let b = advertisement("chess", "leaf-a", 1, &["leaf-a", "hub"]);
301        let c = advertisement("weather", "leaf-b", 0, &["leaf-b"]);
302        let one = RouteSnapshot::canonical(1, vec![c.clone(), b.clone(), a.clone()]);
303        let two = RouteSnapshot::canonical(1, vec![b.clone(), a.clone(), c.clone()]);
304        assert_eq!(one, two);
305        assert_eq!(one.routes, vec![a, b, c]);
306    }
307
308    #[test]
309    fn snapshot_delta_and_ack_payloads_round_trip() {
310        let snapshot =
311            RouteSnapshot::canonical(3, vec![advertisement("chess", "leaf-a", 0, &["leaf-a"])]);
312        let frame = control_frame(
313            Kind::RouteSnapshot,
314            Envelope::encode_payload(&serde_json::to_value(&snapshot).unwrap()),
315        );
316        let decoded = Envelope::decode(frame.encode()).unwrap();
317        assert_eq!(decoded.parse_payload::<RouteSnapshot>().unwrap(), snapshot);
318
319        let delta = RouteDelta {
320            generation: 4,
321            upsert: vec![advertisement("go", "leaf-b", 1, &["leaf-b", "hub"])],
322            withdraw: vec![RouteWithdrawal {
323                subject: "chess".into(),
324                owner: "leaf-a".into(),
325                owner_instance: "leaf-a-inst".into(),
326                owner_epoch: 1,
327                owner_revision: 0,
328            }],
329        };
330        let frame = control_frame(
331            Kind::RouteDelta,
332            Envelope::encode_payload(&serde_json::to_value(&delta).unwrap()),
333        );
334        let decoded = Envelope::decode(frame.encode()).unwrap();
335        assert_eq!(decoded.parse_payload::<RouteDelta>().unwrap(), delta);
336
337        let ack = RouteAck {
338            generation: 4,
339            status: RouteAckStatus::ResyncRequired,
340        };
341        let frame = control_frame(
342            Kind::RouteAck,
343            Envelope::encode_payload(&serde_json::to_value(&ack).unwrap()),
344        );
345        let decoded = Envelope::decode(frame.encode()).unwrap();
346        assert_eq!(decoded.parse_payload::<RouteAck>().unwrap(), ack);
347    }
348
349    #[test]
350    fn v2_control_kinds_encode_snake_case() {
351        for (kind, name) in [
352            (Kind::Identify, "\"identify\""),
353            (Kind::IdentityAccepted, "\"identity_accepted\""),
354            (Kind::RouteSnapshot, "\"route_snapshot\""),
355            (Kind::RouteDelta, "\"route_delta\""),
356            (Kind::RouteAck, "\"route_ack\""),
357        ] {
358            assert_eq!(serde_json::to_string(&kind).unwrap(), name);
359        }
360    }
361
362    #[test]
363    fn a_missing_payload_on_a_control_frame_is_malformed() {
364        let frame = control_frame(Kind::Identify, Bytes::new());
365        let decoded = Envelope::decode(frame.encode()).unwrap();
366        assert!(matches!(
367            decoded.parse_payload::<NodeIdentity>(),
368            Err(CoreError::Malformed(_))
369        ));
370    }
371
372    #[test]
373    fn the_legal_establishment_sequence_reaches_ready() {
374        let mut establishment = Establishment::default();
375        assert!(!establishment.ready());
376        establishment.identify_sent().unwrap();
377        establishment.on_identify(identity("peer")).unwrap();
378        assert_eq!(establishment.local_accept().unwrap().node_id, "peer");
379        establishment.on_identity_accepted().unwrap();
380        assert!(!establishment.ready());
381        establishment.snapshot_sent(7).unwrap();
382        establishment.on_snapshot_applied().unwrap();
383        assert!(!establishment.ready());
384        establishment
385            .on_route_ack(&RouteAck {
386                generation: 7,
387                status: RouteAckStatus::Applied,
388            })
389            .unwrap();
390        assert!(establishment.ready());
391    }
392
393    #[test]
394    fn identify_twice_is_rejected() {
395        let mut establishment = Establishment::default();
396        establishment.identify_sent().unwrap();
397        assert!(matches!(
398            establishment.identify_sent(),
399            Err(CoreError::Establishment(_))
400        ));
401        establishment.on_identify(identity("peer")).unwrap();
402        assert!(matches!(
403            establishment.on_identify(identity("peer")),
404            Err(CoreError::Establishment(_))
405        ));
406    }
407
408    #[test]
409    fn acceptance_before_identify_is_rejected() {
410        let mut establishment = Establishment::default();
411        assert!(matches!(
412            establishment.on_identity_accepted(),
413            Err(CoreError::Establishment(_))
414        ));
415        assert!(matches!(
416            establishment.local_accept(),
417            Err(CoreError::Establishment(_))
418        ));
419    }
420
421    #[test]
422    fn snapshots_before_bilateral_identity_acceptance_are_rejected() {
423        let mut establishment = Establishment::default();
424        establishment.identify_sent().unwrap();
425        establishment.on_identify(identity("peer")).unwrap();
426        establishment.local_accept().unwrap();
427        assert!(matches!(
428            establishment.snapshot_sent(1),
429            Err(CoreError::Establishment(_))
430        ));
431        assert!(matches!(
432            establishment.on_snapshot_applied(),
433            Err(CoreError::Establishment(_))
434        ));
435    }
436
437    #[test]
438    fn an_ack_before_the_snapshot_was_sent_is_rejected() {
439        let mut establishment = Establishment::default();
440        establishment.identify_sent().unwrap();
441        establishment.on_identify(identity("peer")).unwrap();
442        establishment.local_accept().unwrap();
443        establishment.on_identity_accepted().unwrap();
444        assert!(matches!(
445            establishment.on_route_ack(&RouteAck {
446                generation: 1,
447                status: RouteAckStatus::Applied,
448            }),
449            Err(CoreError::Establishment(_))
450        ));
451    }
452
453    #[test]
454    fn initial_snapshot_requires_an_applied_ack_for_the_exact_generation() {
455        let ready_for_snapshot = || {
456            let mut establishment = Establishment::default();
457            establishment.identify_sent().unwrap();
458            establishment.on_identify(identity("peer")).unwrap();
459            establishment.local_accept().unwrap();
460            establishment.on_identity_accepted().unwrap();
461            establishment.snapshot_sent(4).unwrap();
462            establishment
463        };
464        for ack in [
465            RouteAck {
466                generation: 3,
467                status: RouteAckStatus::Applied,
468            },
469            RouteAck {
470                generation: 4,
471                status: RouteAckStatus::ResyncRequired,
472            },
473        ] {
474            assert!(matches!(
475                ready_for_snapshot().on_route_ack(&ack),
476                Err(CoreError::Establishment(_))
477            ));
478        }
479    }
480
481    #[test]
482    fn the_first_frame_classifies_the_session() {
483        assert_eq!(
484            SessionClass::classify_first(Kind::Identify),
485            SessionClass::NodeCandidate
486        );
487        for kind in [
488            Kind::Request,
489            Kind::Subscribe,
490            Kind::Channel,
491            Kind::Discover,
492        ] {
493            assert_eq!(SessionClass::classify_first(kind), SessionClass::Client);
494        }
495        for kind in [Kind::Ping, Kind::Pong] {
496            assert_eq!(
497                SessionClass::classify_first(kind),
498                SessionClass::Unclassified
499            );
500        }
501    }
502
503    #[test]
504    fn the_hop_budget_matches_the_spec_and_is_independent_of_the_route_path_limit() {
505        assert_eq!(DEFAULT_HOPS, 8);
506        assert_eq!(MAX_ROUTE_PATH, 8);
507        assert_eq!(MAX_ROUTES_PER_UPDATE, 4096);
508        assert_eq!(MAX_SUBJECT_LEN, 256);
509    }
510}