Separate DNS options and toggle

5 files changed, 111 insertions, 80 deletions

diff --git a/mullvad-cli/src/cmds/dns.rs b/mullvad-cli/src/cmds/dns.rs
index 91eb7854f6..9ffa0ec09c 100644
--- a/mullvad-cli/src/cmds/dns.rs
+++ b/mullvad-cli/src/cmds/dns.rs
@@ -1,7 +1,7 @@
 use crate::{new_rpc_client, Command, Result};
 use mullvad_management_interface::types;
-use mullvad_types::settings::DnsOptions;
-use std::convert::TryFrom;
+use mullvad_types::settings::{DnsOptions, DnsState};
+use std::convert::TryInto;
 
 pub struct Dns;
 
@@ -75,13 +75,14 @@ impl Command for Dns {
 impl Dns {
     async fn set_default(&self, block_ads: bool, block_trackers: bool) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
+        let settings = rpc.get_settings(()).await?.into_inner();
         rpc.set_dns_options(types::DnsOptions {
-            r#type: Some(types::dns_options::Type::Default(
-                types::DefaultDnsOptions {
-                    block_ads,
-                    block_trackers,
-                },
-            )),
+            state: types::dns_options::DnsState::Default as i32,
+            default_options: Some(types::DefaultDnsOptions {
+                block_ads,
+                block_trackers,
+            }),
+            ..settings.tunnel_options.unwrap().dns_options.unwrap()
         })
         .await?;
         println!("Updated DNS settings");
@@ -90,10 +91,13 @@ impl Dns {
 
     async fn set_custom(&self, servers: Option<Vec<String>>) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
+        let settings = rpc.get_settings(()).await?.into_inner();
         rpc.set_dns_options(types::DnsOptions {
-            r#type: Some(types::dns_options::Type::Custom(types::CustomDnsOptions {
+            state: types::dns_options::DnsState::Custom as i32,
+            custom_options: Some(types::CustomDnsOptions {
                 addresses: servers.unwrap_or_default(),
-            })),
+            }),
+            ..settings.tunnel_options.unwrap().dns_options.unwrap()
         })
         .await?;
         println!("Updated DNS settings");
