Add server ownership constraint to the daemon and CLI

2 files changed, 86 insertions, 4 deletions

diff --git a/mullvad-cli/src/cmds/bridge.rs b/mullvad-cli/src/cmds/bridge.rs
index 40afc93b3b..3f00a7f69a 100644
--- a/mullvad-cli/src/cmds/bridge.rs
+++ b/mullvad-cli/src/cmds/bridge.rs
@@ -59,6 +59,19 @@ fn create_bridge_set_subcommand() -> clap::App<'static> {
                         .required(true),
                 ),
         )
+        .subcommand(
+            clap::App::new("ownership")
+                .about(
+                    "Filters bridges based on ownership. The 'list' \
+                       command shows the available relays and whether they're rented.",
+                )
+                .arg(
+                    clap::Arg::new("ownership")
+                        .help("Ownership preference, or 'any' for no preference.")
+                        .possible_values(&["any", "owned", "rented"])
+                        .required(true),
+                ),
+        )
         .subcommand(location::get_subcommand().about(
             "Set country or city to select bridge relays from. Use the 'list' \
              command to show available alternatives.",
@@ -186,6 +199,9 @@ impl Bridge {
             Some(("provider", provider_matches)) => {
                 Self::handle_set_bridge_provider(provider_matches).await
             }
+            Some(("ownership", ownership_matches)) => {
+                Self::handle_set_bridge_ownership(ownership_matches).await
+            }
             Some(("custom", custom_matches)) => {
                 Self::handle_bridge_set_custom_settings(custom_matches).await
             }
@@ -220,7 +236,12 @@ impl Bridge {
     }
 
     async fn handle_set_bridge_location(matches: &clap::ArgMatches) -> Result<()> {
-        Self::update_bridge_settings(Some(location::get_constraint_from_args(matches)), None).await
+        Self::update_bridge_settings(
+            Some(location::get_constraint_from_args(matches)),
+            None,
+            None,
+        )
+        .await
     }
 
     async fn handle_set_bridge_provider(matches: &clap::ArgMatches) -> Result<()> {
@@ -231,12 +252,19 @@ impl Bridge {
             providers
         };
 
-        Self::update_bridge_settings(None, Some(providers)).await
+        Self::update_bridge_settings(None, Some(providers), None).await
+    }
+
+    async fn handle_set_bridge_ownership(matches: &clap::ArgMatches) -> Result<()> {
+        let ownership =
+            super::relay::parse_ownership_constraint(matches.value_of("ownership").unwrap());
+        Self::update_bridge_settings(None, None, Some(ownership)).await
     }
 
     async fn update_bridge_settings(
         location: Option<types::RelayLocation>,
         providers: Option<Vec<String>>,
+        ownership: Option<types::Ownership>,
     ) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
         let settings = rpc.get_settings(()).await?.into_inner();
@@ -251,6 +279,9 @@ impl Bridge {
                     constraints.providers =
                         types::try_providers_constraint_from_proto(&new_providers).unwrap();
                 }
+                if let Some(new_ownership) = ownership {
+                    constraints.ownership = types::ownership_constraint_from_proto(new_ownership);
+                }
                 constraints
             }
             _ => {
@@ -258,10 +289,14 @@ impl Bridge {
                 let providers =
                     types::try_providers_constraint_from_proto(&providers.unwrap_or_default())
                         .unwrap();
+                let ownership = ownership
+                    .map(types::ownership_constraint_from_proto)
+                    .unwrap_or_default();
 
                 BridgeConstraints {
                     location,
                     providers,
+                    ownership,
                 }
             }
         };
@@ -433,8 +468,13 @@ impl Bridge {
                     city.name, city.code, city.latitude, city.longitude
                 );
                 for relay in &city.relays {
+                    let ownership = if relay.owned {
+                        "Mullvad-owned"
+                    } else {
+                        "rented"
+                    };
                     println!(
-                        "\t\t{} ({}) - hosted by {}",
+                        "\t\t{} ({}) - hosted by {} ({ownership})",
                         relay.hostname, relay.ipv4_addr_in, relay.provider
                     );
                 }
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index c11ee82158..d0d75dea20 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -124,6 +124,17 @@ impl Command for Relay {
                             )
                     )
                     .subcommand(
+                        clap::App::new("ownership")
+                            .about("Filters relays based on ownership. The 'list' \
+                                   command shows the available relays and whether they're rented.")
+                            .arg(
+                                clap::Arg::new("ownership")
+                                .help("Ownership preference, or 'any' for no preference.")
+                                .possible_values(&["any", "owned", "rented"])
+                                .required(true)
+                            )
+                    )
+                    .subcommand(
                         clap::App::new("tunnel")
                             .about("Set tunnel protocol-specific constraints.")
                             .setting(clap::AppSettings::SubcommandRequiredElseHelp)
@@ -226,6 +237,8 @@ impl Relay {
             self.set_hostname(relay_matches).await
         } else if let Some(providers_matches) = matches.subcommand_matches("provider") {
             self.set_providers(providers_matches).await
+        } else if let Some(ownership_matches) = matches.subcommand_matches("ownership") {
+            self.set_ownership(ownership_matches).await
         } else if let Some(matches) = matches.subcommand_matches("tunnel") {
             if let Some(tunnel_matches) = matches.subcommand_matches("openvpn") {
                 self.set_openvpn_constraints(tunnel_matches).await
@@ -483,6 +496,21 @@ impl Relay {
         .await
     }
 
+    async fn set_ownership(&self, matches: &clap::ArgMatches) -> Result<()> {
+        let ownership = parse_ownership_constraint(matches.value_of("ownership").unwrap());
+        self.update_constraints(types::RelaySettingsUpdate {
+            r#type: Some(types::relay_settings_update::Type::Normal(
+                types::NormalRelaySettingsUpdate {
+                    ownership: Some(types::OwnershipUpdate {
+                        ownership: ownership as i32,
+                    }),
+                    ..Default::default()
+                },
+            )),
+        })
+        .await
+    }
+
     async fn set_openvpn_constraints(&self, matches: &clap::ArgMatches) -> Result<()> {
         let mut openvpn_constraints = {
             let mut rpc = new_rpc_client().await?;
@@ -646,12 +674,17 @@ impl Relay {
                         (false, true) => "WireGuard",
                         _ => unreachable!("Bug in relay filtering earlier on"),
                     };
+                    let ownership = if relay.owned {
+                        "Mullvad-owned"
+                    } else {
+                        "rented"
+                    };
                     let mut addresses = vec![&relay.ipv4_addr_in];
                     if !relay.ipv6_addr_in.is_empty() {
                         addresses.push(&relay.ipv6_addr_in);
                     }
                     println!(
-                        "\t\t{} ({}) - {}, hosted by {}",
+                        "\t\t{} ({}) - {}, hosted by {} ({ownership})",
                         relay.hostname,
                         addresses.iter().join(", "),
                         support_msg,
@@ -801,3 +834,12 @@ fn parse_transport_port(
         )),
     }
 }
+
+pub fn parse_ownership_constraint(constraint: &str) -> types::Ownership {
+    match constraint {
+        "any" => types::Ownership::Any,
+        "owned" => types::Ownership::MullvadOwned,
+        "rented" => types::Ownership::Rented,
+        _ => unreachable!(),
+    }
+}