Merge branch 'reject-invalid-ports'

7 files changed, 148 insertions, 84 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1c51fdb530..9e9bd14c13 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,8 @@ Line wrap the file at 100 chars.                                              Th
   the option to log other devices out when the account already has five devices.
 
 ### Changed
+- Reject invalid WireGuard ports in the CLI.
+
 #### Android
 - Lowered default MTU to 1280 on Android.
 - Disable app icon badge for tunnel state notification/status.
diff --git a/gui/src/main/daemon-rpc.ts b/gui/src/main/daemon-rpc.ts
index a868208ac2..310a9a5c83 100644
--- a/gui/src/main/daemon-rpc.ts
+++ b/gui/src/main/daemon-rpc.ts
@@ -246,17 +246,10 @@ export class DaemonRpc {
     }
   }
 
-  public getRelayLocations(): Promise<IRelayList> {
+  public async getRelayLocations(): Promise<IRelayList> {
     if (this.isConnected) {
-      return new Promise((resolve, reject) => {
-        const relayLocations: IRelayListCountry[] = [];
-        const stream = this.client.getRelayLocations(new Empty());
-        stream.on('data', (country: grpcTypes.RelayListCountry) =>
-          relayLocations.push(convertFromRelayListCountry(country.toObject())),
-        );
-        stream.on('end', () => resolve({ countries: relayLocations }));
-        stream.on('close', reject);
-      });
+      const response = await this.callEmpty<grpcTypes.RelayList>(this.client.getRelayLocations);
+      return convertFromRelayList(response);
     } else {
       throw noConnectionError;
     }
@@ -711,6 +704,16 @@ function liftConstraint<T>(constraint: Constraint<T> | undefined): T | undefined
   return undefined;
 }
 
+function convertFromRelayList(relayList: grpcTypes.RelayList): IRelayList {
+  return {
+    countries: relayList
+      .getCountriesList()
+      .map((country: grpcTypes.RelayListCountry) =>
+        convertFromRelayListCountry(country.toObject()),
+      ),
+  };
+}
+
 function convertFromRelayListCountry(
   country: grpcTypes.RelayListCountry.AsObject,
 ): IRelayListCountry {
@@ -1160,15 +1163,7 @@ function convertFromDaemonEvent(data: grpcTypes.DaemonEvent): DaemonEvent {
 
   const relayList = data.getRelayList();
   if (relayList !== undefined) {
-    return {
-      relayList: {
-        countries: relayList
-          .getCountriesList()
-          ?.map((country: grpcTypes.RelayListCountry) =>
-            convertFromRelayListCountry(country.toObject()),
-          ),
-      },
-    };
+    return { relayList: convertFromRelayList(relayList) };
   }
 
   const deviceConfig = data.getDevice();
diff --git a/mullvad-cli/src/cmds/bridge.rs b/mullvad-cli/src/cmds/bridge.rs
index 27532f1416..6e09b75084 100644
--- a/mullvad-cli/src/cmds/bridge.rs
+++ b/mullvad-cli/src/cmds/bridge.rs
@@ -424,7 +424,7 @@ impl Bridge {
 
     async fn list_bridge_relays() -> Result<()> {
         let mut rpc = new_rpc_client().await?;
-        let mut locations = rpc
+        let relay_list = rpc
             .get_relay_locations(())
             .await
             .map_err(|error| Error::RpcFailedExt("Failed to obtain relay locations", error))?
@@ -432,7 +432,7 @@ impl Bridge {
 
         let mut countries = Vec::new();
 
-        while let Some(mut country) = locations.message().await? {
+        for mut country in relay_list.countries {
             country.cities = country
                 .cities
                 .into_iter()
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 8e68adda6c..e48c4681f9 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -554,12 +554,30 @@ impl Relay {
 
     async fn set_wireguard_constraints(&self, matches: &clap::ArgMatches) -> Result<()> {
         let mut rpc = new_rpc_client().await?;
+        let relay_list = rpc
+            .get_relay_locations(())
+            .await?
+            .into_inner()
+            .wireguard
+            .unwrap();
         let mut wireguard_constraints = self.get_wireguard_constraints(&mut rpc).await?;
 
         if let Some(port) = matches.value_of("port") {
             wireguard_constraints.port = match parse_port_constraint(port)? {
                 Constraint::Any => 0,
-                Constraint::Only(specific_port) => u32::from(specific_port),
+                Constraint::Only(specific_port) => {
+                    let specific_port = u32::from(specific_port);
+
+                    let is_valid_port = relay_list
+                        .port_ranges
+                        .iter()
+                        .any(|range| range.first <= specific_port && specific_port <= range.last);
+                    if !is_valid_port {
+                        return Err(Error::CommandFailed("The specified port is invalid"));
+                    }
+
+                    specific_port
+                }
             }
         }
 
@@ -718,7 +736,7 @@ impl Relay {
 
     async fn get_filtered_relays() -> Result<Vec<types::RelayListCountry>> {
         let mut rpc = new_rpc_client().await?;
-        let mut locations = rpc
+        let relay_list = rpc
             .get_relay_locations(())
             .await
             .map_err(|error| Error::RpcFailedExt("Failed to obtain relay locations", error))?
@@ -726,7 +744,7 @@ impl Relay {
 
         let mut countries = Vec::new();
 
-        while let Some(mut country) = locations.message().await? {
+        for mut country in relay_list.countries {
             country.cities = country
                 .cities
                 .into_iter()
@@ -755,7 +773,7 @@ fn parse_port_constraint(raw_port: &str) -> Result<Constraint<u16>> {
     match raw_port.to_lowercase().as_str() {
         "any" => Ok(Constraint::Any),
         port => Ok(Constraint::Only(u16::from_str(port).map_err(|_| {
-            Error::InvalidCommand("Invalid port. Must be \"any\" or [0-65535].")
+            Error::InvalidCommand("Invalid port. Must be \"any\" or 0-65535.")
         })?)),
     }
 }
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 26ad5e39c0..68ceed4c16 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -24,13 +24,12 @@ use parking_lot::RwLock;
 #[cfg(windows)]
 use std::path::PathBuf;
 use std::{
-    cmp,
     convert::{TryFrom, TryInto},
     sync::Arc,
     time::Duration,
 };
 use talpid_types::ErrorExt;
-use tokio_stream::wrappers::{ReceiverStream, UnboundedReceiverStream};
+use tokio_stream::wrappers::UnboundedReceiverStream;
 
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
@@ -54,7 +53,6 @@ const USED_VOUCHER_MESSAGE: &str = "This voucher code has already been used";
 
 #[mullvad_management_interface::async_trait]
 impl ManagementService for ManagementServiceImpl {
-    type GetRelayLocationsStream = ReceiverStream<Result<types::RelayListCountry, Status>>;
     type GetSplitTunnelProcessesStream = UnboundedReceiverStream<Result<i32, Status>>;
     type EventsListenStream = EventsListenerReceiver;
 
@@ -189,34 +187,14 @@ impl ManagementService for ManagementServiceImpl {
             .map_err(map_settings_error)
     }
 
-    async fn get_relay_locations(
-        &self,
-        _: Request<()>,
-    ) -> ServiceResult<Self::GetRelayLocationsStream> {
+    async fn get_relay_locations(&self, _: Request<()>) -> ServiceResult<types::RelayList> {
         log::debug!("get_relay_locations");
 
         let (tx, rx) = oneshot::channel();
         self.send_command_to_daemon(DaemonCommand::GetRelayLocations(tx))?;
-        let locations = self.wait_for_result(rx).await?;
-
-        let (stream_tx, stream_rx) =
-            tokio::sync::mpsc::channel(cmp::max(1, locations.countries.len()));
-
-        tokio::spawn(async move {
-            for country in locations.countries.into_iter() {
-                if let Err(error) = stream_tx
-                    .send(Ok(types::RelayListCountry::from(country)))
-                    .await
-                {
-                    log::error!(
-                        "Error while sending relays to client: {}",
-                        error.display_chain()
-                    );
-                }
-            }
-        });
-
-        Ok(Response::new(ReceiverStream::new(stream_rx)))
+        self.wait_for_result(rx)
+            .await
+            .map(|relays| Response::new(types::RelayList::from(relays)))
     }
 
     async fn get_current_location(&self, _: Request<()>) -> ServiceResult<types::GeoIpLocation> {
@@ -919,17 +897,10 @@ impl EventListener for ManagementInterfaceEventBroadcaster {
     /// Sends relays to all subscribers of the management interface.
     fn notify_relay_list(&self, relay_list: RelayList) {
         log::debug!("Broadcasting new relay list");
-        let mut new_list = types::RelayList {
-            countries: Vec::new(),
-        };
-        new_list.countries.reserve(relay_list.countries.len());
-        for country in relay_list.countries.into_iter() {
-            new_list
-                .countries
-                .push(types::RelayListCountry::from(country));
-        }
         self.notify(types::DaemonEvent {
-            event: Some(daemon_event::Event::RelayList(new_list)),
+            event: Some(daemon_event::Event::RelayList(types::RelayList::from(
+                relay_list,
+            ))),
         })
     }
 
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 4c6ee5ab71..8ec2bdf529 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -29,7 +29,7 @@ service ManagementService {
 	// Relays and tunnel constraints
 	rpc UpdateRelayLocations(google.protobuf.Empty) returns (google.protobuf.Empty) {}
 	rpc UpdateRelaySettings(RelaySettingsUpdate) returns (google.protobuf.Empty) {}
-	rpc GetRelayLocations(google.protobuf.Empty) returns (stream RelayListCountry) {}
+	rpc GetRelayLocations(google.protobuf.Empty) returns (RelayList) {}
 	rpc GetCurrentLocation(google.protobuf.Empty) returns (GeoIpLocation) {}
 	rpc SetBridgeSettings(BridgeSettings) returns (google.protobuf.Empty) {}
 	rpc SetBridgeState(BridgeState) returns (google.protobuf.Empty) {}
@@ -547,23 +547,49 @@ enum TransportProtocol {
 	TCP = 1;
 }
 
-message ShadowsocksEndpointData {
-    uint32 port = 1;
-    string cipher = 2;
-    string password = 3;
-    TransportProtocol protocol = 4;
+message DaemonEvent {
+	oneof event {
+		TunnelState tunnel_state = 1;
+		Settings settings = 2;
+		RelayList relay_list = 3;
+		AppVersionInfo version_info = 4;
+		DeviceEvent device = 5;
+		RemoveDeviceEvent remove_device = 6;
+	}
+}
+
+message RelayList {
+	repeated RelayListCountry countries = 1;
+	OpenVpnEndpointData openvpn = 2;
+	BridgeEndpointData bridge = 3;
+	WireguardEndpointData wireguard = 4;
 }
 
 message OpenVpnEndpointData {
+    repeated OpenVpnEndpoint endpoints = 1;
+}
+
+message OpenVpnEndpoint {
     uint32 port = 1;
     TransportProtocol protocol = 2;
 }
 
+message BridgeEndpointData {
+    repeated ShadowsocksEndpointData shadowsocks = 1;
+}
+
+message ShadowsocksEndpointData {
+    uint32 port = 1;
+    string cipher = 2;
+    string password = 3;
+    TransportProtocol protocol = 4;
+}
+
 message WireguardEndpointData {
 	repeated PortRange port_ranges = 1;
 	string ipv4_gateway = 2;
 	string ipv6_gateway = 3;
-	bytes public_key = 4;
+	repeated uint32 udp2tcp_ports = 4;
 }
 
 message PortRange {
@@ -571,21 +597,6 @@ message PortRange {
 	uint32 last = 2;
 }
 
-message DaemonEvent {
-	oneof event {
-		TunnelState tunnel_state = 1;
-		Settings settings = 2;
-		RelayList relay_list = 3;
-		AppVersionInfo version_info = 4;
-		DeviceEvent device = 5;
-		RemoveDeviceEvent remove_device = 6;
-	}
-}
-
-message RelayList {
-	repeated RelayListCountry countries = 1;
-}
-
 message AccountAndDevice {
 	string account_token = 1;
 	Device device = 2;
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index d3404e2298..4ea8a9729e 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -694,6 +694,73 @@ impl From<&mullvad_types::settings::TunnelOptions> for TunnelOptions {
     }
 }
 
+impl From<mullvad_types::relay_list::RelayList> for RelayList {
+    fn from(relay_list: mullvad_types::relay_list::RelayList) -> Self {
+        let mut proto_list = RelayList {
+            countries: vec![],
+            openvpn: Some(OpenVpnEndpointData::from(relay_list.openvpn)),
+            bridge: Some(BridgeEndpointData::from(relay_list.bridge)),
+            wireguard: Some(WireguardEndpointData::from(relay_list.wireguard)),
+        };
+        proto_list.countries = relay_list
+            .countries
+            .into_iter()
+            .map(RelayListCountry::from)
+            .collect();
+        proto_list
+    }
+}
+
+impl From<mullvad_types::relay_list::OpenVpnEndpointData> for OpenVpnEndpointData {
+    fn from(openvpn: mullvad_types::relay_list::OpenVpnEndpointData) -> Self {
+        OpenVpnEndpointData {
+            endpoints: openvpn
+                .ports
+                .into_iter()
+                .map(|endpoint| OpenVpnEndpoint {
+                    port: u32::from(endpoint.port),
+                    protocol: TransportProtocol::from(endpoint.protocol) as i32,
+                })
+                .collect(),
+        }
+    }
+}
+
+impl From<mullvad_types::relay_list::BridgeEndpointData> for BridgeEndpointData {
+    fn from(bridge: mullvad_types::relay_list::BridgeEndpointData) -> Self {
+        BridgeEndpointData {
+            shadowsocks: bridge
+                .shadowsocks
+                .into_iter()
+                .map(|endpoint| ShadowsocksEndpointData {
+                    port: u32::from(endpoint.port),
+                    cipher: endpoint.cipher,
+                    password: endpoint.password,
+                    protocol: TransportProtocol::from(endpoint.protocol) as i32,
+                })
+                .collect(),
+        }
+    }
+}
+
+impl From<mullvad_types::relay_list::WireguardEndpointData> for WireguardEndpointData {
+    fn from(wireguard: mullvad_types::relay_list::WireguardEndpointData) -> Self {
+        WireguardEndpointData {
+            port_ranges: wireguard
+                .port_ranges
+                .into_iter()
+                .map(|(first, last)| PortRange {
+                    first: u32::from(first),
+                    last: u32::from(last),
+                })
+                .collect(),
+            ipv4_gateway: wireguard.ipv4_gateway.to_string(),
+            ipv6_gateway: wireguard.ipv6_gateway.to_string(),
+            udp2tcp_ports: wireguard.udp2tcp_ports.into_iter().map(u32::from).collect(),
+        }
+    }
+}
+
 impl From<mullvad_types::relay_list::RelayListCountry> for RelayListCountry {
     fn from(country: mullvad_types::relay_list::RelayListCountry) -> Self {
         let mut proto_country = RelayListCountry {