fleck/src/test/encryption.rs

use std::{cell::Cell, rc::Rc};

use crate::fleck_core;

const SECRET : u32 = 0x12345678;

#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct EncryptedMessage {
    secret_payload: u32
}

impl fleck_core::MessageParams for EncryptedMessage {
    const NAME: &'static str = "EncryptedMessage";
    const SIGNED: bool = false;
    const ENCRYPTED: bool = true;
}

#[derive(Default)]
struct ReceiveService {
    counter: Cell<u32>,
}

impl crate::DefaultService for ReceiveService {
    fn setup(self: &Rc<Self>, api: Rc<crate::API>) {
        api.channel::<fleck_core::MessageChannel<EncryptedMessage>>().sub_eat(self, Self::message);
    }
}

impl ReceiveService {
    fn message(&self, msg: (fleck_core::msg::Metadata, EncryptedMessage)) {
        if msg.1 .secret_payload == SECRET {
            self.counter.set(self.counter.get() + 1);
        }
    }
}

#[test_log::test]
fn one_way_test() {
    let skeleton = super::TestSkeleton::new();

    let span_f1 = tracing::span!(tracing::Level::TRACE, "f1");
    let span_f2 = tracing::span!(tracing::Level::TRACE, "f2");

    let fvec = skeleton.initialize_nodes(2);
    let f1 = fvec[0].clone();
    let f2 = fvec[1].clone();

    f1.with_service(|ms: &fleck_core::MessageService| {
        ms.add_message_type::<EncryptedMessage>();
    });
    f2.with_service(|ms: &fleck_core::MessageService| {
        ms.add_message_type::<EncryptedMessage>();
    });

    f2.add_service::<ReceiveService>();

    // let the encryption handshake happen
    for _ in 0..5 {
        {
            let _guard = span_f1.enter();
            f1.run_once();
        }
        {
            let _guard = span_f2.enter();
            f2.run_once();
        }
    }

    // grab f1's node for f2
    let node_for_f2 = f1.with_service(|ns: &fleck_core::NodeService|
        ns.node_by_pubkey(&f2.with_service(|ns: &fleck_core::NodeService| *ns.self_node().pubkey()))).expect("handshake failure?");

    log::info!("--- test triggering a message to be sent ---");

    // trigger f1 to send a packet
    f1.queue::<fleck_core::SendPacketChannel>(fleck_core::msg::Message::build(EncryptedMessage { secret_payload: SECRET }).with_peer(node_for_f2.peer().unwrap()).with_node(node_for_f2));

    // let the communication commence...
    for _ in 0..5 {
        {
            let _guard = span_f1.enter();
            f1.run_once();
        }
        {
            let _guard = span_f2.enter();
            f2.run_once();
        }
    }

    assert_eq!(1, f2.with_service(|rs: &ReceiveService| rs.counter.get()));
}