Use TransportPort for OpenVPN constraints

1 files changed, 38 insertions, 63 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",
+        )),
+    }
+}