Add custom DNS tunnel command

6 files changed, 51 insertions, 0 deletions

diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index b561fc7a83..c4b7458960 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -172,6 +172,31 @@ impl ConnectedState {
                     }
                 }
             }
+            Some(TunnelCommand::CustomDns(servers)) => {
+                if shared_values.custom_dns != servers {
+                    shared_values.custom_dns = servers;
+
+                    if let Err(error) = self.set_firewall_policy(shared_values) {
+                        return self.disconnect(
+                            shared_values,
+                            AfterDisconnect::Block(ErrorStateCause::SetFirewallPolicyError(error)),
+                        );
+                    }
+
+                    match self.set_dns(shared_values) {
+                        Ok(()) => SameState(self.into()),
+                        Err(error) => {
+                            log::error!("{}", error.display_chain_with_msg("Failed to set DNS"));
+                            self.disconnect(
+                                shared_values,
+                                AfterDisconnect::Block(ErrorStateCause::SetDnsError),
+                            )
+                        }
+                    }
+                } else {
+                    SameState(self.into())
+                }
+            }
             Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                 shared_values.block_when_disconnected = block_when_disconnected;
                 SameState(self.into())
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index bfccac6572..8fde71cc98 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -227,6 +227,10 @@ impl ConnectingState {
                     }
                 }
             }
+            Some(TunnelCommand::CustomDns(servers)) => {
+                shared_values.custom_dns = servers;
+                SameState(self.into())
+            }
             Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                 shared_values.block_when_disconnected = block_when_disconnected;
                 SameState(self.into())
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index faeac0a45f..685925b922 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -82,6 +82,10 @@ impl TunnelState for DisconnectedState {
                 }
                 SameState(self.into())
             }
+            Some(TunnelCommand::CustomDns(servers)) => {
+                shared_values.custom_dns = servers;
+                SameState(self.into())
+            }
             Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                 if shared_values.block_when_disconnected != block_when_disconnected {
                     shared_values.block_when_disconnected = block_when_disconnected;
diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index 33a09ca31a..9ecdd6dc88 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -32,6 +32,10 @@ impl DisconnectingState {
                     let _ = shared_values.set_allow_lan(allow_lan);
                     AfterDisconnect::Nothing
                 }
+                Some(TunnelCommand::CustomDns(servers)) => {
+                    shared_values.custom_dns = servers;
+                    AfterDisconnect::Nothing
+                }
                 Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                     shared_values.block_when_disconnected = block_when_disconnected;
                     AfterDisconnect::Nothing
@@ -49,6 +53,10 @@ impl DisconnectingState {
                     let _ = shared_values.set_allow_lan(allow_lan);
                     AfterDisconnect::Block(reason)
                 }
+                Some(TunnelCommand::CustomDns(servers)) => {
+                    shared_values.custom_dns = servers;
+                    AfterDisconnect::Block(reason)
+                }
                 Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                     shared_values.block_when_disconnected = block_when_disconnected;
                     AfterDisconnect::Block(reason)
@@ -71,6 +79,10 @@ impl DisconnectingState {
                     let _ = shared_values.set_allow_lan(allow_lan);
                     AfterDisconnect::Reconnect(retry_attempt)
                 }
+                Some(TunnelCommand::CustomDns(servers)) => {
+                    shared_values.custom_dns = servers;
+                    AfterDisconnect::Reconnect(retry_attempt)
+                }
                 Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                     shared_values.block_when_disconnected = block_when_disconnected;
                     AfterDisconnect::Reconnect(retry_attempt)
diff --git a/talpid-core/src/tunnel_state_machine/error_state.rs b/talpid-core/src/tunnel_state_machine/error_state.rs
index a9861e788e..bf545258fe 100644
--- a/talpid-core/src/tunnel_state_machine/error_state.rs
+++ b/talpid-core/src/tunnel_state_machine/error_state.rs
@@ -102,6 +102,10 @@ impl TunnelState for ErrorState {
                     SameState(self.into())
                 }
             }
+            Some(TunnelCommand::CustomDns(servers)) => {
+                shared_values.custom_dns = servers;
+                SameState(self.into())
+            }
             Some(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                 shared_values.block_when_disconnected = block_when_disconnected;
                 SameState(self.into())
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 1793093580..21119aa6ad 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -150,6 +150,8 @@ pub async fn spawn(
 pub enum TunnelCommand {
     /// Enable or disable LAN access in the firewall.
     AllowLan(bool),
+    /// Set custom DNS servers to use.
+    CustomDns(Option<Vec<IpAddr>>),
     /// Enable or disable the block_when_disconnected feature.
     BlockWhenDisconnected(bool),
     /// Notify the state machine of the connectivity of the device.