fedimint_server/config/
dkg.rs

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
20// Implementation of the classic Pedersen DKG.
21
22struct 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    /// Runs a single step of the DKG algorithm
61    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                // Once everyone has send their commitments, send out the key shares...
95
96                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
156/// Runs the DKG algorithms with our peers. We do not handle any unexpected
157/// messages and all peers are expected to be cooperative.
158pub 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}