Improve protocol implementation

unshell-protocol/src/tests/oneshot.rs

@@ -0,0 +1,231 @@
+use crate::{
+    endpoint::{Endpoint, error::EndpointError},
+    leaf::Leaf,
+    packet::Packet,
+};
+
+use alloc::{boxed::Box, string::ToString, vec, vec::Vec};
+use crossbeam_channel::{Receiver, Sender};
+
+const ENDPOINT_A: u32 = 0;
+const ENDPOINT_B: u32 = 1;
+
+const LEAF_CONTROLLER: u32 = 100;
+const LEAF_COMMS: u32 = 101;
+const LEAF_RESPONDER: u32 = 102;
+// const HOOK_ECHO: u16 = 500;
+
+fn echo_packet(path: Vec<u32>, hook_id: u16) -> Packet {
+    Packet {
+        hook_id,
+        end_hook: false,
+        path,
+        procedure_id: "echo".to_string(),
+        data: "ABC123".as_bytes().to_vec(),
+    }
+}
+
+struct ControllerLeaf {
+    has_run: bool,
+}
+struct CommsLeaf {
+    tx: Sender<Vec<u8>>,
+    rx: Receiver<Vec<u8>>,
+
+    remote_id: u32,
+    is_authority: bool,
+    started: bool,
+}
+struct ResponderLeaf;
+
+impl Leaf for ControllerLeaf {
+    fn get_id(&self) -> u32 {
+        LEAF_CONTROLLER
+    }
+
+    fn update(&mut self, endpoint: &mut Endpoint) {
+        if !self.has_run {
+            // Get next free available hook id
+            let hook_id = endpoint.get_hook_id();
+
+            // Create packet
+            let packet = echo_packet(vec![ENDPOINT_A, ENDPOINT_B], hook_id);
+
+            // Add packet to queue
+            let _ = endpoint.add_outbound(packet);
+
+            // Don't run again
+            self.has_run = true;
+        }
+    }
+}
+
+impl Leaf for CommsLeaf {
+    fn get_id(&self) -> u32 {
+        LEAF_COMMS
+    }
+
+    fn update(&mut self, endpoint: &mut Endpoint) {
+        if !self.started {
+            endpoint
+                .connections
+                .insert((self.remote_id, self.is_authority));
+        }
+
+        while !self.rx.is_empty() {
+            let packet = Packet::deserialize(&self.rx.recv().unwrap()).unwrap();
+
+            let _ = endpoint.add_inbound(packet);
+        }
+
+        endpoint.take_outbound_clear(self.remote_id, |packet| {
+            let data = packet.serialize().unwrap();
+            let _ = self.tx.send(data);
+        });
+    }
+}
+
+impl Leaf for ResponderLeaf {
+    fn get_id(&self) -> u32 {
+        LEAF_RESPONDER
+    }
+
+    fn update(&mut self, endpoint: &mut Endpoint) {
+        let local_id = endpoint.path.last().cloned().unwrap_or(0);
+        let mut packets = Vec::new();
+
+        endpoint.take_inbound_clear(local_id, |packet| {
+            let mut response = echo_packet(vec![ENDPOINT_A], packet.hook_id);
+            response.hook_id = packet.hook_id;
+            response.data = packet.data.clone();
+            packets.push(response);
+        });
+
+        for packet in packets {
+            endpoint.hooks.insert(packet.hook_id, 0);
+            let _ = endpoint.add_outbound(packet);
+        }
+    }
+}
+
+#[test]
+fn test_oneshot() {
+    let (tx_a, rx_a) = crossbeam_channel::unbounded();
+    let (tx_b, rx_b) = crossbeam_channel::unbounded();
+
+    let mut endpoint_a = crate::endpoint::Endpoint::new(
+        ENDPOINT_A,
+        vec![
+            Box::new(ControllerLeaf { has_run: false }),
+            Box::new(CommsLeaf {
+                tx: tx_b,
+                rx: rx_a,
+                remote_id: ENDPOINT_B,
+                is_authority: false,
+                started: false,
+            }),
+        ],
+    );
+    endpoint_a.path = vec![ENDPOINT_A];
+
+    let mut endpoint_b = crate::endpoint::Endpoint::new(
+        ENDPOINT_B,
+        vec![
+            Box::new(ResponderLeaf),
+            Box::new(CommsLeaf {
+                tx: tx_a,
+                rx: rx_b,
+                remote_id: ENDPOINT_A,
+                is_authority: true,
+                started: false,
+            }),
+        ],
+    );
+    endpoint_b.path = vec![ENDPOINT_A, ENDPOINT_B];
+
+    // Connections are registered routing state. The comms leaves also insert them
+    // during updates, but the first application packet should not depend on leaf order.
+    endpoint_a.connections.insert((ENDPOINT_B, false));
+    endpoint_b.connections.insert((ENDPOINT_A, true));
+
+    // Cycle 1: A sends request to B
+    endpoint_a.update();
+    endpoint_b.update();
+
+    // Cycle 2: B receives request and sends response to A
+    endpoint_b.update();
+    endpoint_a.update();
+
+    // Cycle 3: A's CommsLeaf needs one more update to pull the packet from the channel
+    // and put it into the inbound queue.
+    endpoint_a.update();
+
+    // Assertions on state
+    assert!(
+        endpoint_a.inbound.contains_key(&ENDPOINT_A),
+        "Endpoint A should have received response"
+    );
+    assert_eq!(
+        endpoint_a.inbound.get(&ENDPOINT_A).unwrap().len(),
+        1,
+        "Endpoint A should have exactly one packet"
+    );
+    let response = &endpoint_a
+        .inbound
+        .get(&ENDPOINT_A)
+        .unwrap()
+        .front()
+        .unwrap();
+    assert_eq!(response.data, "ABC123".as_bytes());
+    // assert_eq!(response.hook_id, HOOK_ECHO);
+}
+
+#[test]
+fn upward_outbound_without_hook_is_rejected() {
+    let mut endpoint = Endpoint::new(ENDPOINT_B, vec![]);
+    endpoint.path = vec![ENDPOINT_A, ENDPOINT_B];
+    endpoint.connections.insert((ENDPOINT_A, true));
+
+    let new_hook = endpoint.get_hook_id();
+
+    let error = endpoint
+        .add_outbound(echo_packet(vec![ENDPOINT_A], new_hook))
+        .unwrap_err();
+
+    assert!(matches!(error, EndpointError::HookNotExist));
+    assert!(endpoint.outbound.is_empty());
+}
+
+#[test]
+fn downward_outbound_without_hook_is_allowed() {
+    let mut endpoint = crate::endpoint::Endpoint::new(ENDPOINT_A, vec![]);
+    endpoint.path = vec![ENDPOINT_A];
+    endpoint.connections.insert((ENDPOINT_B, false));
+
+    let new_hook = endpoint.get_hook_id();
+
+    endpoint
+        .add_outbound(echo_packet(vec![ENDPOINT_A, ENDPOINT_B], new_hook))
+        .unwrap();
+
+    assert_eq!(endpoint.outbound.get(&ENDPOINT_B).unwrap().len(), 1);
+}
+
+#[test]
+fn deeper_upward_route_uses_parent_as_next_hop() {
+    const ENDPOINT_C: u32 = 2;
+
+    let mut endpoint = crate::endpoint::Endpoint::new(ENDPOINT_C, vec![]);
+    let new_hook = endpoint.get_hook_id();
+
+    endpoint.path = vec![ENDPOINT_A, ENDPOINT_B, ENDPOINT_C];
+    endpoint.hooks.insert(new_hook, ENDPOINT_A);
+    endpoint.connections.insert((ENDPOINT_B, true));
+
+    endpoint
+        .add_outbound(echo_packet(vec![ENDPOINT_A], new_hook))
+        .unwrap();
+
+    assert!(endpoint.outbound.contains_key(&ENDPOINT_B));
+    assert!(!endpoint.outbound.contains_key(&ENDPOINT_A));
+}