Add configurable transport protocol to local SOCKS5 access method

9 files changed, 113 insertions, 51 deletions

diff --git a/mullvad-api/src/lib.rs b/mullvad-api/src/lib.rs
index 91e2bc524a..267fdfd811 100644
--- a/mullvad-api/src/lib.rs
+++ b/mullvad-api/src/lib.rs
@@ -21,7 +21,7 @@ use std::{
     ops::Deref,
     path::Path,
 };
-use talpid_types::ErrorExt;
+use talpid_types::{net::Endpoint, ErrorExt};
 
 pub mod availability;
 use availability::{ApiAvailability, ApiAvailabilityHandle};
@@ -221,13 +221,13 @@ pub enum Error {
 
 /// Closure that receives the next API (real or proxy) endpoint to use for `api.mullvad.net`.
 /// It should return a future that determines whether to reject the new endpoint or not.
-pub trait ApiEndpointUpdateCallback: Fn(SocketAddr) -> Self::AcceptedNewEndpoint {
+pub trait ApiEndpointUpdateCallback: Fn(Endpoint) -> Self::AcceptedNewEndpoint {
     type AcceptedNewEndpoint: Future<Output = bool> + Send;
 }
 
 impl<U, T: Future<Output = bool> + Send> ApiEndpointUpdateCallback for U
 where
-    U: Fn(SocketAddr) -> T,
+    U: Fn(Endpoint) -> T,
 {
     type AcceptedNewEndpoint = T;
 }
diff --git a/mullvad-api/src/proxy.rs b/mullvad-api/src/proxy.rs
index 3193fdc19b..0727628596 100644
--- a/mullvad-api/src/proxy.rs
+++ b/mullvad-api/src/proxy.rs
@@ -4,12 +4,14 @@ use mullvad_types::access_method;
 use serde::{Deserialize, Serialize};
 use std::{
     fmt, io,
-    net::SocketAddr,
     path::Path,
     pin::Pin,
     task::{self, Poll},
 };
-use talpid_types::ErrorExt;
+use talpid_types::{
+    net::{Endpoint, TransportProtocol},
+    ErrorExt,
+};
 use tokio::{
     fs,
     io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf},
@@ -41,13 +43,17 @@ pub enum ProxyConfig {
 }
 
 impl ProxyConfig {
-    /// Returns the remote address to reach the proxy.
-    fn get_endpoint(&self) -> SocketAddr {
+    /// Returns the remote endpoint describing how to reach the proxy.
+    fn get_endpoint(&self) -> Endpoint {
         match self {
-            ProxyConfig::Shadowsocks(ss) => ss.peer,
+            ProxyConfig::Shadowsocks(shadowsocks) => {
+                Endpoint::from_socket_address(shadowsocks.peer, TransportProtocol::Tcp)
+            }
             ProxyConfig::Socks(socks) => match socks {
-                access_method::Socks5::Local(s) => s.remote_peer,
-                access_method::Socks5::Remote(s) => s.peer,
+                access_method::Socks5::Local(local) => local.remote_endpoint,
+                access_method::Socks5::Remote(remote) => {
+                    Endpoint::from_socket_address(remote.peer, TransportProtocol::Tcp)
+                }
             },
         }
     }
@@ -55,18 +61,14 @@ impl ProxyConfig {
 
 impl fmt::Display for ProxyConfig {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        let endpoint = self.get_endpoint();
         match self {
-            // TODO: Do not hardcode TCP
-            ProxyConfig::Shadowsocks(ss) => write!(f, "Shadowsocks {}/TCP", ss.peer),
+            ProxyConfig::Shadowsocks(_) => write!(f, "Shadowsocks {}", endpoint),
             ProxyConfig::Socks(socks) => match socks {
-                access_method::Socks5::Local(s) => {
-                    write!(
-                        f,
-                        "Socks5 {}/TCP via localhost:{}",
-                        s.remote_peer, s.local_port
-                    )
+                access_method::Socks5::Remote(_) => write!(f, "Socks5 {}", endpoint),
+                access_method::Socks5::Local(local) => {
+                    write!(f, "Socks5 {} via localhost:{}", endpoint, local.local_port)
                 }
-                access_method::Socks5::Remote(s) => write!(f, "Socks5 {}/TCP", s.peer),
             },
         }
     }
@@ -132,8 +134,9 @@ impl ApiConnectionMode {
         }
     }
 
-    /// Returns the remote address required to reach the API, or `None` for `ApiConnectionMode::Direct`.
-    pub fn get_endpoint(&self) -> Option<SocketAddr> {
+    /// Returns the remote endpoint required to reach the API, or `None` for
+    /// `ApiConnectionMode::Direct`.
+    pub fn get_endpoint(&self) -> Option<Endpoint> {
         match self {
             ApiConnectionMode::Direct => None,
             ApiConnectionMode::Proxied(proxy_config) => Some(proxy_config.get_endpoint()),
diff --git a/mullvad-api/src/rest.rs b/mullvad-api/src/rest.rs
index 63c909507c..a973c32f2a 100644
--- a/mullvad-api/src/rest.rs
+++ b/mullvad-api/src/rest.rs
@@ -24,7 +24,10 @@ use std::{
     sync::{Arc, Weak},
     time::Duration,
 };
-use talpid_types::ErrorExt;
+use talpid_types::{
+    net::{Endpoint, TransportProtocol},
+    ErrorExt,
+};
 
 #[cfg(feature = "api-override")]
 use crate::API;
@@ -209,7 +212,10 @@ impl<
                 if let Some(new_config) = self.proxy_config_provider.next().await {
                     let endpoint = match new_config.get_endpoint() {
                         Some(endpoint) => endpoint,
-                        None => self.address_cache.get_address().await,
+                        None => Endpoint::from_socket_address(
+                            self.address_cache.get_address().await,
+                            TransportProtocol::Tcp,
+                        ),
                     };
                     // Switch to new connection mode unless rejected by address change callback
                     if (self.new_address_callback)(endpoint).await {
diff --git a/mullvad-cli/src/cmds/api_access.rs b/mullvad-cli/src/cmds/api_access.rs
index 55c4a9b516..9ad481c4ef 100644
--- a/mullvad-cli/src/cmds/api_access.rs
+++ b/mullvad-cli/src/cmds/api_access.rs
@@ -4,7 +4,7 @@ use mullvad_types::access_method::{AccessMethod, AccessMethodSetting, CustomAcce
 use std::net::IpAddr;
 
 use clap::{Args, Subcommand};
-use talpid_types::net::openvpn::SHADOWSOCKS_CIPHERS;
+use talpid_types::net::{openvpn::SHADOWSOCKS_CIPHERS, TransportProtocol};
 
 #[derive(Subcommand, Debug, Clone)]
 pub enum ApiAccess {
@@ -118,10 +118,21 @@ impl ApiAccess {
                 }
                 CustomAccessMethod::Socks5(socks) => match socks {
                     Socks5::Local(local) => {
-                        let remote_ip = cmd.params.ip.unwrap_or(local.remote_peer.ip());
-                        let remote_port = cmd.params.port.unwrap_or(local.remote_peer.port());
+                        let remote_ip = cmd.params.ip.unwrap_or(local.remote_endpoint.address.ip());
+                        let remote_port = cmd
+                            .params
+                            .port
+                            .unwrap_or(local.remote_endpoint.address.port());
                         let local_port = cmd.params.local_port.unwrap_or(local.local_port);
-                        AccessMethod::from(Socks5Local::new((remote_ip, remote_port), local_port))
+                        let remote_peer_transport_protocol = cmd
+                            .params
+                            .transport_protocol
+                            .unwrap_or(local.remote_endpoint.protocol);
+                        AccessMethod::from(Socks5Local::new_with_transport_protocol(
+                            (remote_ip, remote_port),
+                            local_port,
+                            remote_peer_transport_protocol,
+                        ))
                     }
                     Socks5::Remote(remote) => {
                         let ip = cmd.params.ip.unwrap_or(remote.peer.ip());
@@ -306,6 +317,14 @@ pub enum AddSocks5Commands {
         remote_ip: IpAddr,
         /// The port of the remote peer
         remote_port: u16,
+        /// The Mullvad App can not know which transport protocol that the
+        /// remote peer accepts, but it needs to know this in order to correctly
+        /// exempt the connection traffic in the firewall.
+        ///
+        /// By default, the transport protocol is assumed to be `TCP`, but it
+        /// can optionally be set to `UDP` as well.
+        #[arg(long, default_value_t = TransportProtocol::Tcp)]
+        transport_protocol: TransportProtocol,
         /// Disable the use of this custom access method. It has to be manually
         /// enabled at a later stage to be used when accessing the Mullvad API.
         #[arg(default_value_t = false, short, long)]
@@ -398,6 +417,9 @@ pub struct EditParams {
     /// The port that the server on localhost is listening on [Socks5 (Local proxy)]
     #[arg(long)]
     local_port: Option<u16>,
+    /// The transport protocol used by the remote proxy [Socks5 (Local proxy)]
+    #[arg(long)]
+    transport_protocol: Option<TransportProtocol>,
 }
 
 /// Implement conversions from CLI types to Daemon types.
@@ -418,9 +440,15 @@ mod conversions {
                         remote_port,
                         name: _,
                         disabled: _,
+                        transport_protocol,
                     } => {
-                        println!("Adding SOCKS5-proxy: localhost:{local_port} => {remote_ip}:{remote_port}");
-                        daemon_types::Socks5Local::new((remote_ip, remote_port), local_port).into()
+                        println!("Adding SOCKS5-proxy: localhost:{local_port} => {remote_ip}:{remote_port}/{transport_protocol}");
+                        daemon_types::Socks5Local::new_with_transport_protocol(
+                            (remote_ip, remote_port),
+                            local_port,
+                            transport_protocol,
+                        )
+                        .into()
                     }
                     AddSocks5Commands::Remote {
                         remote_ip,
@@ -559,7 +587,13 @@ mod pp {
                             }
                             writeln!(f)?;
                             print_option!("Protocol", "Socks5 (local)");
-                            print_option!("Peer", local.remote_peer);
+                            print_option!(
+                                "Peer",
+                                format!(
+                                    "{}/{}",
+                                    local.remote_endpoint.address, local.remote_endpoint.protocol
+                                )
+                            );
                             print_option!("Local port", local.local_port);
                             Ok(())
                         }
diff --git a/mullvad-daemon/src/api.rs b/mullvad-daemon/src/api.rs
index c548f0a293..a6dd0dead6 100644
--- a/mullvad-daemon/src/api.rs
+++ b/mullvad-daemon/src/api.rs
@@ -12,7 +12,6 @@ use mullvad_api::{
 use mullvad_relay_selector::RelaySelector;
 use mullvad_types::access_method::{AccessMethod, AccessMethodSetting, BuiltInAccessMethod};
 use std::{
-    net::SocketAddr,
     path::PathBuf,
     pin::Pin,
     sync::{Arc, Mutex, Weak},
@@ -22,7 +21,7 @@ use std::{
 use talpid_core::mpsc::Sender;
 use talpid_core::tunnel_state_machine::TunnelCommand;
 use talpid_types::{
-    net::{openvpn::ProxySettings, AllowedEndpoint, Endpoint, TransportProtocol},
+    net::{openvpn::ProxySettings, AllowedEndpoint, Endpoint},
     ErrorExt,
 };
 
@@ -240,7 +239,7 @@ impl ApiEndpointUpdaterHandle {
 
     pub fn callback(&self) -> impl ApiEndpointUpdateCallback {
         let tunnel_tx = self.tunnel_cmd_tx.clone();
-        move |address: SocketAddr| {
+        move |endpoint: Endpoint| {
             let inner_tx = tunnel_tx.clone();
             async move {
                 let tunnel_tx = if let Some(tunnel_tx) = { inner_tx.lock().unwrap().as_ref() }
@@ -253,21 +252,19 @@ impl ApiEndpointUpdaterHandle {
                 };
                 let (result_tx, result_rx) = oneshot::channel();
                 let _ = tunnel_tx.unbounded_send(TunnelCommand::AllowEndpoint(
-                    get_allowed_endpoint(address),
+                    get_allowed_endpoint(endpoint),
                     result_tx,
                 ));
                 // Wait for the firewall policy to be updated.
                 let _ = result_rx.await;
-                log::debug!("API endpoint: {}", address);
+                log::debug!("API endpoint: {endpoint}");
                 true
             }
         }
     }
 }
 
-pub(super) fn get_allowed_endpoint(api_address: SocketAddr) -> AllowedEndpoint {
-    let endpoint = Endpoint::from_socket_address(api_address, TransportProtocol::Tcp);
-
+pub(super) fn get_allowed_endpoint(endpoint: Endpoint) -> AllowedEndpoint {
     #[cfg(windows)]
     let daemon_exe = std::env::current_exe().expect("failed to obtain executable path");
     #[cfg(windows)]
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 8ee8e58fd6..77304ae75f 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -726,7 +726,10 @@ where
         };
 
         let initial_api_endpoint =
-            api::get_allowed_endpoint(api_runtime.address_cache.get_address().await);
+            api::get_allowed_endpoint(talpid_types::net::Endpoint::from_socket_address(
+                api_runtime.address_cache.get_address().await,
+                talpid_types::net::TransportProtocol::Tcp,
+            ));
         let parameters_generator = tunnel::ParametersGenerator::new(
             account_manager.clone(),
             relay_selector.clone(),
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 057718bdb1..0eefd9a1db 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -335,7 +335,8 @@ message AccessMethod {
   message Socks5Local {
     string remote_ip = 1;
     uint32 remote_port = 2;
-    uint32 local_port = 3;
+    TransportProtocol remote_transport_protocol = 3;
+    uint32 local_port = 4;
   }
   message SocksAuth {
     string username = 1;
diff --git a/mullvad-management-interface/src/types/conversions/access_method.rs b/mullvad-management-interface/src/types/conversions/access_method.rs
index b4a4547fdf..425e0a70da 100644
--- a/mullvad-management-interface/src/types/conversions/access_method.rs
+++ b/mullvad-management-interface/src/types/conversions/access_method.rs
@@ -144,15 +144,19 @@ mod data {
         type Error = FromProtobufTypeError;
 
         fn try_from(value: proto::access_method::Socks5Local) -> Result<Self, Self::Error> {
+            use crate::types::conversions::net::try_transport_protocol_from_i32;
             let remote_ip = value.remote_ip.parse::<Ipv4Addr>().map_err(|_| {
                 FromProtobufTypeError::InvalidArgument(
                     "Could not parse Socks5 (local) message from protobuf",
                 )
             })?;
-            Ok(AccessMethod::from(Socks5Local::new(
-                (remote_ip, value.remote_port as u16),
-                value.local_port as u16,
-            )))
+            Ok(AccessMethod::from(
+                Socks5Local::new_with_transport_protocol(
+                    (remote_ip, value.remote_port as u16),
+                    value.local_port as u16,
+                    try_transport_protocol_from_i32(value.remote_transport_protocol)?,
+                ),
+            ))
         }
     }
 
@@ -229,13 +233,16 @@ mod data {
                     )
                 }
                 CustomAccessMethod::Socks5(Socks5::Local(Socks5Local {
-                    remote_peer: peer,
-                    local_port: port,
+                    remote_endpoint,
+                    local_port,
                 })) => proto::access_method::AccessMethod::Socks5local(
                     proto::access_method::Socks5Local {
-                        remote_ip: peer.ip().to_string(),
-                        remote_port: peer.port() as u32,
-                        local_port: port as u32,
+                        remote_ip: remote_endpoint.address.ip().to_string(),
+                        remote_port: remote_endpoint.address.port() as u32,
+                        remote_transport_protocol: i32::from(proto::TransportProtocol::from(
+                            remote_endpoint.protocol,
+                        )),
+                        local_port: local_port as u32,
                     },
                 ),
                 CustomAccessMethod::Socks5(Socks5::Remote(Socks5Remote {
diff --git a/mullvad-types/src/access_method.rs b/mullvad-types/src/access_method.rs
index f4a76392fa..b714a79edd 100644
--- a/mullvad-types/src/access_method.rs
+++ b/mullvad-types/src/access_method.rs
@@ -2,6 +2,7 @@ use std::str::FromStr;
 
 use serde::{Deserialize, Serialize};
 use std::net::SocketAddr;
+use talpid_types::net::{Endpoint, TransportProtocol};
 
 /// Daemon settings for API access methods.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -205,7 +206,7 @@ pub struct Shadowsocks {
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct Socks5Local {
-    pub remote_peer: SocketAddr,
+    pub remote_endpoint: Endpoint,
     /// Port on localhost where the SOCKS5-proxy listens to.
     pub local_port: u16,
 }
@@ -252,8 +253,18 @@ impl Shadowsocks {
 
 impl Socks5Local {
     pub fn new<I: Into<SocketAddr>>(remote_peer: I, local_port: u16) -> Self {
+        let transport_protocol = TransportProtocol::Tcp;
+        Self::new_with_transport_protocol(remote_peer, local_port, transport_protocol)
+    }
+
+    pub fn new_with_transport_protocol<I: Into<SocketAddr>>(
+        remote_peer: I,
+        local_port: u16,
+        transport_protocol: TransportProtocol,
+    ) -> Self {
+        let remote_endpoint = Endpoint::from_socket_address(remote_peer.into(), transport_protocol);
         Self {
-            remote_peer: remote_peer.into(),
+            remote_endpoint,
             local_port,
         }
     }