Merge branch 'debug-rollout-cmds'

5 files changed, 185 insertions, 1 deletions

diff --git a/mullvad-cli/src/cmds/debug.rs b/mullvad-cli/src/cmds/debug.rs
index 933f1212e8..37c2d579e3 100644
--- a/mullvad-cli/src/cmds/debug.rs
+++ b/mullvad-cli/src/cmds/debug.rs
@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Result, bail};
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_types::{
     constraints::Constraint,
@@ -12,6 +12,9 @@ pub enum DebugCommands {
     /// Relay
     #[clap(subcommand)]
     Relay(RelayDebugCommands),
+    /// Handy commands for interacting with the app release rollout system.
+    #[clap(subcommand)]
+    Rollout(RolloutDebugCommands),
 }
 
 #[derive(clap::Subcommand, Debug)]
@@ -24,6 +27,18 @@ pub enum RelayDebugCommands {
     Enable { relay: String },
 }
 
+#[derive(clap::Subcommand, Debug)]
+pub enum RolloutDebugCommands {
+    /// Print your rollout threshold.
+    Get,
+    /// Generate a new rollout threshold (overwrites the current threshold value)
+    Reroll,
+    /// Set your rollout threshold seed to a known value.
+    ///
+    /// The seed is used to generate a rollout threshold.
+    Seed { value: u32 },
+}
+
 impl DebugCommands {
     pub async fn handle(self) -> Result<()> {
         match self {
@@ -66,6 +81,35 @@ impl DebugCommands {
                 println!("{relay} is now marked as active");
                 Ok(())
             }
+            DebugCommands::Rollout(rollout_cmd) => rollout_cmd.handle().await,
+        }
+    }
+}
+
+impl RolloutDebugCommands {
+    pub async fn handle(self) -> Result<()> {
+        let mut rpc = MullvadProxyClient::new().await?;
+        match self {
+            RolloutDebugCommands::Get => {
+                let Ok(threshold) = rpc.get_rollout_threshold().await else {
+                    bail!("Failed to get rollout");
+                };
+                println!("{threshold}");
+                Ok(())
+            }
+            RolloutDebugCommands::Reroll => {
+                let Ok(threshold) = rpc.generate_new_rollout_threshold().await else {
+                    bail!("Failed to get rollout");
+                };
+                println!("{threshold}");
+                Ok(())
+            }
+            RolloutDebugCommands::Seed { value: seed } => {
+                if rpc.set_new_rollout_threshold_seed(seed).await.is_err() {
+                    bail!("Failed to update rollout seed");
+                }
+                Ok(())
+            }
         }
     }
 }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index f1116e4875..db1528d917 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -430,6 +430,20 @@ pub enum DaemonCommand {
         relay: String,
         tx: oneshot::Sender<()>,
     },
+    /// Calculate and return the rollout threshold for this client.
+    #[cfg(not(target_os = "android"))]
+    GetRolloutThreshold(oneshot::Sender<f32>),
+    /// Generate a new rollout threshold seed and update settings. Returns the new rollout
+    /// threshold.
+    #[cfg(not(target_os = "android"))]
+    GenerateNewRolloutSeed(oneshot::Sender<f32>),
+    /// Set the rollout threshold seed to the provided value and update settings.
+    #[cfg(not(target_os = "android"))]
+    SetRolloutThresholdSeed {
+        seed: u32,
+        tx: oneshot::Sender<()>,
+    },
+
     // App upgrade
     /// Prompt the daemon to start an app version upgrade.
     ///
@@ -1558,6 +1572,19 @@ impl Daemon {
             GetFeatureIndicators(tx) => self.on_get_feature_indicators(tx),
             DisableRelay { relay, tx } => self.on_toggle_relay(relay, false, tx),
             EnableRelay { relay, tx } => self.on_toggle_relay(relay, true, tx),
+            #[cfg(not(target_os = "android"))]
+            GetRolloutThreshold(tx) => self.on_get_rollout_threshold(tx).await,
+            #[cfg(not(target_os = "android"))]
+            GenerateNewRolloutSeed(tx) => {
+                let seed = self.generate_and_set().await;
+                let threshold = Self::calculate_rollout_threshold(seed);
+                let _ = tx.send(threshold);
+            }
+            #[cfg(not(target_os = "android"))]
+            SetRolloutThresholdSeed { seed, tx } => {
+                self.set_rollout_threshold_seed(seed).await;
+                let _ = tx.send(());
+            }
             AppUpgrade(tx) => self.on_app_upgrade(tx).await,
             AppUpgradeAbort(tx) => self.on_app_upgrade_abort(tx).await,
             GetAppUpgradeCacheDir(tx) => self.on_get_app_upgrade_cache_dir(tx).await,
@@ -3262,6 +3289,48 @@ impl Daemon {
         self.reconnect_tunnel();
     }
 
+    #[cfg(not(target_os = "android"))]
+    async fn on_get_rollout_threshold(&mut self, reply: oneshot::Sender<f32>) {
+        let seed = match self.settings.rollout_threshold_seed {
+            Some(seed) => seed,
+            None => self.generate_and_set().await,
+        };
+        let _ = reply.send(Self::calculate_rollout_threshold(seed));
+    }
+
+    #[cfg(not(target_os = "android"))]
+    fn calculate_rollout_threshold(seed: u32) -> f32 {
+        let version = mullvad_version::VERSION
+            .parse::<mullvad_version::Version>()
+            .expect("Failed to parse version");
+        let threshold = mullvad_update::version::Rollout::threshold(seed, version);
+        // a tiny bit hacky way to map Rollout -> f32, but it works.
+        threshold
+            .to_string()
+            .parse()
+            .expect("threshold is a valid Rollout is a valid f32")
+    }
+
+    // Regenrate a new seed and store it to settings.
+    #[cfg(not(target_os = "android"))]
+    async fn generate_and_set(&mut self) -> u32 {
+        let seed = generate_rollout_seed();
+        self.set_rollout_threshold_seed(seed).await;
+        seed
+    }
+
+    // Store the given seed to settings.
+    #[cfg(not(target_os = "android"))]
+    async fn set_rollout_threshold_seed(&mut self, seed: u32) {
+        if let Err(err) = self
+            .settings
+            .update(|settings| settings.rollout_threshold_seed = Some(seed))
+            .await
+        {
+            log::warn!("Failed to save settings when updating rollout seed: {err}");
+        }
+    }
+
     fn oneshot_send<T>(tx: oneshot::Sender<T>, t: T, msg: &'static str) {
         if tx.send(t).is_err() {
             log::warn!("Unable to send {} to the daemon command sender", msg);
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 4ed05df442..04511d902d 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -1168,6 +1168,51 @@ impl ManagementService for ManagementServiceImpl {
         Ok(Response::new(()))
     }
 
+    #[cfg(not(target_os = "android"))]
+    async fn get_rollout_threshold(&self, _: Request<()>) -> ServiceResult<types::Rollout> {
+        log::debug!("get_rollout_threshold");
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::GetRolloutThreshold(tx))?;
+        let threshold = self.wait_for_result(rx).await?;
+        let rollout = types::Rollout { threshold };
+        Ok(Response::new(rollout))
+    }
+
+    #[cfg(not(target_os = "android"))]
+    async fn set_rollout_threshold_seed(&self, seed: Request<types::Seed>) -> ServiceResult<()> {
+        log::debug!("set_rollout_threshold_seed");
+        let seed = seed.into_inner().seed;
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::SetRolloutThresholdSeed { seed, tx })?;
+        self.wait_for_result(rx).await?;
+        Ok(Response::new(()))
+    }
+
+    #[cfg(not(target_os = "android"))]
+    async fn regenerate_rollout_threshold(&self, _: Request<()>) -> ServiceResult<types::Rollout> {
+        log::debug!("regenerate_rollout_threshold");
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::GenerateNewRolloutSeed(tx))?;
+        let threshold = self.wait_for_result(rx).await?;
+        let rollout = types::Rollout { threshold };
+        Ok(Response::new(rollout))
+    }
+
+    #[cfg(target_os = "android")]
+    async fn get_rollout_threshold(&self, _: Request<()>) -> ServiceResult<types::Rollout> {
+        unreachable!("You should not call get_rollout_threshold");
+    }
+
+    #[cfg(target_os = "android")]
+    async fn set_rollout_threshold_seed(&self, _: Request<types::Seed>) -> ServiceResult<()> {
+        unreachable!("You should not call set_rollout_threshold_seed");
+    }
+
+    #[cfg(target_os = "android")]
+    async fn regenerate_rollout_threshold(&self, _: Request<()>) -> ServiceResult<types::Rollout> {
+        unreachable!("You should not call regenerate_rollout_threshold");
+    }
+
     // App upgrade
 
     async fn app_upgrade(&self, _: Request<()>) -> ServiceResult<()> {
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 4d9116e6c8..8556699929 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -136,6 +136,10 @@ service ManagementService {
   rpc DisableRelay(google.protobuf.StringValue) returns (google.protobuf.Empty) {}
   rpc EnableRelay(google.protobuf.StringValue) returns (google.protobuf.Empty) {}
 
+  rpc GetRolloutThreshold(google.protobuf.Empty) returns (Rollout) {}
+  rpc RegenerateRolloutThreshold(google.protobuf.Empty) returns (Rollout) {}
+  rpc SetRolloutThresholdSeed(Seed) returns (google.protobuf.Empty) {}
+
   // App upgrade
   rpc AppUpgrade(google.protobuf.Empty) returns (google.protobuf.Empty) {}
   rpc AppUpgradeAbort(google.protobuf.Empty) returns (google.protobuf.Empty) {}
@@ -172,6 +176,9 @@ message AppUpgradeError {
   Error error = 1;
 }
 
+message Seed { uint32 seed = 1; }
+message Rollout { float threshold = 1; }
+
 message UUID { string value = 1; }
 
 message AccountData {
diff --git a/mullvad-management-interface/src/client.rs b/mullvad-management-interface/src/client.rs
index aa8b67d930..abc5d34af3 100644
--- a/mullvad-management-interface/src/client.rs
+++ b/mullvad-management-interface/src/client.rs
@@ -635,6 +635,25 @@ impl MullvadProxyClient {
         Ok(())
     }
 
+    pub async fn get_rollout_threshold(&mut self) -> Result<f32> {
+        let rollout = self.0.get_rollout_threshold(()).await?;
+        let threshold = rollout.into_inner().threshold;
+        Ok(threshold)
+    }
+
+    pub async fn generate_new_rollout_threshold(&mut self) -> Result<f32> {
+        let rollout = self.0.regenerate_rollout_threshold(()).await?;
+        let threshold = rollout.into_inner().threshold;
+        Ok(threshold)
+    }
+
+    pub async fn set_new_rollout_threshold_seed(&mut self, seed: u32) -> Result<()> {
+        self.0
+            .set_rollout_threshold_seed(types::Seed { seed })
+            .await?;
+        Ok(())
+    }
+
     pub async fn set_wireguard_allowed_ips(&mut self, allowed_ips: AllowedIps) -> Result<()> {
         self.0
             .set_wireguard_allowed_ips(types::AllowedIpsList {