p2panda_encryption/data_scheme/test_utils/
network.rs

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

p2panda_encryption/data_scheme/test_utils/
network.rs
