p2panda_encryption/message_scheme/test_utils/
network.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use std::collections::{HashMap, VecDeque};
4
5use crate::Rng;
6use crate::key_manager::KeyManager;
7use crate::key_registry::KeyRegistry;
8use crate::message_scheme::GroupError;
9use crate::message_scheme::group::{GroupConfig, GroupEvent, GroupState, MessageGroup};
10use crate::message_scheme::test_utils::dcgka::init_dcgka_state;
11use crate::message_scheme::test_utils::dgm::AckedTestDgm;
12use crate::message_scheme::test_utils::ordering::{ForwardSecureOrderer, TestMessage};
13use crate::test_utils::{MemberId, MessageId};
14use crate::traits::ForwardSecureGroupMessage;
15
16pub type TestGroupState = GroupState<
17    MemberId,
18    MessageId,
19    KeyRegistry<MemberId>,
20    AckedTestDgm<MemberId, MessageId>,
21    KeyManager,
22    ForwardSecureOrderer<AckedTestDgm<MemberId, MessageId>>,
23>;
24
25pub type TestGroupError = GroupError<
26    MemberId,
27    MessageId,
28    KeyRegistry<MemberId>,
29    AckedTestDgm<MemberId, MessageId>,
30    KeyManager,
31    ForwardSecureOrderer<AckedTestDgm<MemberId, MessageId>>,
32>;
33
34pub fn init_group_state<const N: usize>(
35    member_ids: [MemberId; N],
36    rng: &Rng,
37) -> [TestGroupState; N] {
38    init_dcgka_state(member_ids, rng)
39        .into_iter()
40        .map(|dcgka| {
41            let orderer =
42                ForwardSecureOrderer::<AckedTestDgm<MemberId, MessageId>>::init(dcgka.my_id);
43            TestGroupState {
44                my_id: dcgka.my_id,
45                dcgka,
46                orderer,
47                welcome: None,
48                ratchet: None,
49                decryption_ratchet: HashMap::new(),
50                config: GroupConfig::default(),
51            }
52        })
53        .collect::<Vec<TestGroupState>>()
54        .try_into()
55        .unwrap()
56}
57
58pub struct Network {
59    rng: Rng,
60    members: HashMap<MemberId, TestGroupState>,
61    queue: VecDeque<TestMessage<AckedTestDgm<MemberId, MessageId>>>,
62}
63
64impl Network {
65    pub fn new<const N: usize>(members: [MemberId; N], rng: Rng) -> Self {
66        let members = init_group_state(members, &rng);
67        Self {
68            members: HashMap::from_iter(members.into_iter().map(|state| (state.my_id, state))),
69            rng,
70            queue: VecDeque::new(),
71        }
72    }
73
74    pub fn create(&mut self, creator: MemberId, initial_members: Vec<MemberId>) {
75        let y = self.get_y(&creator);
76        let (y_i, message) = MessageGroup::create(y, initial_members, &self.rng).unwrap();
77        self.queue.push_back(message);
78        self.set_y(y_i);
79    }
80
81    pub fn add(&mut self, adder: MemberId, added: MemberId) {
82        let y = self.get_y(&adder);
83        let (y_i, message) = MessageGroup::add(y, added, &self.rng).unwrap();
84        self.queue.push_back(message);
85        self.set_y(y_i);
86    }
87
88    pub fn remove(&mut self, remover: MemberId, removed: MemberId) {
89        let y = self.get_y(&remover);
90        let (y_i, message) = MessageGroup::remove(y, removed, &self.rng).unwrap();
91        self.queue.push_back(message);
92        self.set_y(y_i);
93        self.get_y(&removed);
94    }
95
96    pub fn update(&mut self, updater: MemberId) {
97        let y = self.get_y(&updater);
98        let (y_i, message) = MessageGroup::update(y, &self.rng).unwrap();
99        self.queue.push_back(message);
100        self.set_y(y_i);
101    }
102
103    pub fn send(&mut self, sender: MemberId, plaintext: &[u8]) {
104        let y = self.get_y(&sender);
105        let (y_i, message) = MessageGroup::send(y, plaintext).unwrap();
106        self.queue.push_back(message);
107        self.set_y(y_i);
108    }
109
110    pub fn process(&mut self) -> Vec<(MemberId, MemberId, Vec<u8>)> {
111        if self.queue.is_empty() {
112            return Vec::new();
113        }
114
115        let mut decrypted_messages = Vec::new();
116        let member_ids: Vec<MemberId> = self.members.keys().cloned().collect();
117
118        while let Some(message) = self.queue.pop_front() {
119            for id in &member_ids {
120                // Do not process our own messages.
121                if &message.sender() == id {
122                    continue;
123                }
124
125                // Member processes each message broadcast to the group.
126                let y = self.get_y(id);
127                let (y_i, result) = MessageGroup::receive(y, &message, &self.rng).unwrap();
128                self.set_y(y_i);
129
130                let Some(result) = result else {
131                    continue;
132                };
133
134                for event in result.events {
135                    match event {
136                        GroupEvent::Control(control_message) => {
137                            // Processing messages might yield new ones, process these as well.
138                            self.queue.push_back(control_message);
139                        }
140                        GroupEvent::Application { plaintext, .. } => decrypted_messages.push((
141                            message.sender(), // Sender
142                            *id,              // Receiver
143                            plaintext,        // Decrypted content
144                        )),
145                        GroupEvent::RemovedOurselves => (),
146                    }
147                }
148            }
149        }
150
151        decrypted_messages.sort();
152        decrypted_messages
153    }
154
155    pub fn members(&self, member: &MemberId) -> Vec<MemberId> {
156        let y = self.members.get(member).expect("member exists");
157        let mut members = Vec::from_iter(MessageGroup::members(y).unwrap());
158        members.sort();
159        members
160    }
161
162    fn get_y(&mut self, member: &MemberId) -> TestGroupState {
163        self.members.remove(member).expect("member exists")
164    }
165
166    fn set_y(&mut self, y: TestGroupState) {
167        assert!(
168            self.members.insert(y.my_id, y).is_none(),
169            "state was removed before insertion"
170        );
171    }
172}
p2panda_encryption/message_scheme/test_utils/network.rs

p2panda_encryption/message_scheme/test_utils/
network.rs