@@ -102,24 +106,26 @@ impl Dns {
 
     async fn get(&self) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        let options = rpc
+        let options: DnsOptions = rpc
             .get_settings(())
             .await?
             .into_inner()
             .tunnel_options
             .unwrap()
             .dns_options
+            .unwrap()
+            .try_into()
             .unwrap();
 
-        match DnsOptions::try_from(options).unwrap() {
-            DnsOptions::Default(options) => {
+        match options.state {
+            DnsState::Default => {
                 println!("Custom DNS: no");
-                println!("Block ads: {}", options.block_ads);
-                println!("Block trackers: {}", options.block_trackers);
+                println!("Block ads: {}", options.default_options.block_ads);
+                println!("Block trackers: {}", options.default_options.block_trackers);
             }
-            DnsOptions::Custom(options) => {
+            DnsState::Custom => {
                 println!("Custom DNS: yes\nServers:");
-                for server in &options.addresses {
+                for server in &options.custom_options.addresses {
                     println!("{}", server);
                 }
             }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 01635186cd..d44ce55153 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -35,7 +35,7 @@ use mullvad_types::{
         RelaySettingsUpdate,
     },
     relay_list::{Relay, RelayList},
-    settings::{DnsOptions, Settings},
+    settings::{DnsOptions, DnsState, Settings},
     states::{TargetState, TunnelState},
     version::{AppVersion, AppVersionInfo},
     wireguard::{KeygenEvent, RotationInterval},
@@ -700,26 +700,26 @@ where
         Ok(daemon)
     }
 
-    fn get_dns_resolvers(dns_options: &DnsOptions) -> Option<Vec<IpAddr>> {
-        match dns_options {
-            DnsOptions::Default(options) => {
-                if options.block_ads {
-                    if options.block_trackers {
+    fn get_dns_resolvers(options: &DnsOptions) -> Option<Vec<IpAddr>> {
+        match options.state {
+            DnsState::Default => {
+                if options.default_options.block_ads {
+                    if options.default_options.block_trackers {
                         Some(DNS_AD_TRACKER_BLOCKING_SERVERS.to_vec())
                     } else {
                         Some(DNS_AD_BLOCKING_SERVERS.to_vec())
                     }
-                } else if options.block_trackers {
+                } else if options.default_options.block_trackers {
                     Some(DNS_TRACKER_BLOCKING_SERVERS.to_vec())
                 } else {
                     None
                 }
             }
-            DnsOptions::Custom(options) => {
-                if options.addresses.is_empty() {
+            DnsState::Custom => {
+                if options.custom_options.addresses.is_empty() {
                     None
                 } else {
-                    Some(options.addresses.clone())
+                    Some(options.custom_options.addresses.clone())
                 }
             }
         }
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index a5b63794e1..41ba8c96a9 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -386,10 +386,13 @@ message CustomDnsOptions {
 }
 
 message DnsOptions {
-	oneof type {
-		DefaultDnsOptions default = 1;
-		CustomDnsOptions custom = 2;
+	enum DnsState {
+		DEFAULT = 0;
+		CUSTOM = 1;
 	}
+	DnsState state = 1;
+	DefaultDnsOptions default_options = 2;
+	CustomDnsOptions custom_options = 3;
 }
 
 message PublicKey {
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index e7d0fd8063..aff3a30789 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -474,26 +474,23 @@ impl From<mullvad_types::relay_constraints::RelaySettings> for RelaySettings {
 
 impl From<&mullvad_types::settings::DnsOptions> for DnsOptions {
     fn from(options: &mullvad_types::settings::DnsOptions) -> Self {
-        let inner_options = match options {
-            mullvad_types::settings::DnsOptions::Default(options) => {
-                dns_options::Type::Default(DefaultDnsOptions {
-                    block_ads: options.block_ads,
-                    block_trackers: options.block_trackers,
-                })
-            }
-            mullvad_types::settings::DnsOptions::Custom(options) => {
-                dns_options::Type::Custom(CustomDnsOptions {
-                    addresses: options
-                        .addresses
-                        .iter()
-                        .map(|addr| addr.to_string())
-                        .collect(),
-                })
-            }
-        };
-
         DnsOptions {
-            r#type: Some(inner_options),
+            state: match options.state {
+                mullvad_types::settings::DnsState::Default => dns_options::DnsState::Default as i32,
+                mullvad_types::settings::DnsState::Custom => dns_options::DnsState::Custom as i32,
+            },
+            default_options: Some(DefaultDnsOptions {
+                block_ads: options.default_options.block_ads,
+                block_trackers: options.default_options.block_trackers,
+            }),
+            custom_options: Some(CustomDnsOptions {
+                addresses: options
+                    .custom_options
+                    .addresses
+                    .iter()
+                    .map(|addr| addr.to_string())
+                    .collect(),
+            }),
         }
     }
 }
@@ -1034,32 +1031,50 @@ impl TryFrom<DnsOptions> for mullvad_types::settings::DnsOptions {
         use mullvad_types::settings::{
             CustomDnsOptions as MullvadCustomDnsOptions,
             DefaultDnsOptions as MullvadDefaultDnsOptions, DnsOptions as MullvadDnsOptions,
+            DnsState as MullvadDnsState,
         };
 
-        match options.r#type {
-            Some(dns_options::Type::Default(options)) => {
-                Ok(MullvadDnsOptions::Default(MullvadDefaultDnsOptions {
-                    block_ads: options.block_ads,
-                    block_trackers: options.block_trackers,
-                }))
+        let state = match dns_options::DnsState::from_i32(options.state) {
+            Some(dns_options::DnsState::Default) => MullvadDnsState::Default,
+            Some(dns_options::DnsState::Custom) => MullvadDnsState::Custom,
+            None => {
+                return Err(FromProtobufTypeError::InvalidArgument(
+                    "invalid DNS options state",
+                ))
             }
-            Some(dns_options::Type::Custom(options)) => {
-                Ok(MullvadDnsOptions::Custom(MullvadCustomDnsOptions {
-                    addresses: options
-                        .addresses
-                        .into_iter()
-                        .map(|addr| {
-                            addr.parse().map_err(|_| {
-                                FromProtobufTypeError::InvalidArgument("invalid IP address")
-                            })
+        };
+
+        let default_options =
+            options
+                .default_options
+                .ok_or(FromProtobufTypeError::InvalidArgument(
+                    "missing default DNS options",
+                ))?;
+        let custom_options =
+            options
+                .custom_options
+                .ok_or(FromProtobufTypeError::InvalidArgument(
+                    "missing default DNS options",
+                ))?;
+
+        Ok(MullvadDnsOptions {
+            state,
+            default_options: MullvadDefaultDnsOptions {
+                block_ads: default_options.block_ads,
+                block_trackers: default_options.block_trackers,
+            },
+            custom_options: MullvadCustomDnsOptions {
+                addresses: custom_options
+                    .addresses
+                    .into_iter()
+                    .map(|addr| {
+                        addr.parse().map_err(|_| {
+                            FromProtobufTypeError::InvalidArgument("invalid IP address")
                         })
-                        .collect::<Result<Vec<_>, _>>()?,
-                }))
-            }
-            None => Err(FromProtobufTypeError::InvalidArgument(
-                "invalid DNS setting",
-            )),
-        }
+                    })
+                    .collect::<Result<Vec<_>, _>>()?,
+            },
+        })
     }
 }
 
diff --git a/mullvad-types/src/settings/mod.rs b/mullvad-types/src/settings/mod.rs
index d615b835ba..b4d4c6ceb7 100644
--- a/mullvad-types/src/settings/mod.rs
+++ b/mullvad-types/src/settings/mod.rs
@@ -175,25 +175,32 @@ pub struct TunnelOptions {
     #[cfg_attr(target_os = "android", jnix(skip))]
     pub generic: GenericTunnelOptions,
     /// DNS options.
-    #[serde(default = "DnsOptions::default")]
     pub dns_options: DnsOptions,
 }
 
-/// DNS config
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
-#[cfg_attr(target_os = "android", derive(FromJava, IntoJava))]
-#[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
-pub enum DnsOptions {
-    Default(DefaultDnsOptions),
-    Custom(CustomDnsOptions),
+#[serde(rename_all = "snake_case")]
+pub enum DnsState {
+    Default,
+    Custom,
 }
 
-impl Default for DnsOptions {
+impl Default for DnsState {
     fn default() -> Self {
-        DnsOptions::Default(DefaultDnsOptions::default())
+        Self::Default
     }
 }
 
+/// DNS config
+#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
+#[cfg_attr(target_os = "android", derive(FromJava, IntoJava))]
+#[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
+pub struct DnsOptions {
+    pub state: DnsState,
+    pub default_options: DefaultDnsOptions,
+    pub custom_options: CustomDnsOptions,
+}
+
 /// Default DNS config
 #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
 pub struct DefaultDnsOptions {