1
- use super :: { config:: TEST_CONFIG , Error , WAIT_FOR_TUNNEL_STATE_TIMEOUT } ;
1
+ use super :: { config:: TEST_CONFIG , Error , TestContext , WAIT_FOR_TUNNEL_STATE_TIMEOUT } ;
2
2
use crate :: network_monitor:: {
3
3
self , start_packet_monitor, MonitorOptions , MonitorUnexpectedlyStopped , PacketMonitor ,
4
4
} ;
@@ -22,6 +22,9 @@ use std::{
22
22
use talpid_types:: net:: wireguard:: { PeerConfig , PrivateKey , TunnelConfig } ;
23
23
use test_rpc:: { package:: Package , AmIMullvad , ServiceClient } ;
24
24
25
+ /// TODO: Document
26
+ pub const THROTTLE_RETRY_DELAY : Duration = Duration :: from_secs ( 120 ) ;
27
+
25
28
#[ macro_export]
26
29
macro_rules! assert_tunnel_state {
27
30
( $mullvad_client: expr, $pattern: pat) => { {
@@ -199,19 +202,44 @@ pub async fn ping_sized_with_timeout(
199
202
. map_err ( Error :: Rpc )
200
203
}
201
204
205
+ /// Log in and retry if it fails due to throttling
206
+ pub async fn login_with_retries (
207
+ mullvad_client : & mut MullvadProxyClient ,
208
+ ) -> Result < ( ) , mullvad_management_interface:: Error > {
209
+ loop {
210
+ match mullvad_client
211
+ . login_account ( TEST_CONFIG . account_number . clone ( ) )
212
+ . await
213
+ {
214
+ Err ( mullvad_management_interface:: Error :: Rpc ( status) )
215
+ if status. message ( ) . to_uppercase ( ) . contains ( "THROTTLED" ) =>
216
+ {
217
+ // Work around throttling errors by sleeping
218
+ log:: debug!(
219
+ "Login failed due to throttling. Sleeping for {} seconds" ,
220
+ THROTTLE_RETRY_DELAY . as_secs( )
221
+ ) ;
222
+
223
+ tokio:: time:: sleep ( THROTTLE_RETRY_DELAY ) . await ;
224
+ }
225
+ Err ( err) => break Err ( err) ,
226
+ Ok ( _) => break Ok ( ( ) ) ,
227
+ }
228
+ }
229
+ }
230
+
202
231
/// Try to connect to a Mullvad Tunnel.
203
232
///
204
- /// If that fails to begin to connect, [`Error::DaemonError`] is returned. If it fails to connect
205
- /// after that, the daemon ends up in the [`TunnelState::Error`] state, and
206
- /// [`Error::UnexpectedErrorState`] is returned.
233
+ /// # Returns
234
+ /// - `Result::Ok` if the daemon successfully connected to a tunnel
235
+ /// - `Result::Err` if:
236
+ /// - The daemon failed to even begin connecting. Then [`Error::Rpc`] is returned.
237
+ /// - The daemon started to connect but ended up in the [`TunnelState::Error`] state.
238
+ /// Then [`Error::UnexpectedErrorState`] is returned
207
239
pub async fn connect_and_wait ( mullvad_client : & mut MullvadProxyClient ) -> Result < ( ) , Error > {
208
240
log:: info!( "Connecting" ) ;
209
241
210
- mullvad_client
211
- . connect_tunnel ( )
212
- . await
213
- . map_err ( |error| Error :: Daemon ( format ! ( "failed to begin connecting: {}" , error) ) ) ?;
214
-
242
+ mullvad_client. connect_tunnel ( ) . await ?;
215
243
let new_state = wait_for_tunnel_state ( mullvad_client. clone ( ) , |state| {
216
244
matches ! (
217
245
state,
@@ -231,11 +259,8 @@ pub async fn connect_and_wait(mullvad_client: &mut MullvadProxyClient) -> Result
231
259
232
260
pub async fn disconnect_and_wait ( mullvad_client : & mut MullvadProxyClient ) -> Result < ( ) , Error > {
233
261
log:: info!( "Disconnecting" ) ;
262
+ mullvad_client. disconnect_tunnel ( ) . await ?;
234
263
235
- mullvad_client
236
- . disconnect_tunnel ( )
237
- . await
238
- . map_err ( |error| Error :: Daemon ( format ! ( "failed to begin disconnecting: {}" , error) ) ) ?;
239
264
wait_for_tunnel_state ( mullvad_client. clone ( ) , |state| {
240
265
matches ! ( state, TunnelState :: Disconnected { .. } )
241
266
} )
@@ -308,6 +333,30 @@ where
308
333
}
309
334
}
310
335
336
+ /// Set environment variables specified by `env` and restart the Mullvad daemon.
337
+ /// Returns a new [rpc client][`MullvadProxyClient`], since the old client *probably*
338
+ /// can't communicate with the new daemon.
339
+ ///
340
+ /// # Note
341
+ /// This is just a thin wrapper around [`ServiceClient::set_daemon_environment`] which also
342
+ /// invalides the old [`MullvadProxyClient`].
343
+ pub async fn restart_daemon_with < K , V , Env > (
344
+ rpc : & ServiceClient ,
345
+ test_context : & TestContext ,
346
+ _: MullvadProxyClient , // Just consume the old proxy client
347
+ env : Env ,
348
+ ) -> Result < MullvadProxyClient , Error >
349
+ where
350
+ Env : IntoIterator < Item = ( K , V ) > ,
351
+ K : Into < String > ,
352
+ V : Into < String > ,
353
+ {
354
+ rpc. set_daemon_environment ( env) . await ?;
355
+ // Need to create a new `mullvad_client` here after the restart
356
+ // otherwise we *probably* can't communicate with the daemon.
357
+ Ok ( test_context. rpc_provider . new_client ( ) . await )
358
+ }
359
+
311
360
pub async fn geoip_lookup_with_retries ( rpc : & ServiceClient ) -> Result < AmIMullvad , Error > {
312
361
const MAX_ATTEMPTS : usize = 5 ;
313
362
const BEFORE_RETRY_DELAY : Duration = Duration :: from_secs ( 2 ) ;
@@ -409,6 +458,11 @@ pub fn unreachable_wireguard_tunnel() -> talpid_types::net::wireguard::Connectio
409
458
}
410
459
}
411
460
461
+ /// Return the current `MULLVAD_API_HOST` et al.
462
+ ///
463
+ /// # Note
464
+ /// This is independent of the running daemon's environment.
465
+ /// It is solely dependant on the current value of [`TEST_CONFIG`].
412
466
pub fn get_app_env ( ) -> HashMap < String , String > {
413
467
use mullvad_api:: env;
414
468
use std:: net:: ToSocketAddrs ;
@@ -420,10 +474,10 @@ pub fn get_app_env() -> HashMap<String, String> {
420
474
. next ( )
421
475
. unwrap ( ) ;
422
476
423
- let api_host_env = ( env :: API_HOST_VAR . to_string ( ) , api_host ) ;
424
- let api_addr_env = ( env:: API_ADDR_VAR . to_string ( ) , api_addr . to_string ( ) ) ;
425
-
426
- [ api_host_env , api_addr_env ] . into_iter ( ) . collect ( )
477
+ HashMap :: from_iter ( vec ! [
478
+ ( env:: API_HOST_VAR . to_string( ) , api_host ) ,
479
+ ( env :: API_ADDR_VAR . to_string ( ) , api_addr . to_string ( ) ) ,
480
+ ] )
427
481
}
428
482
429
483
/// Return a filtered version of the daemon's relay list.
0 commit comments