hotshot_new_protocol/
network.rs

1use std::{
2    collections::{HashMap, HashSet},
3    num::NonZeroUsize,
4    sync::Arc,
5};
6
7pub use cliquenet::{Config as CliquenetConfig, NetAddr, Role};
8use cliquenet::{NetworkReceiver, NetworkSender, Slot, noise::Protocol, x25519::PublicKey};
9use hotshot_types::{
10    PeerConnectInfo,
11    data::{EpochNumber, ViewNumber},
12    epoch_membership::EpochMembershipCoordinator,
13    message::{EXTERNAL_MESSAGE_VERSION, MessageKind, UpgradeLock},
14    traits::{
15        metrics::{Counter, CounterFamily, Gauge, GaugeFamily, Metrics},
16        node_implementation::NodeType,
17    },
18    x25519::{self, Keypair},
19};
20use hotshot_utils::anytrace;
21use parking_lot::RwLock;
22use tracing::{error, info};
23
24use crate::message::{Message, MessageType, Unchecked, Validated};
25
26/// A predicate that discards inbound messages.
27///
28/// For tests that need to lose specific traffic rather than a whole connection.
29#[cfg(any(test, feature = "testing"))]
30pub type InboundFilter<T> = Box<dyn Fn(&Message<T, Unchecked>) -> bool + Send + Sync>;
31
32pub struct Cliquenet<T: NodeType> {
33    inner: Sender<T>,
34    receiver: NetworkReceiver,
35    #[cfg(any(test, feature = "testing"))]
36    ifilter: InboundFilter<T>,
37}
38
39impl<T: NodeType> std::fmt::Debug for Cliquenet<T> {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("Cliquenet")
42            .field("inner", &self.inner)
43            .field("receiver", &self.receiver)
44            .finish_non_exhaustive()
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct Sender<T: NodeType> {
50    my_keys: (T::SignatureKey, PublicKey),
51    sender: NetworkSender,
52    shared: Arc<RwLock<Shared<T::SignatureKey>>>,
53    upgrade_lock: UpgradeLock<T>,
54}
55
56#[derive(Debug)]
57struct Shared<K> {
58    peers: HashMap<K, PeerConnectInfo>,
59    epoch: EpochNumber,
60}
61
62/// Lower bound of `message_limit`.
63pub const MIN_MESSAGE_LIMIT: NonZeroUsize =
64    NonZeroUsize::new(10 * 1024 * 1024).expect("10 MiB > 0");
65
66/// Room above a full block for message envelopes.
67const MESSAGE_HEADROOM: usize = 64 * 1024;
68
69/// Message limit for blocks of at most `max_block_size`.
70pub fn message_limit(max_block_size: u64) -> NonZeroUsize {
71    let block = usize::try_from(max_block_size).unwrap_or(usize::MAX);
72    let limit = block.saturating_add(MESSAGE_HEADROOM);
73    NonZeroUsize::new(limit).map_or(MIN_MESSAGE_LIMIT, |n| n.max(MIN_MESSAGE_LIMIT))
74}
75
76impl<T: NodeType> Cliquenet<T> {
77    #[expect(clippy::too_many_arguments)]
78    pub async fn create<A, P, S>(
79        name: S,
80        signing_key: T::SignatureKey,
81        keypair: Keypair,
82        addr: A,
83        parties: P,
84        max_message_size: Option<NonZeroUsize>,
85        upgrade_lock: UpgradeLock<T>,
86        metrics: Box<dyn Metrics>,
87    ) -> Result<Self, NetworkError>
88    where
89        A: Into<cliquenet::NetAddr>,
90        P: IntoIterator<Item = (T::SignatureKey, PeerConnectInfo)>,
91        S: Into<String>,
92    {
93        let parties: HashMap<T::SignatureKey, PeerConnectInfo> = parties.into_iter().collect();
94
95        let cfg = cliquenet::Config::builder()
96            .name(name)
97            .keypair(keypair.into())
98            .bind(addr.into())
99            .parties(
100                parties
101                    .values()
102                    .map(|info| (info.x25519_key.into(), info.p2p_addr.clone())),
103            )
104            .noise_protocols([(1.into(), Protocol::IK_25519_AesGcm_Blake2s)])
105            .maybe_max_message_size(max_message_size)
106            .build();
107
108        Self::create_with_config(signing_key, upgrade_lock, cfg, parties, metrics).await
109    }
110
111    pub(crate) async fn create_with_config<P>(
112        signing_key: T::SignatureKey,
113        upgrade_lock: UpgradeLock<T>,
114        config: cliquenet::Config,
115        parties: P,
116        metrics: Box<dyn Metrics>,
117    ) -> Result<Self, NetworkError>
118    where
119        P: IntoIterator<Item = (T::SignatureKey, PeerConnectInfo)>,
120    {
121        let public_key = config.public_key();
122        let metrics = CliquenetMetrics::new(metrics);
123        let network = cliquenet::Network::create(config.with_metrics(metrics)).await?;
124        let peers: HashMap<_, _> = parties.into_iter().collect();
125
126        info!(peers = %peers.len(), "cliquenet created");
127
128        let (send, recv) = network.split_into();
129
130        Ok(Self {
131            #[cfg(any(test, feature = "testing"))]
132            ifilter: Box::new(|_| false),
133            inner: Sender {
134                my_keys: (signing_key, public_key),
135                sender: send,
136                shared: Arc::new(RwLock::new(Shared {
137                    peers,
138                    epoch: EpochNumber::new(0),
139                })),
140                upgrade_lock,
141            },
142            receiver: recv,
143        })
144    }
145
146    pub fn sender(&self) -> &Sender<T> {
147        &self.inner
148    }
149
150    /// The epoch whose peer window is applied.
151    #[cfg(test)]
152    pub(crate) fn epoch(&self) -> EpochNumber {
153        self.inner.shared.read().epoch
154    }
155
156    /// Discard inbound messages the predicate matches.
157    #[cfg(any(test, feature = "testing"))]
158    pub fn drop_inbound(&mut self, filter: InboundFilter<T>) {
159        self.ifilter = filter;
160    }
161
162    pub async fn receive(&mut self) -> Result<Message<T, Unchecked>, NetworkError> {
163        #[allow(clippy::never_loop)]
164        loop {
165            let (src, bytes) = self
166                .receiver
167                .receive()
168                .await
169                .ok_or(cliquenet::NetworkError::ChannelClosed)?;
170            let msg = self.deserialize(&bytes)?;
171            let key = self
172                .inner
173                .shared
174                .read()
175                .peers
176                .get(&msg.sender)
177                .map(|info| info.x25519_key)
178                .or_else(|| {
179                    (msg.sender == self.inner.my_keys.0).then_some(self.inner.my_keys.1.into())
180                });
181            if Some(src.into()) != key {
182                return Err(NetworkError::InvalidSender {
183                    msg: key,
184                    src: src.into(),
185                });
186            }
187            #[cfg(any(test, feature = "testing"))]
188            if (self.ifilter)(&msg) {
189                continue;
190            }
191            return Ok(msg);
192        }
193    }
194
195    pub async fn shutdown(&mut self) {
196        if let Ok(done) = self.inner.sender.shutdown() {
197            done.await
198        }
199    }
200
201    pub fn gc(&mut self, v: ViewNumber) -> Result<(), NetworkError> {
202        self.inner.sender.gc(Slot::new(*v))?;
203        Ok(())
204    }
205
206    pub fn add_peers(
207        &mut self,
208        r: Role,
209        ps: Vec<(T::SignatureKey, (PublicKey, NetAddr))>,
210    ) -> Result<(), NetworkError> {
211        let mut targets = Vec::new();
212        {
213            let mut shared = self.inner.shared.write();
214            for (k, (x, a)) in ps {
215                shared.peers.insert(
216                    k,
217                    PeerConnectInfo {
218                        x25519_key: x.into(),
219                        p2p_addr: a.clone(),
220                    },
221                );
222                targets.push((x, a))
223            }
224        }
225        self.inner.sender.add_peers(r, targets)?;
226        Ok(())
227    }
228
229    pub fn remove_peers(&mut self, ps: Vec<&T::SignatureKey>) -> Result<(), NetworkError> {
230        let mut targets = Vec::new();
231        {
232            let mut shared = self.inner.shared.write();
233            for k in ps {
234                if let Some(info) = shared.peers.remove(k) {
235                    targets.push(info.x25519_key.into())
236                }
237            }
238        }
239        self.inner.sender.remove_peers(targets)?;
240        Ok(())
241    }
242
243    pub fn assign_role(&mut self, r: Role, ps: Vec<&T::SignatureKey>) -> Result<(), NetworkError> {
244        let mut targets = Vec::new();
245        {
246            let shared = self.inner.shared.read();
247            for k in ps {
248                if let Some(info) = shared.peers.get(k) {
249                    targets.push(info.x25519_key.into())
250                }
251            }
252        }
253        self.inner.sender.assign_peers(r, targets)?;
254        Ok(())
255    }
256
257    /// Update peers on every epoch change.
258    ///
259    /// For any given epoch `e` we collect the validators of `e`, `e-1` and
260    /// `e+1` from the stake tables and merge their connection information.
261    ///
262    /// We keep validators that were in `e-1` but not in `e` for one additional
263    /// epoch and eagerly connect to new validators of `e+1`.
264    pub fn apply_epoch(
265        &mut self,
266        epoch: EpochNumber,
267        coord: &EpochMembershipCoordinator<T>,
268    ) -> Result<(), NetworkError> {
269        let ours = self.inner.shared.read().epoch;
270        if epoch <= ours {
271            info!(%epoch, %ours, "epoch already seen");
272            return Ok(());
273        }
274
275        // Validators of the new epoch.
276        let Some(curr_infos) = coord.epoch_peers(Some(epoch)) else {
277            error!(%epoch, "no stake table available");
278            return Ok(());
279        };
280
281        // Validators leaving are retained as peers for one additional epoch.
282        let prev_infos = if *epoch > 0 {
283            coord.epoch_peers(Some(epoch - 1)).unwrap_or_else(|| {
284                info!(%epoch, "previous epoch's stake table unavailable");
285                HashMap::new()
286            })
287        } else {
288            HashMap::new()
289        };
290
291        // Validators joining in the next epoch are connected to early.
292        let next_infos = coord.epoch_peers(Some(epoch + 1)).unwrap_or_else(|| {
293            info!(%epoch, "next epoch's stake table not available");
294            HashMap::new()
295        });
296
297        // Since connection information may be updated, we need to merge them,
298        // preferring the newest epoch's data, i.e. `next(curr(prev))`.
299        let mut merged_infos = prev_infos.clone();
300        for (k, v) in curr_infos.iter().chain(&next_infos) {
301            merged_infos.insert(k.clone(), v.clone());
302        }
303
304        let wanted: HashSet<T::SignatureKey> = curr_infos
305            .keys()
306            .chain(next_infos.keys())
307            .cloned()
308            .collect();
309
310        let retained: HashSet<T::SignatureKey> = curr_infos
311            .keys()
312            .chain(prev_infos.keys())
313            .cloned()
314            .collect();
315
316        let mut to_add: Vec<(T::SignatureKey, PeerConnectInfo)> = Vec::new();
317        let mut to_del: Vec<(T::SignatureKey, PeerConnectInfo)> = Vec::new();
318
319        for k in &wanted {
320            if let Some(Some(new_info)) = merged_infos.get(k) {
321                if Some(new_info) != self.inner.shared.read().peers.get(k) {
322                    info!(%epoch, peer = %k, "adding/updating network peer");
323                    to_add.push((k.clone(), new_info.clone()));
324                } else {
325                    info!(%epoch, peer = %k, "peer unchanged");
326                }
327            } else {
328                info!(%epoch, peer = %k, "ignoring peer without connection info");
329            }
330        }
331
332        // Remove peers that have left both the current and previous epochs.
333        for (k, info) in &self.inner.shared.read().peers {
334            if !(retained.contains(k) || wanted.contains(k)) {
335                info!(%epoch, peer = %k, "removing network peer");
336                to_del.push((k.clone(), info.clone()));
337            }
338        }
339
340        {
341            let peers = &mut self.inner.shared.write().peers;
342            for (k, _) in &to_del {
343                peers.remove(k);
344            }
345            for (k, info) in &to_add {
346                peers.insert(k.clone(), info.clone());
347            }
348        }
349
350        let add_targets: Vec<(PublicKey, NetAddr)> = to_add
351            .iter()
352            .map(|(_, i)| (i.x25519_key.into(), i.p2p_addr.clone()))
353            .collect();
354        let del_targets: Vec<PublicKey> = to_del.iter().map(|(_, i)| i.x25519_key.into()).collect();
355
356        if let Err(err) = self.inner.sender.add_peers(Role::Active, add_targets) {
357            error!(%epoch, %err, "network down; could not add peers to network");
358            return Err(err.into());
359        }
360
361        if let Err(err) = self.inner.sender.remove_peers(del_targets) {
362            error!(%epoch, %err, "network down; could not remove peers from network");
363            return Err(err.into());
364        }
365
366        info!(%epoch, peers = %self.inner.shared.read().peers.len());
367
368        self.inner.shared.write().epoch = epoch;
369
370        Ok(())
371    }
372
373    fn deserialize(&self, bytes: &[u8]) -> Result<Message<T, Unchecked>, NetworkError> {
374        match self
375            .inner
376            .upgrade_lock
377            .deserialize::<Message<T, Unchecked>>(bytes)
378        {
379            Ok((m, v)) => {
380                if v == EXTERNAL_MESSAGE_VERSION && !m.is_external() {
381                    let e = anytrace::warn!("received a non-external message with version 0.0");
382                    return Err(NetworkError::Serialize(e));
383                }
384                Ok(m)
385            },
386            Err(primary_err) => {
387                // Fallback: bytes may be a hotshot-types `Message<T>` carrying
388                // an `External` payload (this is how `Leaf2Fetcher` in the
389                // membership layer frames leaf-catchup requests/responses).
390                // If so, surface it as `MessageType::External` so the
391                // Coordinator can route it to the membership external
392                // channel just like a native new-protocol external message.
393                if let Ok((_v, hs_msg)) =
394                    versions::decode::<hotshot_types::message::Message<T>>(bytes)
395                    && let MessageKind::External(data) = hs_msg.kind
396                {
397                    return Ok(Message {
398                        sender: hs_msg.sender,
399                        message_type: MessageType::External(data),
400                    });
401                }
402                Err(NetworkError::Serialize(primary_err))
403            },
404        }
405    }
406}
407
408impl<T: NodeType> Sender<T> {
409    pub fn unicast(
410        &self,
411        v: ViewNumber,
412        to: &T::SignatureKey,
413        m: &Message<T, Validated>,
414    ) -> Result<(), NetworkError> {
415        let target = if *to == self.my_keys.0 {
416            self.my_keys.1
417        } else if let Some(info) = self.shared.read().peers.get(to) {
418            info.x25519_key.into()
419        } else {
420            error!(peer = %to, "unicast target not found");
421            return Ok(());
422        };
423        let bytes = self.serialize(m)?;
424        self.sender.unicast(Slot::new(*v), target, bytes)?;
425        Ok(())
426    }
427
428    pub fn multicast(
429        &self,
430        v: ViewNumber,
431        to: Vec<&T::SignatureKey>,
432        m: &Message<T, Validated>,
433    ) -> Result<(), NetworkError> {
434        let bytes = self.serialize(m)?;
435        let mut targets = Vec::new();
436        {
437            let shared = self.shared.read();
438            for t in to {
439                if let Some(info) = shared.peers.get(t) {
440                    targets.push(info.x25519_key.into())
441                } else if *t == self.my_keys.0 {
442                    targets.push(self.my_keys.1)
443                } else {
444                    error!(peer = %t, "multicast target not found");
445                }
446            }
447        }
448        self.sender.multicast(Slot::new(*v), targets, bytes)?;
449        Ok(())
450    }
451
452    pub fn broadcast(&self, v: ViewNumber, m: &Message<T, Validated>) -> Result<(), NetworkError> {
453        let bytes = self.serialize(m)?;
454        self.sender.broadcast(Slot::new(*v), bytes)?;
455        Ok(())
456    }
457
458    pub fn max_message_size(&self) -> NonZeroUsize {
459        self.sender.config().max_message_size()
460    }
461
462    fn serialize(&self, m: &Message<T, Validated>) -> Result<Vec<u8>, NetworkError> {
463        if let MessageType::External(bytes) = &m.message_type {
464            return Ok(bytes.clone());
465        }
466        let v = self.upgrade_lock.serialize(m)?;
467        Ok(v)
468    }
469}
470
471#[derive(Debug, thiserror::Error)]
472pub enum NetworkError {
473    #[error("cliquenet: {0}")]
474    Cliquenet(#[from] cliquenet::NetworkError),
475
476    #[error("serialization: {0}")]
477    Serialize(#[from] anytrace::Error),
478
479    #[error("message sender {msg:?} != message source {src}")]
480    InvalidSender {
481        msg: Option<x25519::PublicKey>,
482        src: x25519::PublicKey,
483    },
484}
485
486impl NetworkError {
487    pub fn is_critical(&self) -> bool {
488        matches!(
489            self,
490            Self::Cliquenet(
491                cliquenet::NetworkError::Bind(..)
492                    | cliquenet::NetworkError::ChannelClosed
493                    | cliquenet::NetworkError::BudgetClosed
494            )
495        )
496    }
497}
498
499struct CliquenetMetrics {
500    metrics: Box<dyn Metrics>,
501    gauges: RwLock<Gauges>,
502    counters: RwLock<Counters>,
503}
504
505#[derive(Default)]
506struct Gauges {
507    gauges: HashMap<PublicKey, HashMap<String, Box<dyn Gauge>>>,
508    family: HashMap<String, Box<dyn GaugeFamily>>,
509}
510
511#[derive(Default)]
512struct Counters {
513    counters: HashMap<PublicKey, HashMap<String, Box<dyn Counter>>>,
514    family: HashMap<String, Box<dyn CounterFamily>>,
515}
516
517impl CliquenetMetrics {
518    pub fn new(m: Box<dyn Metrics>) -> Self {
519        Self {
520            metrics: m.subgroup("cliquenet".to_string()),
521            gauges: RwLock::new(Gauges::default()),
522            counters: RwLock::new(Counters::default()),
523        }
524    }
525}
526
527// In here we lazily create counters and gauges based on their labels.
528// If not found, we create a family using the label, e.g. "connect_attempts",
529// indexed by the peer (key). Afterwards we create the actual counter or gauge,
530// and update its value. On the next call, the metric would be found and
531// updated right away.
532impl cliquenet::Metrics for CliquenetMetrics {
533    fn set(&self, key: &PublicKey, label: &str, val: usize) {
534        if let Some(g) = self
535            .gauges
536            .read()
537            .gauges
538            .get(key)
539            .and_then(|m| m.get(label))
540        {
541            return g.set(val);
542        }
543
544        let mut gauges = self.gauges.write();
545
546        // Check again, in case a concurrent write has created the gauge:
547        if let Some(g) = gauges.gauges.get(key).and_then(|m| m.get(label)) {
548            return g.set(val);
549        }
550
551        let g = gauges
552            .family
553            .entry(label.to_string())
554            .or_insert_with(|| {
555                self.metrics
556                    .gauge_family(label.to_string(), vec!["peer".to_string()])
557            })
558            .create(vec![key.to_string()]);
559
560        gauges
561            .gauges
562            .entry(*key)
563            .or_default()
564            .entry(label.to_string())
565            .or_insert(g)
566            .set(val)
567    }
568
569    fn add(&self, key: &PublicKey, label: &str, val: usize) {
570        if let Some(c) = self
571            .counters
572            .read()
573            .counters
574            .get(key)
575            .and_then(|m| m.get(label))
576        {
577            return c.add(val);
578        }
579
580        let mut counters = self.counters.write();
581
582        // Check again, in case a concurrent write has created the counter:
583        if let Some(c) = counters.counters.get(key).and_then(|m| m.get(label)) {
584            return c.add(val);
585        }
586
587        let c = counters
588            .family
589            .entry(label.to_string())
590            .or_insert_with(|| {
591                self.metrics
592                    .counter_family(label.to_string(), vec!["peer".to_string()])
593            })
594            .create(vec![key.to_string()]);
595
596        counters
597            .counters
598            .entry(*key)
599            .or_default()
600            .entry(label.to_string())
601            .or_insert(c)
602            .add(val)
603    }
604
605    fn del(&self, key: &PublicKey) {
606        let key_string = key.to_string();
607
608        {
609            let mut gauges = self.gauges.write();
610            for (label, _) in gauges.gauges.remove(key).into_iter().flatten() {
611                if let Some(f) = gauges.family.get(&label) {
612                    f.destroy(&[&key_string]);
613                }
614            }
615        }
616
617        {
618            let mut counters = self.counters.write();
619            for (label, _) in counters.counters.remove(key).into_iter().flatten() {
620                if let Some(f) = counters.family.get(&label) {
621                    f.destroy(&[&key_string]);
622                }
623            }
624        }
625    }
626}
hotshot_new_protocol/network.rs

hotshot_new_protocol/
network.rs
