Use TransportPort for OpenVPN constraints

5 files changed, 112 insertions, 120 deletions

diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index c3d990403e..80ee09a6c9 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -13,8 +13,7 @@ use mullvad_management_interface::types::{
     relay_settings, relay_settings_update, ConnectionConfig, CustomRelaySettings, IpVersion,
     IpVersionConstraint, NormalRelaySettingsUpdate, OpenvpnConstraints, ProviderUpdate,
     RelayListCountry, RelayLocation, RelaySettingsUpdate, TransportPort, TransportProtocol,
-    TransportProtocolConstraint, TunnelType, TunnelTypeConstraint, TunnelTypeUpdate,
-    WireguardConstraints,
+    TunnelType, TunnelTypeConstraint, TunnelTypeUpdate, WireguardConstraints,
 };
 use mullvad_types::relay_constraints::Constraint;
 use talpid_types::net::all_of_the_internet;
@@ -149,13 +148,15 @@ impl Command for Relay {
                                     .arg(
                                         clap::Arg::with_name("port")
                                             .help("Port to use. Either 'any' or a specific port")
-                                            .required(true)
+                                            .long("port")
+                                            .default_value("any"),
                                     )
                                     .arg(
                                         clap::Arg::with_name("transport protocol")
+                                            .help("Transport protocol")
                                             .long("protocol")
-                                            .default_value("any")
-                                            .possible_values(&["any", "udp", "tcp"]),
+                                            .possible_values(&["any", "udp", "tcp"])
+                                            .default_value("any"),
                                     )
                             )
                             .subcommand(
@@ -165,17 +166,15 @@ impl Command for Relay {
                                         clap::Arg::with_name("port")
                                             .help("Port to use. Either 'any' or a specific port")
                                             .long("port")
-                                            .takes_value(true)
-                                            .requires("transport protocol"),
+                                            .default_value("any"),
                                     )
                                     .arg(
                                         clap::Arg::with_name("transport protocol")
                                             .help("Transport protocol. If TCP is selected, traffic is \
                                                    sent over TCP using a udp-over-tcp proxy")
                                             .long("protocol")
-                                            .takes_value(true)
-                                            .possible_values(&["udp", "tcp"])
-                                            .requires("port"),
+                                            .possible_values(&["any", "udp", "tcp"])
+                                            .default_value("any"),
                                     )
                                     .arg(
                                         clap::Arg::with_name("ip version")
@@ -516,20 +515,11 @@ impl Relay {
     }
 
     async fn set_openvpn_constraints(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
-        let port = parse_port_constraint(matches.value_of("port").unwrap())?;
-        let protocol = parse_protocol_constraint(matches.value_of("transport protocol").unwrap());
-
+        let port = parse_transport_port(matches)?;
         self.update_constraints(RelaySettingsUpdate {
             r#type: Some(relay_settings_update::Type::Normal(
                 NormalRelaySettingsUpdate {
-                    openvpn_constraints: Some(OpenvpnConstraints {
-                        port: port.unwrap_or(0) as u32,
-                        protocol: protocol
-                            .option()
-                            .map(|protocol| TransportProtocolConstraint {
-                                protocol: protocol as i32,
-                            }),
-                    }),
+                    openvpn_constraints: Some(OpenvpnConstraints { port }),
                     ..Default::default()
                 },
             )),
@@ -538,19 +528,7 @@ impl Relay {
     }
 
     async fn set_wireguard_constraints(&self, matches: &clap::ArgMatches<'_>) -> Result<()> {
-        let protocol = matches.value_of("transport protocol").map(parse_protocol);
-        let port = match matches.value_of("port").map(parse_port_constraint) {
-            None => None,
-            Some(Err(error)) => return Err(error),
-            Some(Ok(Constraint::Any)) => Some(TransportPort {
-                protocol: protocol.unwrap() as i32,
-                port: 0,
-            }),
-            Some(Ok(Constraint::Only(port))) => Some(TransportPort {
-                protocol: protocol.unwrap() as i32,
-                port: u32::from(port),
-            }),
-        };
+        let port = parse_transport_port(matches)?;
         let ip_version = parse_ip_version_constraint(matches.value_of("ip version").unwrap());
         let entry_location =
             parse_entry_location_constraint(matches.values_of("entry location").unwrap());
@@ -722,26 +700,12 @@ impl Relay {
         }
     }
 
-    fn format_port(port: u32) -> String {
-        if port != 0 {
-            format!("port {}", port)
-        } else {
-            "any port".to_string()
-        }
-    }
-
     fn format_openvpn_constraints(constraints: Option<&OpenvpnConstraints>) -> String {
         if let Some(constraints) = constraints {
-            format!(
-                "{} over {}",
-                Self::format_port(constraints.port),
-                Self::format_transport_protocol(
-                    constraints
-                        .protocol
-                        .clone()
-                        .map(|protocol| TransportProtocol::from_i32(protocol.protocol).unwrap())
-                )
-            )
+            let ovpn_constraints =
+                mullvad_types::relay_constraints::OpenVpnConstraints::try_from(constraints)
+                    .unwrap();
+            format!("{}", ovpn_constraints)
         } else {
             "any port over any transport protocol".to_string()
         }
@@ -805,19 +769,11 @@ fn parse_port_constraint(raw_port: &str) -> Result<Constraint<u16>> {
     }
 }
 
-/// Parses a protocol constraint string. Can be infallible because the possible values are limited
-/// with clap.
-fn parse_protocol_constraint(raw_protocol: &str) -> Constraint<TransportProtocol> {
+fn parse_protocol(raw_protocol: &str) -> Constraint<TransportProtocol> {
     match raw_protocol {
         "any" => Constraint::Any,
-        protocol => Constraint::Only(parse_protocol(protocol)),
-    }
-}
-
-fn parse_protocol(raw_protocol: &str) -> TransportProtocol {
-    match raw_protocol {
-        "udp" => TransportProtocol::Udp,
-        "tcp" => TransportProtocol::Tcp,
+        "udp" => Constraint::Only(TransportProtocol::Udp),
+        "tcp" => Constraint::Only(TransportProtocol::Tcp),
         _ => unreachable!(),
     }
 }
@@ -846,3 +802,22 @@ fn parse_entry_location_constraint<'a, T: Iterator<Item = &'a str>>(
         location.next(),
     ))
 }
+
+fn parse_transport_port(matches: &clap::ArgMatches<'_>) -> Result<Option<TransportPort>> {
+    let port = parse_port_constraint(matches.value_of("port").unwrap())?;
+    let protocol = parse_protocol(matches.value_of("transport protocol").unwrap());
+    match (port, protocol) {
+        (Constraint::Any, Constraint::Any) => Ok(None),
+        (Constraint::Any, Constraint::Only(protocol)) => Ok(Some(TransportPort {
+            protocol: protocol as i32,
+            ..TransportPort::default()
+        })),
+        (Constraint::Only(port), Constraint::Only(protocol)) => Ok(Some(TransportPort {
+            protocol: protocol as i32,
+            port: u32::from(port),
+        })),
+        (Constraint::Only(_), Constraint::Any) => Err(Error::InvalidCommand(
+            "a transport protocol must be given to select a specific port",
+        )),
+    }
+}
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 8417926c3d..ad1a038817 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -389,12 +389,13 @@ impl RelaySelector {
             // If no tunnel protocol is selected, use preferred constraints
             Constraint::Any => {
                 if original_constraints.openvpn_constraints.port.is_any()
-                    && original_constraints.openvpn_constraints.protocol.is_any()
                     || bridge_state == BridgeState::On
                 {
                     relay_constraints.openvpn_constraints = OpenVpnConstraints {
-                        port: preferred_port,
-                        protocol: Constraint::Only(preferred_protocol),
+                        port: Constraint::Only(TransportPort {
+                            protocol: preferred_protocol,
+                            port: preferred_port,
+                        }),
                     };
                 } else {
                     relay_constraints.openvpn_constraints =
@@ -414,15 +415,19 @@ impl RelaySelector {
             Constraint::Only(TunnelType::OpenVpn) => {
                 let openvpn_constraints = &mut relay_constraints.openvpn_constraints;
                 *openvpn_constraints = original_constraints.openvpn_constraints;
-                if bridge_state == BridgeState::On && openvpn_constraints.protocol.is_any() {
+                if bridge_state == BridgeState::On && openvpn_constraints.port.is_any() {
                     // FIXME: This is temporary while talpid-core only supports TCP proxies
-                    openvpn_constraints.protocol = Constraint::Only(TransportProtocol::Tcp);
-                } else if openvpn_constraints.port.is_any() && openvpn_constraints.protocol.is_any()
-                {
+                    openvpn_constraints.port = Constraint::Only(TransportPort {
+                        protocol: TransportProtocol::Tcp,
+                        port: Constraint::Any,
+                    });
+                } else if openvpn_constraints.port.is_any() {
                     let (preferred_port, preferred_protocol) =
                         Self::preferred_openvpn_constraints(retry_attempt);
-                    openvpn_constraints.port = preferred_port;
-                    openvpn_constraints.protocol = Constraint::Only(preferred_protocol);
+                    openvpn_constraints.port = Constraint::Only(TransportPort {
+                        protocol: preferred_protocol,
+                        port: preferred_port,
+                    });
                 }
             }
             Constraint::Only(TunnelType::Wireguard) => {
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 489b8982d3..aa59458979 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -313,18 +313,13 @@ message TunnelTypeUpdate {
 	TunnelTypeConstraint tunnel_type = 2;
 }
 
-message TransportProtocolConstraint {
-	TransportProtocol protocol = 1;
-}
-
 message TransportPort {
 	TransportProtocol protocol = 1;
 	uint32 port = 2;
 }
 
 message OpenvpnConstraints {
-	uint32 port = 1;
-	TransportProtocolConstraint protocol = 2;
+	TransportPort port = 1;
 }
 
 enum IpVersion {
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index 390560fe49..38e5031df1 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -294,14 +294,6 @@ impl From<talpid_types::net::TransportProtocol> for TransportProtocol {
     }
 }
 
-impl From<TransportProtocol> for TransportProtocolConstraint {
-    fn from(protocol: TransportProtocol) -> Self {
-        Self {
-            protocol: i32::from(protocol),
-        }
-    }
-}
-
 impl From<talpid_types::net::IpVersion> for IpVersion {
     fn from(version: talpid_types::net::IpVersion) -> Self {
         match version {
@@ -518,14 +510,11 @@ impl From<mullvad_types::relay_constraints::RelaySettings> for RelaySettings {
                     }),
 
                     openvpn_constraints: Some(OpenvpnConstraints {
-                        port: u32::from(constraints.openvpn_constraints.port.unwrap_or(0)),
-                        protocol: constraints
+                        port: constraints
                             .openvpn_constraints
-                            .protocol
-                            .as_ref()
+                            .port
                             .option()
-                            .map(|protocol| TransportProtocol::from(*protocol))
-                            .map(TransportProtocolConstraint::from),
+                            .map(TransportPort::from),
                     }),
                 })
             }
@@ -728,6 +717,23 @@ impl TryFrom<&WireguardConstraints> for mullvad_types::relay_constraints::Wiregu
     }
 }
 
+impl TryFrom<&OpenvpnConstraints> for mullvad_types::relay_constraints::OpenVpnConstraints {
+    type Error = FromProtobufTypeError;
+
+    fn try_from(
+        constraints: &OpenvpnConstraints,
+    ) -> Result<mullvad_types::relay_constraints::OpenVpnConstraints, Self::Error> {
+        use mullvad_types::relay_constraints as mullvad_constraints;
+
+        Ok(mullvad_constraints::OpenVpnConstraints {
+            port: Constraint::from(match &constraints.port {
+                Some(port) => Some(mullvad_constraints::TransportPort::try_from(port.clone())?),
+                None => None,
+            }),
+        })
+    }
+}
+
 impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySettingsUpdate {
     type Error = FromProtobufTypeError;
 
@@ -912,11 +918,11 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                     None
                 };
 
-                let openvpn_transport_protocol =
+                let openvpn_transport_port =
                     if let Some(ref constraints) = settings.openvpn_constraints {
-                        match &constraints.protocol {
-                            Some(constraint) => {
-                                Some(try_transport_protocol_from_i32(constraint.protocol)?)
+                        match &constraints.port {
+                            Some(port) => {
+                                Some(mullvad_constraints::TransportPort::try_from(port.clone())?)
                             }
                             None => None,
                         }
@@ -986,14 +992,9 @@ impl TryFrom<RelaySettingsUpdate> for mullvad_types::relay_constraints::RelaySet
                                 ),
                             }
                         }),
-                        openvpn_constraints: settings.openvpn_constraints.map(|constraints| {
+                        openvpn_constraints: settings.openvpn_constraints.map(|_constraints| {
                             mullvad_constraints::OpenVpnConstraints {
-                                port: if constraints.port != 0 {
-                                    Constraint::Only(constraints.port as u16)
-                                } else {
-                                    Constraint::Any
-                                },
-                                protocol: Constraint::from(openvpn_transport_protocol),
+                                port: Constraint::from(openvpn_transport_port),
                             }
                         }),
                     },
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index fdefb25fce..d1d3a87f71 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -71,6 +71,10 @@ impl<T: fmt::Debug + Clone + Eq + PartialEq> Constraint<T> {
         }
     }
 
+    pub fn is_only(&self) -> bool {
+        !self.is_any()
+    }
+
     pub fn as_ref(&self) -> Constraint<&T> {
         match self {
             Constraint::Any => Constraint::Any,
@@ -176,13 +180,12 @@ impl RelaySettings {
                 if constraints.tunnel_protocol == Constraint::Only(TunnelType::Wireguard) {
                     constraints.tunnel_protocol = Constraint::Any;
                 }
-                if constraints.openvpn_constraints.protocol
-                    == Constraint::Only(TransportProtocol::Udp)
+                if let Constraint::Only(TransportPort {
+                    protocol: TransportProtocol::Udp,
+                    ..
+                }) = constraints.openvpn_constraints.port
                 {
-                    constraints.openvpn_constraints = OpenVpnConstraints {
-                        protocol: Constraint::Any,
-                        port: Constraint::Any,
-                    }
+                    constraints.openvpn_constraints.port = Constraint::Any;
                 }
             }
             RelaySettings::CustomTunnelEndpoint(config) => {
@@ -454,27 +457,36 @@ pub struct TransportPort {
 /// [`Constraint`]s applicable to OpenVPN relay servers.
 #[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Deserialize, Serialize)]
 pub struct OpenVpnConstraints {
-    pub port: Constraint<u16>,
-    pub protocol: Constraint<TransportProtocol>,
+    pub port: Constraint<TransportPort>,
 }
 
 impl fmt::Display for OpenVpnConstraints {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         match self.port {
-            Constraint::Any => write!(f, "any port")?,
-            Constraint::Only(port) => write!(f, "port {}", port)?,
-        }
-        write!(f, " over ")?;
-        match self.protocol {
-            Constraint::Any => write!(f, "any protocol"),
-            Constraint::Only(protocol) => write!(f, "{}", protocol),
+            Constraint::Any => write!(f, "any port"),
+            Constraint::Only(port) => {
+                match port.port {
+                    Constraint::Any => write!(f, "any port")?,
+                    Constraint::Only(port) => write!(f, "port {}", port)?,
+                }
+                write!(f, " over {}", port.protocol)
+            }
         }
     }
 }
 
 impl Match<OpenVpnEndpointData> for OpenVpnConstraints {
     fn matches(&self, endpoint: &OpenVpnEndpointData) -> bool {
-        self.port.matches_eq(&endpoint.port) && self.protocol.matches_eq(&endpoint.protocol)
+        match self.port {
+            Constraint::Any => true,
+            Constraint::Only(transport_port) => {
+                transport_port.protocol == endpoint.protocol
+                    && match transport_port.port {
+                        Constraint::Any => true,
+                        Constraint::Only(port) => port == endpoint.port,
+                    }
+            }
+        }
     }
 }
 
@@ -613,7 +625,11 @@ impl RelaySettingsUpdate {
                 if let Some(Constraint::Only(TunnelType::Wireguard)) = &update.tunnel_protocol {
                     false
                 } else if let Some(constraints) = &update.openvpn_constraints {
-                    if let Constraint::Only(TransportProtocol::Udp) = &constraints.protocol {
+                    if let Constraint::Only(TransportPort {
+                        protocol: TransportProtocol::Udp,
+                        ..
+                    }) = &constraints.port
+                    {
                         false
                     } else {
                         true