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