From 4fb3d5f271e7ff6439ba7d5e5fe996f7f1c9ad31 Mon Sep 17 00:00:00 2001
From: Harold <harold@mrtz.fr>
Date: Tue, 3 Jul 2018 13:09:22 +0200
Subject: [PATCH] Add ap-port option to specify preferred port

---
 core/src/apresolve.rs | 17 ++++++++++++++---
 core/src/config.rs    |  2 ++
 core/src/session.rs   |  2 +-
 src/main.rs           | 24 ++++++++++++++----------
 4 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/core/src/apresolve.rs b/core/src/apresolve.rs
index 9ff23cf3..bcb7bfe1 100644
--- a/core/src/apresolve.rs
+++ b/core/src/apresolve.rs
@@ -17,7 +17,11 @@ pub struct APResolveData {
     ap_list: Vec<String>,
 }
 
-fn apresolve(handle: &Handle, proxy: &Option<Url>) -> Box<Future<Item = String, Error = Error>> {
+fn apresolve(
+    handle: &Handle,
+    proxy: &Option<Url>,
+    ap_port: &Option<u16>,
+) -> Box<Future<Item = String, Error = Error>> {
     let url = Uri::from_str(APRESOLVE_ENDPOINT).expect("invalid AP resolve URL");
     let use_proxy = proxy.is_some();
 
@@ -53,9 +57,15 @@ fn apresolve(handle: &Handle, proxy: &Option<Url>) -> Box<Future<Item = String,
     let data =
         body.and_then(|body| serde_json::from_str::<APResolveData>(&body).chain_err(|| "invalid JSON"));
 
+    let p = ap_port.clone();
+
     let ap = data.and_then(move |data| {
         let mut aps = data.ap_list.iter().filter(|ap| {
-            if use_proxy {
+            if p.is_some() {
+                Uri::from_str(ap)
+                    .ok()
+                    .map_or(false, |uri| uri.port().map_or(false, |port| port == p.unwrap()))
+            } else if use_proxy {
                 // It is unlikely that the proxy will accept CONNECT on anything other than 443.
                 Uri::from_str(ap)
                     .ok()
@@ -75,11 +85,12 @@ fn apresolve(handle: &Handle, proxy: &Option<Url>) -> Box<Future<Item = String,
 pub(crate) fn apresolve_or_fallback<E>(
     handle: &Handle,
     proxy: &Option<Url>,
+    ap_port: &Option<u16>,
 ) -> Box<Future<Item = String, Error = E>>
 where
     E: 'static,
 {
-    let ap = apresolve(handle, proxy).or_else(|e| {
+    let ap = apresolve(handle, proxy, ap_port).or_else(|e| {
         warn!("Failed to resolve Access Point: {}", e.description());
         warn!("Using fallback \"{}\"", AP_FALLBACK);
         Ok(AP_FALLBACK.into())
diff --git a/core/src/config.rs b/core/src/config.rs
index 7baff40b..283b7c83 100644
--- a/core/src/config.rs
+++ b/core/src/config.rs
@@ -10,6 +10,7 @@ pub struct SessionConfig {
     pub user_agent: String,
     pub device_id: String,
     pub proxy: Option<Url>,
+    pub ap_port: Option<u16>,
 }
 
 impl Default for SessionConfig {
@@ -19,6 +20,7 @@ impl Default for SessionConfig {
             user_agent: version::version_string(),
             device_id: device_id,
             proxy: None,
+            ap_port: None,
         }
     }
 }
diff --git a/core/src/session.rs b/core/src/session.rs
index 335cf0e3..931b60c7 100644
--- a/core/src/session.rs
+++ b/core/src/session.rs
@@ -51,7 +51,7 @@ impl Session {
         cache: Option<Cache>,
         handle: Handle,
     ) -> Box<Future<Item = Session, Error = io::Error>> {
-        let access_point = apresolve_or_fallback::<io::Error>(&handle, &config.proxy);
+        let access_point = apresolve_or_fallback::<io::Error>(&handle, &config.proxy, &config.ap_port);
 
         let handle_ = handle.clone();
         let proxy = config.proxy.clone();
diff --git a/src/main.rs b/src/main.rs
index 3665e615..a18d7b5c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -129,6 +129,7 @@ fn setup(args: &[String]) -> Setup {
         .optopt("u", "username", "Username to sign in with", "USERNAME")
         .optopt("p", "password", "Password", "PASSWORD")
         .optopt("", "proxy", "HTTP proxy to use when connecting", "PROXY")
+        .optopt("", "ap-port", "Connect to AP with specified port. If no AP with that port are present fallback AP will be used. Available ports are usually 80, 443 and 4070", "AP_PORT")
         .optflag("", "disable-discovery", "Disable discovery mode")
         .optopt(
             "",
@@ -255,20 +256,23 @@ fn setup(args: &[String]) -> Setup {
             proxy: matches.opt_str("proxy").or(std::env::var("http_proxy").ok()).map(
                 |s| {
                     match Url::parse(&s) {
-                Ok(url) => {
-                    if url.host().is_none() || url.port().is_none() {
-                        panic!("Invalid proxy url, only urls on the format \"http://host:port\" are allowed");
-                    }
+                        Ok(url) => {
+                            if url.host().is_none() || url.port().is_none() {
+                                panic!("Invalid proxy url, only urls on the format \"http://host:port\" are allowed");
+                            }
 
-                    if url.scheme() != "http" {
-                        panic!("Only unsecure http:// proxies are supported");
+                            if url.scheme() != "http" {
+                                panic!("Only unsecure http:// proxies are supported");
+                            }
+                            url
+                        },
+                    Err(err) => panic!("Invalid proxy url: {}, only urls on the format \"http://host:port\" are allowed", err)
                     }
-                    url
-                },
-                Err(err) => panic!("Invalid proxy url: {}, only urls on the format \"http://host:port\" are allowed", err)
-            }
                 },
             ),
+            ap_port: matches
+                .opt_str("ap-port")
+                .map(|port| port.parse::<u16>().expect("Invalid port")),
         }
     };