devimint/vars/
net_overrides.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
3
4use fedimint_core::envs::{FM_GW_IROH_CONNECT_OVERRIDES_ENV, 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 FedimintIrohEndpoint {
16    node_id: iroh_base::NodeId,
17    secret_key: iroh_base::SecretKey,
18    port: u16,
19}
20
21impl FedimintIrohEndpoint {
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    pub fn node_id(&self) -> iroh_base::NodeId {
41        self.node_id
42    }
43
44    fn to_override(&self) -> String {
45        let node_addr = NodeAddr {
46            node_id: self.node_id,
47            relay_url: None,
48            direct_addresses: BTreeSet::from([SocketAddr::V4(SocketAddrV4::new(
49                Ipv4Addr::LOCALHOST,
50                self.port,
51            ))]),
52        };
53        format!(
54            "{}={},",
55            self.node_id,
56            Ticket::serialize(&NodeTicket::from(node_addr))
57        )
58    }
59}
60
61#[derive(Debug, Clone)]
62pub struct FedimintdPeerOverrides {
63    pub p2p: FedimintIrohEndpoint,
64    pub api: FedimintIrohEndpoint,
65    pub base_port: u16,
66}
67
68impl FedimintdPeerOverrides {
69    fn new(base_port: u16) -> Self {
70        Self {
71            p2p: FedimintIrohEndpoint::new(base_port + FEDIMINTD_P2P_PORT_OFFSET),
72            api: FedimintIrohEndpoint::new(base_port + FEDIMINTD_API_PORT_OFFSET),
73            base_port,
74        }
75    }
76}
77
78#[derive(Debug, Clone)]
79pub struct FederationNetOverrides {
80    pub base_port: u16,
81    pub num_peers: NumPeers,
82    pub peers: BTreeMap<PeerId, FedimintdPeerOverrides>,
83}
84
85impl FederationNetOverrides {
86    pub fn new(base_port: u16, num_peers: NumPeers) -> Self {
87        let peers = num_peers
88            .peer_ids()
89            .map(|peer_id| {
90                (peer_id, {
91                    FedimintdPeerOverrides::new(
92                        base_port
93                            + u16::try_from(peer_id.to_usize()).expect("Can't fail")
94                                * PORTS_PER_FEDIMINTD,
95                    )
96                })
97            })
98            .collect();
99        Self {
100            base_port,
101            num_peers,
102            peers,
103        }
104    }
105
106    pub fn peer_expect(&self, peer_id: PeerId) -> &FedimintdPeerOverrides {
107        self.peers.get(&peer_id).expect("Wrong peer_id?")
108    }
109}
110
111#[derive(Debug, Clone)]
112pub struct FederationsNetOverrides {
113    federations: Vec<FederationNetOverrides>,
114}
115
116impl FederationsNetOverrides {
117    pub fn new(base_port: u16, num_federations: usize, num_peers: NumPeers) -> Self {
118        Self {
119            federations: (0..num_federations)
120                .map(|fed_i| {
121                    FederationNetOverrides::new(
122                        base_port + fed_i as u16 * PORTS_PER_FEDIMINTD * num_peers.total() as u16,
123                        num_peers,
124                    )
125                })
126                .collect(),
127        }
128    }
129
130    pub fn fed_expect(&self, fed_i: usize) -> &FederationNetOverrides {
131        self.federations.get(fed_i).expect("Wrong fed_i")
132    }
133
134    pub fn peer_expect(&self, fed_i: usize, peer_id: PeerId) -> &FedimintdPeerOverrides {
135        self.federations
136            .get(fed_i)
137            .expect("Wrong fed_i")
138            .peers
139            .get(&peer_id)
140            .expect("Wrong peer_id?")
141    }
142}
143
144impl ToEnvVar for FederationsNetOverrides {
145    fn to_env_values(&self, _base_env: &str) -> impl Iterator<Item = (String, String)> {
146        vec![(
147            FM_IROH_CONNECT_OVERRIDES_ENV.to_string(),
148            self.federations
149                .iter()
150                .flat_map(|f| f.peers.values())
151                .map(|peer| format!("{},{}", peer.p2p.to_override(), peer.api.to_override(),))
152                .collect::<Vec<String>>()
153                .join(","),
154        )]
155        .into_iter()
156    }
157}
158
159#[derive(Debug, Clone)]
160pub struct GatewaydNetOverrides {
161    pub gateway_iroh_endpoints: Vec<FedimintIrohEndpoint>,
162}
163
164impl GatewaydNetOverrides {
165    pub fn new(base_port: u16, num_gateways: usize) -> Self {
166        Self {
167            gateway_iroh_endpoints: (0..num_gateways)
168                .map(|gw_i| FedimintIrohEndpoint::new(base_port + gw_i as u16))
169                .collect(),
170        }
171    }
172}
173
174impl ToEnvVar for GatewaydNetOverrides {
175    fn to_env_values(&self, _base_env: &str) -> impl Iterator<Item = (String, String)> {
176        vec![(
177            FM_GW_IROH_CONNECT_OVERRIDES_ENV.to_string(),
178            self.gateway_iroh_endpoints
179                .iter()
180                .map(|gw| gw.to_override().to_string())
181                .collect::<Vec<String>>()
182                .join(","),
183        )]
184        .into_iter()
185    }
186}
devimint/vars/net_overrides.rs

devimint/vars/
net_overrides.rs
