devimint/vars/
net_overrides.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
3
4use fedimint_core::envs::FM_IROH_CONNECT_OVERRIDES_ENV;
5use fedimint_core::{NumPeers, PeerId};
6use iroh_base::NodeAddr;
7use iroh_base::ticket::{NodeTicket, Ticket};
8
9use super::ToEnvVar;
10use crate::federation::{
11    FEDIMINTD_API_PORT_OFFSET, FEDIMINTD_P2P_PORT_OFFSET, PORTS_PER_FEDIMINTD,
12};
13
14#[derive(Debug, Clone)]
15pub struct FedimintdEndpoint {
16    node_id: iroh_base::NodeId,
17    secret_key: iroh_base::SecretKey,
18    port: u16,
19}
20
21impl FedimintdEndpoint {
22    fn new(port: u16) -> Self {
23        let secret_key = iroh_base::SecretKey::generate(&mut rand::thread_rng());
24
25        Self {
26            node_id: secret_key.public(),
27            secret_key,
28            port,
29        }
30    }
31
32    pub fn secret_key(&self) -> String {
33        self.secret_key.to_string()
34    }
35
36    pub fn port(&self) -> u16 {
37        self.port
38    }
39
40    fn to_override(&self) -> String {
41        let node_addr = NodeAddr {
42            node_id: self.node_id,
43            relay_url: None,
44            direct_addresses: BTreeSet::from([SocketAddr::V4(SocketAddrV4::new(
45                Ipv4Addr::LOCALHOST,
46                self.port,
47            ))]),
48        };
49        format!(
50            "{}={},",
51            self.node_id,
52            Ticket::serialize(&NodeTicket::from(node_addr))
53        )
54    }
55}
56
57#[derive(Debug, Clone)]
58pub struct FedimintdPeerOverrides {
59    pub p2p: FedimintdEndpoint,
60    pub api: FedimintdEndpoint,
61    pub base_port: u16,
62}
63
64impl FedimintdPeerOverrides {
65    fn new(base_port: u16) -> Self {
66        Self {
67            p2p: FedimintdEndpoint::new(base_port + FEDIMINTD_P2P_PORT_OFFSET),
68            api: FedimintdEndpoint::new(base_port + FEDIMINTD_API_PORT_OFFSET),
69            base_port,
70        }
71    }
72}
73
74#[derive(Debug, Clone)]
75pub struct FederationNetOverrides {
76    pub base_port: u16,
77    pub num_peers: NumPeers,
78    pub peers: BTreeMap<PeerId, FedimintdPeerOverrides>,
79}
80
81impl FederationNetOverrides {
82    pub fn new(base_port: u16, num_peers: NumPeers) -> Self {
83        let peers = num_peers
84            .peer_ids()
85            .map(|peer_id| {
86                (peer_id, {
87                    FedimintdPeerOverrides::new(
88                        base_port
89                            + u16::try_from(peer_id.to_usize()).expect("Can't fail")
90                                * PORTS_PER_FEDIMINTD,
91                    )
92                })
93            })
94            .collect();
95        Self {
96            base_port,
97            num_peers,
98            peers,
99        }
100    }
101
102    pub fn peer_expect(&self, peer_id: PeerId) -> &FedimintdPeerOverrides {
103        self.peers.get(&peer_id).expect("Wrong peer_id?")
104    }
105}
106
107#[derive(Debug, Clone)]
108pub struct FederationsNetOverrides {
109    federations: Vec<FederationNetOverrides>,
110}
111
112impl FederationsNetOverrides {
113    pub fn new(base_port: u16, num_federations: usize, num_peers: NumPeers) -> Self {
114        Self {
115            federations: (0..num_federations)
116                .map(|fed_i| {
117                    FederationNetOverrides::new(
118                        base_port + fed_i as u16 * PORTS_PER_FEDIMINTD * num_peers.total() as u16,
119                        num_peers,
120                    )
121                })
122                .collect(),
123        }
124    }
125
126    pub fn fed_expect(&self, fed_i: usize) -> &FederationNetOverrides {
127        self.federations.get(fed_i).expect("Wrong fed_i")
128    }
129
130    pub fn peer_expect(&self, fed_i: usize, peer_id: PeerId) -> &FedimintdPeerOverrides {
131        self.federations
132            .get(fed_i)
133            .expect("Wrong fed_i")
134            .peers
135            .get(&peer_id)
136            .expect("Wrong peer_id?")
137    }
138}
139
140impl ToEnvVar for FederationsNetOverrides {
141    fn to_env_values(&self, _base_env: &str) -> impl Iterator<Item = (String, String)> {
142        vec![(
143            FM_IROH_CONNECT_OVERRIDES_ENV.to_string(),
144            self.federations
145                .iter()
146                .flat_map(|f| f.peers.values())
147                .map(|peer| format!("{},{}", peer.p2p.to_override(), peer.api.to_override(),))
148                .collect::<Vec<String>>()
149                .join(","),
150        )]
151        .into_iter()
152    }
153}
devimint/vars/net_overrides.rs

devimint/vars/
net_overrides.rs
