1use std::collections::BTreeMap;
2use std::iter::once;
3
4use anyhow::{Context, bail, ensure};
5use async_trait::async_trait;
6use bls12_381::{G1Projective, G2Projective, Scalar};
7use fedimint_core::bitcoin::hashes::sha256;
8use fedimint_core::config::{DkgMessage, P2PMessage};
9use fedimint_core::encoding::Encodable as _;
10use fedimint_core::net::peers::{DynP2PConnections, Recipient};
11use fedimint_core::{NumPeers, PeerId};
12use fedimint_logging::LOG_NET_PEER_DKG;
13use fedimint_server_core::config::{PeerHandleOps, g1, g2, scalar};
14use group::ff::Field;
15use rand::rngs::OsRng;
16use tracing::{info, trace};
17
18use super::peer_handle::PeerHandle;
19
20struct Dkg {
23 num_peers: NumPeers,
24 identity: PeerId,
25 polynomial: Vec<Scalar>,
26 hash_commitments: BTreeMap<PeerId, sha256::Hash>,
27 commitments: BTreeMap<PeerId, Vec<(G1Projective, G2Projective)>>,
28 sk_shares: BTreeMap<PeerId, Scalar>,
29}
30
31impl Dkg {
32 fn new(num_peers: NumPeers, identity: PeerId) -> Self {
33 let polynomial = (0..num_peers.threshold())
34 .map(|_| Scalar::random(&mut OsRng))
35 .collect::<Vec<Scalar>>();
36
37 let commitment = polynomial
38 .iter()
39 .map(|f| (g1(f), g2(f)))
40 .collect::<Vec<(G1Projective, G2Projective)>>();
41
42 Dkg {
43 num_peers,
44 identity,
45 polynomial,
46 hash_commitments: once((identity, commitment.consensus_hash_sha256())).collect(),
47 commitments: once((identity, commitment)).collect(),
48 sk_shares: BTreeMap::new(),
49 }
50 }
51
52 fn commitment(&self) -> Vec<(G1Projective, G2Projective)> {
53 self.polynomial.iter().map(|f| (g1(f), g2(f))).collect()
54 }
55
56 fn initial_message(&self) -> DkgMessage {
57 DkgMessage::Hash(self.commitment().consensus_hash_sha256())
58 }
59
60 fn step(&mut self, peer: PeerId, msg: DkgMessage) -> anyhow::Result<DkgStep> {
62 trace!(?peer, ?msg, "Running DKG step");
63 match msg {
64 DkgMessage::Hash(hash) => {
65 ensure!(
66 self.hash_commitments.insert(peer, hash).is_none(),
67 "DKG: peer {peer} sent us two hash commitments."
68 );
69
70 if self.hash_commitments.len() == self.num_peers.total() {
71 return Ok(DkgStep::Broadcast(DkgMessage::Commitment(
72 self.commitment(),
73 )));
74 }
75 }
76 DkgMessage::Commitment(polynomial) => {
77 ensure!(
78 *self.hash_commitments.get(&peer).with_context(|| format!(
79 "DKG: hash commitment not found for peer {peer}"
80 ))? == polynomial.consensus_hash_sha256(),
81 "DKG: polynomial commitment from peer {peer} is of wrong degree."
82 );
83
84 ensure!(
85 self.num_peers.threshold() == polynomial.len(),
86 "DKG: polynomial commitment from peer {peer} is of wrong degree."
87 );
88
89 ensure!(
90 self.commitments.insert(peer, polynomial).is_none(),
91 "DKG: peer {peer} sent us two commitments."
92 );
93
94 if self.commitments.len() == self.num_peers.total() {
97 let mut messages = vec![];
98
99 for peer in self.num_peers.peer_ids() {
100 let s = eval_poly_scalar(&self.polynomial, &scalar(&peer));
101
102 if peer == self.identity {
103 self.sk_shares.insert(self.identity, s);
104 } else {
105 messages.push((peer, DkgMessage::Share(s)));
106 }
107 }
108
109 return Ok(DkgStep::Messages(messages));
110 }
111 }
112 DkgMessage::Share(s) => {
113 let polynomial = self.commitments.get(&peer).with_context(|| {
114 format!("DKG: polynomial commitment not found for peer {peer}.")
115 })?;
116
117 let checksum: (G1Projective, G2Projective) = polynomial
118 .iter()
119 .zip((0..).map(|k| scalar(&self.identity).pow(&[k, 0, 0, 0])))
120 .map(|(c, x)| (c.0 * x, c.1 * x))
121 .reduce(|(a1, a2), (b1, b2)| (a1 + b1, a2 + b2))
122 .expect("DKG: polynomial commitment from peer is empty.");
123
124 ensure!(
125 (g1(&s), g2(&s)) == checksum,
126 "DKG: share from {peer} is invalid."
127 );
128
129 ensure!(
130 self.sk_shares.insert(peer, s).is_none(),
131 "Peer {peer} sent us two sk shares."
132 );
133
134 if self.sk_shares.len() == self.num_peers.total() {
135 let sks = self.sk_shares.values().sum();
136
137 let pks = (0..self.num_peers.threshold())
138 .map(|i| {
139 self.commitments
140 .values()
141 .map(|coefficients| coefficients[i])
142 .reduce(|(a1, a2), (b1, b2)| (a1 + b1, a2 + b2))
143 .expect("DKG: polynomial commitments are empty.")
144 })
145 .collect();
146
147 return Ok(DkgStep::Result((pks, sks)));
148 }
149 }
150 }
151
152 Ok(DkgStep::Messages(vec![]))
153 }
154}
155
156pub async fn run_dkg(
159 num_peers: NumPeers,
160 identity: PeerId,
161 connections: &DynP2PConnections<P2PMessage>,
162) -> anyhow::Result<(Vec<(G1Projective, G2Projective)>, Scalar)> {
163 let mut dkg = Dkg::new(num_peers, identity);
164
165 connections
166 .send(Recipient::Everyone, P2PMessage::Dkg(dkg.initial_message()))
167 .await;
168
169 loop {
170 for peer in num_peers.peer_ids().filter(|p| *p != identity) {
171 let message = connections
172 .receive_from_peer(peer)
173 .await
174 .context("Unexpected shutdown of p2p connections during dkg")?;
175
176 let message = match message {
177 P2PMessage::Dkg(message) => message,
178 _ => bail!("Received unexpected message: {message:?}"),
179 };
180
181 match dkg.step(peer, message)? {
182 DkgStep::Broadcast(message) => {
183 connections
184 .send(Recipient::Everyone, P2PMessage::Dkg(message))
185 .await;
186 }
187 DkgStep::Messages(messages) => {
188 for (peer, message) in messages {
189 connections
190 .send(Recipient::Peer(peer), P2PMessage::Dkg(message))
191 .await;
192 }
193 }
194 DkgStep::Result(result) => {
195 return Ok(result);
196 }
197 }
198 }
199 }
200}
201
202fn eval_poly_scalar(coefficients: &[Scalar], x: &Scalar) -> Scalar {
203 coefficients
204 .iter()
205 .copied()
206 .rev()
207 .reduce(|acc, coefficient| acc * x + coefficient)
208 .expect("We have at least one coefficient")
209}
210
211enum DkgStep {
212 Broadcast(DkgMessage),
213 Messages(Vec<(PeerId, DkgMessage)>),
214 Result((Vec<(G1Projective, G2Projective)>, Scalar)),
215}
216
217#[async_trait]
218impl PeerHandleOps for PeerHandle<'_> {
219 fn num_peers(&self) -> NumPeers {
220 self.num_peers
221 }
222
223 async fn run_dkg_g1(&self) -> anyhow::Result<(Vec<G1Projective>, Scalar)> {
224 info!(
225 target: LOG_NET_PEER_DKG,
226 "Running distributed key generation for group G1..."
227 );
228
229 run_dkg(self.num_peers, self.identity, self.connections)
230 .await
231 .map(|(poly, sk)| (poly.into_iter().map(|c| c.0).collect(), sk))
232 }
233
234 async fn run_dkg_g2(&self) -> anyhow::Result<(Vec<G2Projective>, Scalar)> {
235 info!(
236 target: LOG_NET_PEER_DKG,
237 "Running distributed key generation for group G2..."
238 );
239
240 run_dkg(self.num_peers, self.identity, self.connections)
241 .await
242 .map(|(poly, sk)| (poly.into_iter().map(|c| c.1).collect(), sk))
243 }
244
245 async fn exchange_bytes(&self, bytes: Vec<u8>) -> anyhow::Result<BTreeMap<PeerId, Vec<u8>>> {
246 info!(
247 target: LOG_NET_PEER_DKG,
248 "Exchanging raw bytes..."
249 );
250
251 let mut peer_data: BTreeMap<PeerId, Vec<u8>> = BTreeMap::new();
252
253 self.connections
254 .send(Recipient::Everyone, P2PMessage::Encodable(bytes.clone()))
255 .await;
256
257 peer_data.insert(self.identity, bytes);
258
259 for peer in self.num_peers.peer_ids().filter(|p| *p != self.identity) {
260 let message = self
261 .connections
262 .receive_from_peer(peer)
263 .await
264 .context("Unexpected shutdown of p2p connections")?;
265
266 match message {
267 P2PMessage::Encodable(bytes) => {
268 peer_data.insert(peer, bytes);
269 }
270 message => {
271 anyhow::bail!("Invalid message from {peer}: {message:?}");
272 }
273 }
274 }
275
276 Ok(peer_data)
277 }
278}
279
280#[cfg(test)]
281mod tests {
282 use std::collections::{BTreeMap, VecDeque};
283
284 use bls12_381::{G1Projective, G2Projective};
285 use fedimint_core::{NumPeersExt, PeerId};
286 use fedimint_server_core::config::{eval_poly_g1, eval_poly_g2, g1, g2};
287 use group::Curve;
288
289 use crate::config::dkg::{Dkg, DkgStep};
290
291 #[test_log::test]
292 fn test_dkg() {
293 let peers = (0..7_u16).map(PeerId::from).collect::<Vec<PeerId>>();
294
295 let mut dkgs = peers
296 .iter()
297 .map(|peer| (*peer, Dkg::new(peers.to_num_peers(), *peer)))
298 .collect::<BTreeMap<PeerId, Dkg>>();
299
300 let mut steps = dkgs
301 .iter()
302 .map(|(peer, dkg)| (*peer, DkgStep::Broadcast(dkg.initial_message())))
303 .collect::<VecDeque<(PeerId, DkgStep)>>();
304
305 let mut keys = BTreeMap::new();
306
307 while keys.len() < peers.len() {
308 match steps.pop_front().unwrap() {
309 (send_peer, DkgStep::Broadcast(message)) => {
310 for receive_peer in peers.iter().filter(|p| **p != send_peer) {
311 let step = dkgs
312 .get_mut(receive_peer)
313 .unwrap()
314 .step(send_peer, message.clone());
315
316 steps.push_back((*receive_peer, step.unwrap()));
317 }
318 }
319 (send_peer, DkgStep::Messages(messages)) => {
320 for (receive_peer, messages) in messages {
321 let step = dkgs
322 .get_mut(&receive_peer)
323 .unwrap()
324 .step(send_peer, messages);
325
326 steps.push_back((receive_peer, step.unwrap()));
327 }
328 }
329 (send_peer, DkgStep::Result(step_keys)) => {
330 keys.insert(send_peer, step_keys);
331 }
332 }
333 }
334
335 assert!(steps.is_empty());
336
337 for (peer, (poly, sks)) in keys {
338 let poly_g1: Vec<G1Projective> = poly.clone().into_iter().map(|c| c.0).collect();
339 let poly_g2: Vec<G2Projective> = poly.clone().into_iter().map(|c| c.1).collect();
340
341 assert_eq!(poly_g1.len(), 5);
342 assert_eq!(eval_poly_g1(&poly_g1, &peer), g1(&sks).to_affine());
343
344 assert_eq!(poly_g2.len(), 5);
345 assert_eq!(eval_poly_g2(&poly_g2, &peer), g2(&sks).to_affine());
346 }
347 }
348}