confium_coordinator/
refresh_coordinator.rs1use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct RefreshContribution {
16 pub from_party: u32,
18 pub refresh_shares: HashMap<u32, Vec<u8>>,
20}
21
22#[derive(Debug)]
24pub struct RefreshSession {
25 pub session_id: String,
26 pub threshold: u32,
27 pub party_count: u32,
28 pub contributions: HashMap<u32, RefreshContribution>,
29}
30
31impl RefreshSession {
32 pub fn new(session_id: &str, threshold: u32, party_count: u32) -> Self {
33 Self {
34 session_id: session_id.into(),
35 threshold,
36 party_count,
37 contributions: HashMap::new(),
38 }
39 }
40
41 pub fn submit_contribution(&mut self, contrib: RefreshContribution) -> Result<(), String> {
43 if contrib.from_party == 0 || contrib.from_party > self.party_count {
44 return Err(format!("invalid from_party: {}", contrib.from_party));
45 }
46 if self.contributions.contains_key(&contrib.from_party) {
47 return Err(format!("party {} already contributed", contrib.from_party));
48 }
49 self.contributions.insert(contrib.from_party, contrib);
50 Ok(())
51 }
52
53 pub fn is_complete(&self) -> bool {
55 self.contributions.len() == self.party_count as usize
56 }
57
58 pub fn aggregate_for_party(&self, party_idx: u32) -> Option<Vec<u8>> {
61 if !self.is_complete() {
62 return None;
63 }
64 let mut result: Option<Vec<u8>> = None;
65 for contrib in self.contributions.values() {
66 if let Some(share) = contrib.refresh_shares.get(&party_idx) {
67 result = Some(match result {
68 None => share.clone(),
69 Some(mut existing) => {
70 let len = existing.len().max(share.len());
71 existing.resize(len, 0);
72 for (i, &b) in share.iter().enumerate() {
73 existing[i] ^= b;
74 }
75 existing
76 }
77 });
78 }
79 }
80 result
81 }
82
83 pub fn contribution_count(&self) -> usize {
85 self.contributions.len()
86 }
87
88 pub fn missing_parties(&self) -> Vec<u32> {
90 (1..=self.party_count)
91 .filter(|i| !self.contributions.contains_key(i))
92 .collect()
93 }
94}
95
96pub fn generate_contribution(
100 from_party: u32,
101 party_count: u32,
102 share_size: usize,
103) -> RefreshContribution {
104 use rand_core::{OsRng, RngCore};
105 let mut refresh_shares = HashMap::new();
106 let mut remaining = vec![0u8; share_size];
107
108 for p in 1..party_count {
110 let mut share = vec![0u8; share_size];
111 OsRng.fill_bytes(&mut share);
112 for (i, &b) in share.iter().enumerate() {
113 remaining[i] ^= b;
114 }
115 refresh_shares.insert(p, share);
116 }
117
118 refresh_shares.insert(party_count, remaining);
120
121 RefreshContribution {
122 from_party,
123 refresh_shares,
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130
131 #[test]
132 fn new_session_empty() {
133 let session = RefreshSession::new("r1", 2, 3);
134 assert_eq!(session.contribution_count(), 0);
135 assert!(!session.is_complete());
136 }
137
138 #[test]
139 fn submit_contribution() {
140 let mut session = RefreshSession::new("r1", 2, 3);
141 let contrib = RefreshContribution {
142 from_party: 1,
143 refresh_shares: HashMap::new(),
144 };
145 session.submit_contribution(contrib).unwrap();
146 assert_eq!(session.contribution_count(), 1);
147 }
148
149 #[test]
150 fn complete_when_all_contributed() {
151 let mut session = RefreshSession::new("r1", 2, 3);
152 for p in 1..=3 {
153 session
154 .submit_contribution(RefreshContribution {
155 from_party: p,
156 refresh_shares: HashMap::new(),
157 })
158 .unwrap();
159 }
160 assert!(session.is_complete());
161 }
162
163 #[test]
164 fn duplicate_contribution_rejected() {
165 let mut session = RefreshSession::new("r1", 2, 3);
166 let contrib = RefreshContribution {
167 from_party: 1,
168 refresh_shares: HashMap::new(),
169 };
170 session.submit_contribution(contrib).unwrap();
171 let contrib2 = RefreshContribution {
172 from_party: 1,
173 refresh_shares: HashMap::new(),
174 };
175 assert!(session.submit_contribution(contrib2).is_err());
176 }
177
178 #[test]
179 fn aggregate_requires_complete() {
180 let session = RefreshSession::new("r1", 2, 3);
181 assert!(session.aggregate_for_party(1).is_none());
182 }
183
184 #[test]
185 fn aggregate_xors_shares() {
186 let mut session = RefreshSession::new("r1", 2, 2);
187 let mut shares1 = HashMap::new();
188 shares1.insert(1u32, vec![0xFF]);
189 shares1.insert(2u32, vec![0x0F]);
190 session
191 .submit_contribution(RefreshContribution {
192 from_party: 1,
193 refresh_shares: shares1,
194 })
195 .unwrap();
196
197 let mut shares2 = HashMap::new();
198 shares2.insert(1u32, vec![0xAA]);
199 shares2.insert(2u32, vec![0x55]);
200 session
201 .submit_contribution(RefreshContribution {
202 from_party: 2,
203 refresh_shares: shares2,
204 })
205 .unwrap();
206
207 let agg1 = session.aggregate_for_party(1).unwrap();
208 assert_eq!(agg1, vec![0x55]);
210 }
211
212 #[test]
213 fn generate_contribution_shares_sum_to_zero() {
214 let contrib = generate_contribution(1, 3, 32);
215 let mut xored = vec![0u8; 32];
217 for share in contrib.refresh_shares.values() {
218 for (i, &b) in share.iter().enumerate() {
219 xored[i] ^= b;
220 }
221 }
222 assert_eq!(xored, vec![0u8; 32]);
223 }
224
225 #[test]
226 fn missing_parties_lists_gaps() {
227 let mut session = RefreshSession::new("r1", 2, 5);
228 session
229 .submit_contribution(RefreshContribution {
230 from_party: 1,
231 refresh_shares: HashMap::new(),
232 })
233 .unwrap();
234 session
235 .submit_contribution(RefreshContribution {
236 from_party: 3,
237 refresh_shares: HashMap::new(),
238 })
239 .unwrap();
240 let missing = session.missing_parties();
241 assert_eq!(missing, vec![2, 4, 5]);
242 }
243}