Skip to main content

fedimint_testing/
federation.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3use std::time::Duration;
4
5use fedimint_api_client::api::{DynGlobalApi, FederationApiExt};
6use fedimint_client::module_init::ClientModuleInitRegistry;
7use fedimint_client::{Client, ClientHandleArc, RootSecret};
8use fedimint_client_module::AdminCreds;
9use fedimint_client_module::secret::{PlainRootSecretStrategy, RootSecretStrategy};
10use fedimint_connectors::ConnectorRegistry;
11use fedimint_core::PeerId;
12use fedimint_core::config::{ClientConfig, FederationId};
13use fedimint_core::core::ModuleKind;
14use fedimint_core::db::Database;
15use fedimint_core::db::mem_impl::MemDatabase;
16use fedimint_core::endpoint_constants::SESSION_COUNT_ENDPOINT;
17use fedimint_core::invite_code::InviteCode;
18use fedimint_core::module::{ApiAuth, ApiRequestErased};
19use fedimint_core::net::IP2PConnections;
20use fedimint_core::rustls::install_crypto_provider;
21use fedimint_core::task::{TaskGroup, block_in_place, sleep_in_test};
22use fedimint_gateway_common::ConnectFedPayload;
23use fedimint_gateway_server::{Gateway, IAdminGateway};
24use fedimint_logging::LOG_TEST;
25use fedimint_rocksdb::RocksDb;
26use fedimint_server::config::ServerConfig;
27use fedimint_server::core::ServerModuleInitRegistry;
28use fedimint_server::net::api::ApiSecrets;
29use fedimint_server::net::p2p::{ReconnectP2PConnections, p2p_status_channels};
30use fedimint_server::net::p2p_connector::{IP2PConnector, TlsTcpConnector};
31use fedimint_server::{ConnectionLimits, consensus};
32use fedimint_server_core::bitcoin_rpc::DynServerBitcoinRpc;
33use fedimint_testing_core::config::local_config_gen_params;
34use tracing::info;
35
36/// Test fixture for a running fedimint federation
37#[derive(Clone)]
38pub struct FederationTest {
39    configs: BTreeMap<PeerId, ServerConfig>,
40    server_init: ServerModuleInitRegistry,
41    client_init: ClientModuleInitRegistry,
42    _task: TaskGroup,
43    num_peers: u16,
44    num_offline: u16,
45    connectors: ConnectorRegistry,
46}
47
48impl FederationTest {
49    /// Create two clients, useful for send/receive tests
50    pub async fn two_clients(&self) -> (ClientHandleArc, ClientHandleArc) {
51        (self.new_client().await, self.new_client().await)
52    }
53
54    /// Create a client connected to this fed
55    pub async fn new_client(&self) -> ClientHandleArc {
56        let client_config = self.configs[&PeerId::from(0)]
57            .consensus
58            .to_client_config(&self.server_init)
59            .unwrap();
60
61        self.new_client_with(client_config, MemDatabase::new().into(), None)
62            .await
63    }
64
65    /// Create a client connected to this fed but using RocksDB instead of MemDB
66    pub async fn new_client_rocksdb(&self) -> ClientHandleArc {
67        let client_config = self.configs[&PeerId::from(0)]
68            .consensus
69            .to_client_config(&self.server_init)
70            .unwrap();
71
72        self.new_client_with(
73            client_config,
74            RocksDb::build(tempfile::tempdir().expect("Couldn't create temp dir"))
75                .open()
76                .await
77                .map_err(anyhow::Error::from)
78                .expect("Couldn't open DB")
79                .into(),
80            None,
81        )
82        .await
83    }
84
85    /// Create a new admin api for the given PeerId
86    pub async fn new_admin_api(&self, peer_id: PeerId) -> anyhow::Result<DynGlobalApi> {
87        let config = self.configs.get(&peer_id).expect("peer to have config");
88
89        Ok(DynGlobalApi::new_admin(
90            ConnectorRegistry::build_from_testing_env().bind().await,
91            peer_id,
92            config.consensus.api_endpoints()[&peer_id].url.clone(),
93            None,
94        ))
95    }
96
97    /// Create a new admin client connected to this fed
98    pub async fn new_admin_client(&self, peer_id: PeerId, auth: ApiAuth) -> ClientHandleArc {
99        let client_config = self.configs[&PeerId::from(0)]
100            .consensus
101            .to_client_config(&self.server_init)
102            .unwrap();
103
104        let admin_creds = AdminCreds { peer_id, auth };
105
106        self.new_client_with(client_config, MemDatabase::new().into(), Some(admin_creds))
107            .await
108    }
109
110    pub async fn new_client_with(
111        &self,
112        client_config: ClientConfig,
113        db: Database,
114        admin_creds: Option<AdminCreds>,
115    ) -> ClientHandleArc {
116        info!(target: LOG_TEST, "Setting new client with config");
117        let mut client_builder = Client::builder().await;
118        client_builder.with_module_inits(self.client_init.clone());
119        if let Some(admin_creds) = admin_creds {
120            client_builder.set_admin_creds(admin_creds);
121        }
122        let client_secret = Client::load_or_generate_client_secret(&db).await;
123        client_builder
124            .preview_with_existing_config(self.connectors.clone(), client_config, None)
125            .await
126            .join(
127                db,
128                RootSecret::StandardDoubleDerive(PlainRootSecretStrategy::to_root_secret(
129                    &client_secret,
130                )),
131            )
132            .await
133            .map(Arc::new)
134            .expect("Failed to build client")
135    }
136
137    /// Join a federation with an existing database and root secret
138    pub async fn join_client_with_db(
139        &self,
140        db: Database,
141        root_secret: RootSecret,
142    ) -> ClientHandleArc {
143        let client_config = self.configs[&PeerId::from(0)]
144            .consensus
145            .to_client_config(&self.server_init)
146            .unwrap();
147
148        info!(target: LOG_TEST, "Joining client with existing db");
149        let mut client_builder = Client::builder().await;
150        client_builder.with_module_inits(self.client_init.clone());
151        client_builder
152            .preview_with_existing_config(self.connectors.clone(), client_config, None)
153            .await
154            .join(db, root_secret)
155            .await
156            .map(Arc::new)
157            .expect("Failed to join client")
158    }
159
160    /// Create a recovering client with an existing database and root secret.
161    /// Returns both the client and the database so a new client can be created
162    /// with the same DB after recovery completes.
163    pub async fn recover_client_with_db(
164        &self,
165        db: Database,
166        root_secret: RootSecret,
167    ) -> ClientHandleArc {
168        let client_config = self.configs[&PeerId::from(0)]
169            .consensus
170            .to_client_config(&self.server_init)
171            .unwrap();
172
173        info!(target: LOG_TEST, "Recovering client with existing db");
174        let mut client_builder = Client::builder().await;
175        client_builder.with_module_inits(self.client_init.clone());
176        client_builder
177            .preview_with_existing_config(self.connectors.clone(), client_config, None)
178            .await
179            .recover(db, root_secret, None)
180            .await
181            .map(Arc::new)
182            .expect("Failed to recover client")
183    }
184
185    /// Open an existing client database (e.g., after recovery)
186    pub async fn open_client_with_db(
187        &self,
188        db: Database,
189        root_secret: RootSecret,
190    ) -> ClientHandleArc {
191        info!(target: LOG_TEST, "Opening client with existing db");
192        let mut client_builder = Client::builder().await;
193        client_builder.with_module_inits(self.client_init.clone());
194        client_builder
195            .open(self.connectors.clone(), db, root_secret)
196            .await
197            .map(Arc::new)
198            .expect("Failed to open client")
199    }
200
201    /// Return first invite code for gateways
202    pub fn invite_code(&self) -> InviteCode {
203        let peer_id = PeerId::from(0);
204        let cfg = &self.configs[&peer_id];
205        InviteCode::new(
206            cfg.consensus.api_endpoints()[&peer_id].url.clone(),
207            peer_id,
208            cfg.calculate_federation_id(),
209            None,
210        )
211    }
212
213    ///  Return the federation id
214    pub fn id(&self) -> FederationId {
215        self.configs[&PeerId::from(0)]
216            .consensus
217            .to_client_config(&self.server_init)
218            .unwrap()
219            .global
220            .calculate_federation_id()
221    }
222
223    /// Connects a gateway to this `FederationTest`
224    pub async fn connect_gateway(&self, gw: &Gateway) {
225        gw.handle_connect_federation(ConnectFedPayload {
226            invite_code: self.invite_code().to_string(),
227            use_tor: Some(false),
228            recover: Some(false),
229        })
230        .await
231        .expect("Failed to connect federation");
232    }
233
234    /// Return all online PeerIds
235    pub fn online_peer_ids(&self) -> impl Iterator<Item = PeerId> + use<> {
236        // we can assume this ordering since peers are started in ascending order
237        (0..(self.num_peers - self.num_offline)).map(PeerId::from)
238    }
239
240    /// Returns true if the federation is running in a degraded state
241    pub fn is_degraded(&self) -> bool {
242        self.num_offline > 0
243    }
244}
245
246/// Builder struct for creating a `FederationTest`.
247#[derive(Clone, Debug)]
248pub struct FederationTestBuilder {
249    num_peers: u16,
250    num_offline: u16,
251    base_port: u16,
252    primary_module_kind: ModuleKind,
253    version_hash: String,
254    server_init: ServerModuleInitRegistry,
255    client_init: ClientModuleInitRegistry,
256    bitcoin_rpc_connection: DynServerBitcoinRpc,
257    enable_mint_fees: bool,
258}
259
260impl FederationTestBuilder {
261    pub fn new(
262        server_init: ServerModuleInitRegistry,
263        client_init: ClientModuleInitRegistry,
264        primary_module_kind: ModuleKind,
265        num_offline: u16,
266        bitcoin_rpc_connection: DynServerBitcoinRpc,
267    ) -> FederationTestBuilder {
268        let num_peers = 4;
269        Self {
270            num_peers,
271            num_offline,
272            base_port: block_in_place(|| fedimint_portalloc::port_alloc(num_peers * 3))
273                .expect("Failed to allocate a port range"),
274            primary_module_kind,
275            version_hash: "fedimint-testing-dummy-version-hash".to_owned(),
276            server_init,
277            client_init,
278            bitcoin_rpc_connection,
279            enable_mint_fees: true,
280        }
281    }
282
283    pub fn num_peers(mut self, num_peers: u16) -> FederationTestBuilder {
284        self.num_peers = num_peers;
285        self
286    }
287
288    pub fn num_offline(mut self, num_offline: u16) -> FederationTestBuilder {
289        self.num_offline = num_offline;
290        self
291    }
292
293    pub fn base_port(mut self, base_port: u16) -> FederationTestBuilder {
294        self.base_port = base_port;
295        self
296    }
297
298    pub fn primary_module_kind(mut self, primary_module_kind: ModuleKind) -> FederationTestBuilder {
299        self.primary_module_kind = primary_module_kind;
300        self
301    }
302
303    pub fn version_hash(mut self, version_hash: String) -> FederationTestBuilder {
304        self.version_hash = version_hash;
305        self
306    }
307
308    pub fn disable_mint_fees(mut self) -> FederationTestBuilder {
309        self.enable_mint_fees = false;
310        self
311    }
312
313    #[allow(clippy::too_many_lines)]
314    pub async fn build(self) -> FederationTest {
315        install_crypto_provider().await;
316        let num_offline = self.num_offline;
317        assert!(
318            self.num_peers > 3 * self.num_offline,
319            "too many peers offline ({num_offline}) to reach consensus"
320        );
321        let peers = (0..self.num_peers).map(PeerId::from).collect::<Vec<_>>();
322        let params = local_config_gen_params(
323            &peers,
324            self.base_port,
325            self.enable_mint_fees,
326            &self.server_init,
327        )
328        .expect("Generates local config");
329
330        let configs =
331            ServerConfig::trusted_dealer_gen(&params, &self.server_init, &self.version_hash);
332
333        let task_group = TaskGroup::new();
334        for (peer_id, cfg) in configs.clone() {
335            let peer_port = self.base_port + u16::from(peer_id) * 3;
336
337            let p2p_bind = format!("127.0.0.1:{peer_port}").parse().unwrap();
338            let api_bind = format!("127.0.0.1:{}", peer_port + 1).parse().unwrap();
339            let ui_bind = format!("127.0.0.1:{}", peer_port + 2).parse().unwrap();
340
341            if u16::from(peer_id) >= self.num_peers - self.num_offline {
342                continue;
343            }
344
345            let instances = cfg.consensus.iter_module_instances();
346            let decoders = self.server_init.available_decoders(instances);
347            let db = Database::new(MemDatabase::new(), decoders);
348            let module_init_registry = self.server_init.clone();
349            let subgroup = task_group.make_subgroup();
350            let checkpoint_dir = tempfile::Builder::new().tempdir().unwrap().keep();
351            let code_version_str = env!("CARGO_PKG_VERSION");
352
353            let connector = TlsTcpConnector::new(
354                cfg.tls_config(),
355                p2p_bind,
356                cfg.local.p2p_endpoints.clone(),
357                cfg.local.identity,
358            )
359            .await
360            .into_dyn();
361
362            let (p2p_status_senders, p2p_status_receivers) = p2p_status_channels(connector.peers());
363
364            let connections = ReconnectP2PConnections::new(
365                cfg.local.identity,
366                connector,
367                &task_group,
368                p2p_status_senders,
369                None,
370            )
371            .into_dyn();
372
373            let bitcoin_rpc_connection = self.bitcoin_rpc_connection.clone();
374
375            task_group.spawn("fedimintd", move |_| async move {
376                Box::pin(consensus::run(
377                    ConnectorRegistry::build_from_testing_env().bind().await,
378                    Some(ApiAuth::new("pass".to_string())),
379                    Some(ApiAuth::new("pass".to_string())),
380                    connections,
381                    p2p_status_receivers,
382                    api_bind,
383                    None,
384                    vec![],
385                    cfg.clone(),
386                    db.clone(),
387                    module_init_registry,
388                    &subgroup,
389                    ApiSecrets::default(),
390                    checkpoint_dir,
391                    code_version_str.to_string(),
392                    String::new(),
393                    bitcoin_rpc_connection,
394                    ui_bind,
395                    Box::new(|_| axum::Router::new()),
396                    1,
397                    Duration::from_secs(3600),
398                    ConnectionLimits {
399                        max_connections: 1000,
400                        max_requests_per_connection: 100,
401                    },
402                    None,
403                ))
404                .await
405                .expect("Could not initialise consensus");
406            });
407        }
408
409        for (peer_id, config) in configs.clone() {
410            if u16::from(peer_id) >= self.num_peers - self.num_offline {
411                continue;
412            }
413
414            let connectors = ConnectorRegistry::build_from_testing_env().bind().await;
415            let api = DynGlobalApi::new_admin(
416                connectors,
417                peer_id,
418                config.consensus.api_endpoints()[&peer_id].url.clone(),
419                None,
420            );
421
422            while let Err(e) = api
423                .request_admin_no_auth::<u64>(SESSION_COUNT_ENDPOINT, ApiRequestErased::default())
424                .await
425            {
426                sleep_in_test(
427                    format!("Waiting for api of peer {peer_id} to come online: {e}"),
428                    Duration::from_millis(500),
429                )
430                .await;
431            }
432        }
433
434        FederationTest {
435            configs,
436            server_init: self.server_init,
437            client_init: self.client_init,
438            _task: task_group,
439            num_peers: self.num_peers,
440            num_offline: self.num_offline,
441            connectors: ConnectorRegistry::build_from_testing_env().bind().await,
442        }
443    }
444}