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