PocketRelay · jacobtread · Nov 23, 2023 · Nov 23, 2023 · Nov 23, 2023
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pocket-relay-plugin"
-version = "0.0.2"
+version = "0.0.3"
 edition = "2021"
 description = "ASI plugin for ME3 to allow playing on Pocket Relay servers"
 repository = "https://github.com/PocketRelay/PocketRelayHooks"
@@ -32,6 +32,7 @@ futures-util = { version = "0.3", features = ["sink"] }
 thiserror = "1"
 semver = { version = "1.0", features = ["serde"] }
 hyper = { version = "0.14", features = ["server", "http1", "tcp", "runtime"] }
+url = "2.4.1"
 
 [dependencies.windows-sys]
 version = "0.48"

diff --git a/src/api/mod.rs b/src/api/mod.rs
@@ -1,3 +1,5 @@
+use std::str::FromStr;
+
 use hyper::{
     header::{ACCEPT, USER_AGENT},
     StatusCode,
@@ -7,6 +9,7 @@ use reqwest::Client;
 use semver::Version;
 use serde::Deserialize;
 use thiserror::Error;
+use url::Url;
 
 use crate::constants::{APP_VERSION, MIN_SERVER_VERSION, SERVER_IDENT};
 
@@ -26,21 +29,18 @@ struct ServerDetails {
 /// version obtained from the server
 #[derive(Debug, Clone)]
 pub struct LookupData {
-    pub scheme: String,
-    /// The host address of the server
-    pub host: String,
+    /// The server url
+    pub url: Url,
     /// The server version
     pub version: Version,
-    /// The server port
-    pub port: u16,
 }
 
 /// Errors that can occur while looking up a server
 #[derive(Debug, Error)]
 pub enum LookupError {
-    /// The server url was missing the host portion
-    #[error("Unable to find host portion of provided Connection URL")]
-    InvalidHostTarget,
+    /// The server url was invalid
+    #[error("Invalid Connection URL: {0}")]
+    InvalidHostTarget(#[from] url::ParseError),
     /// The server connection failed
     #[error("Failed to connect to server: {0}")]
     ConnectionFailed(reqwest::Error),
@@ -65,26 +65,32 @@ pub enum LookupError {
 ///
 /// `host` The host to try and lookup
 pub async fn try_lookup_host(host: &str) -> Result<LookupData, LookupError> {
-    let mut url = String::new();
-
-    // Fill in missing host portion
-    if !host.starts_with("http://") && !host.starts_with("https://") {
-        url.push_str("http://");
-        url.push_str(host)
-    } else {
-        url.push_str(host);
-    }
+    let url = {
+        let mut url = String::new();
+
+        // Fill in missing scheme portion
+        if !host.starts_with("http://") && !host.starts_with("https://") {
+            url.push_str("http://");
+            url.push_str(host)
+        } else {
+            url.push_str(host);
+        }
 
-    if !host.ends_with('/') {
-        url.push('/')
-    }
+        // Ensure theres a trailing slash (URL path will be interpeted incorrectly without)
+        if !host.ends_with('/') {
+            url.push('/');
+        }
 
-    url.push_str("api/server");
+        url
+    };
+
+    let url = Url::from_str(&url)?;
+    let info_url = url.join("api/server").expect("Failed to server info URL");
 
     let client = Client::new();
 
     let response = client
-        .get(url)
+        .get(info_url)
         .header(ACCEPT, "application/json")
         .header(USER_AGENT, format!("PocketRelayClient/v{}", APP_VERSION))
         .send()
@@ -112,15 +118,6 @@ pub async fn try_lookup_host(host: &str) -> Result<LookupData, LookupError> {
         }
     };
 
-    let url = response.url();
-    let scheme = url.scheme().to_string();
-
-    let port = url.port_or_known_default().unwrap_or(80);
-    let host = match url.host() {
-        Some(value) => value.to_string(),
-        None => return Err(LookupError::InvalidHostTarget),
-    };
-
     let details = response
         .json::<ServerDetails>()
         .await
@@ -139,9 +136,7 @@ pub async fn try_lookup_host(host: &str) -> Result<LookupData, LookupError> {
     }
 
     Ok(LookupData {
-        scheme,
-        host,
-        port,
+        url,
         version: details.version,
     })
 }
diff --git a/src/interface/mod.rs b/src/interface/mod.rs
@@ -135,12 +135,16 @@ pub fn init(runtime: tokio::runtime::Handle, config: Option<ClientConfig>) {
 
                         debug!(
                             "Connected to server {} {} version v{}",
-                            value.scheme, value.host, value.version
+                            value.url.scheme(),
+                            value.url.authority(),
+                            value.version
                         );
 
                         let message = format!(
                             "Connected: {} {} version v{}",
-                            value.scheme, value.host, value.version
+                            value.url.scheme(),
+                            value.url.authority(),
+                            value.version
                         );
 
                         c_label.set_text(&message);

diff --git a/src/servers/http.rs b/src/servers/http.rs
@@ -52,10 +52,18 @@ async fn proxy_http(
         .map(|value| value.as_str())
         .unwrap_or_default();
 
-    let target_url = format!(
-        "{}://{}:{}{}",
-        target.scheme, target.host, target.port, path
-    );
+    // Remove the leading / to make the path relative
+    let path = path.strip_prefix('/').unwrap_or(path);
+
+    let target_url = match target.url.join(path) {
+        Ok(value) => value,
+        Err(_) => {
+            // Failed to create a path
+            let mut error_response = Response::default();
+            *error_response.status_mut() = StatusCode::SERVICE_UNAVAILABLE;
+            return Ok(error_response);
+        }
+    };
 
     let client = Client::new();
     let proxy_response = match client.get(target_url).send().await {

diff --git a/src/servers/main.rs b/src/servers/main.rs
@@ -1,6 +1,6 @@
 use crate::{
     api::LookupData,
-    constants::{APP_VERSION, MAIN_PORT},
+    constants::{APP_VERSION, HTTP_PORT, MAIN_PORT},
     servers::spawn_task,
 };
 use hyper::header::USER_AGENT;
@@ -47,22 +47,22 @@ pub async fn start_server(target: Arc<LookupData>) {
 }
 
 /// Header for the Pocket Relay connection scheme used by the client
-const HEADER_SCHEME: &str = "X-Pocket-Relay-Scheme";
+const LEGACY_HEADER_SCHEME: &str = "X-Pocket-Relay-Scheme";
 /// Header for the Pocket Relay connection port used by the client
-const HEADER_PORT: &str = "X-Pocket-Relay-Port";
+const LEGACY_HEADER_PORT: &str = "X-Pocket-Relay-Port";
 /// Header for the Pocket Relay connection host used by the client
-const HEADER_HOST: &str = "X-Pocket-Relay-Host";
+const LEGACY_HEADER_HOST: &str = "X-Pocket-Relay-Host";
 /// Header to tell the server to use local HTTP
 const HEADER_LOCAL_HTTP: &str = "X-Pocket-Relay-Local-Http";
 /// Endpoint for upgrading the server connection
-const UPGRADE_ENDPOINT: &str = "/api/server/upgrade";
+const UPGRADE_ENDPOINT: &str = "api/server/upgrade";
 
 async fn handle_blaze(mut client: TcpStream, target: Arc<LookupData>) {
     // Create the upgrade URL
-    let url = format!(
- "{}://{}:{}{}",
- target.scheme, target.host, target.port, UPGRADE_ENDPOINT
-    );
+    let url = target
+ .url
+ .join(UPGRADE_ENDPOINT)
+     .expect("Failed to create upgrade endpoint URL");
 
     let user_agent = format!("PocketRelayClient/v{}", APP_VERSION);
 
@@ -75,24 +75,14 @@ async fn handle_blaze(mut client: TcpStream, target: Arc<LookupData>) {
         HeaderValue::from_str(&user_agent).expect("User agent header was invalid"),
     );
 
-    // TODO: Once users have started updating servers these fields can be removed
-
-    // Append the schema header
-    if let Ok(scheme_value) = HeaderValue::from_str(&target.scheme) {
-        headers.insert(HEADER_SCHEME, scheme_value);
-    }
-
-    // Append the port header
-    headers.insert(HEADER_PORT, HeaderValue::from(target.port));
-
-    // Append the host header
-    if let Ok(host_value) = HeaderValue::from_str(&target.host) {
-        headers.insert(HEADER_HOST, host_value);
-    }
-
     // Append use local http header
     headers.insert(HEADER_LOCAL_HTTP, HeaderValue::from_static("true"));
 
+    // Append legacy http details headers
+    headers.insert(LEGACY_HEADER_SCHEME, HeaderValue::from_static("http"));
+    headers.insert(LEGACY_HEADER_PORT, HeaderValue::from(HTTP_PORT));
+    headers.insert(LEGACY_HEADER_HOST, HeaderValue::from_static("127.0.0.1"));
+
     debug!("Connecting pipe to Pocket Relay server");
 
     // Create the request

diff --git a/src/servers/telemetry.rs b/src/servers/telemetry.rs
@@ -12,7 +12,7 @@ use tokio::{
 use super::spawn_task;
 
 /// Server API endpoint to send telemetry data to
-const TELEMETRY_ENDPOINT: &str = "/api/server/telemetry";
+const TELEMETRY_ENDPOINT: &str = "api/server/telemetry";
 
 pub async fn start_server(target: Arc<LookupData>) {
     // Initializing the underlying TCP listener
@@ -39,10 +39,11 @@ pub async fn start_server(target: Arc<LookupData>) {
 
         spawn_task(async move {
             // Create the telemetry URL
-            let url = format!(
-                "{}://{}:{}{}",
-                target.scheme, target.host, target.port, TELEMETRY_ENDPOINT
-            );
+
+            let url = target
+                .url
+                .join(TELEMETRY_ENDPOINT)
+                .expect("Failed to create telemetry endpoint");
 
             let client = Client::new();
 
@@ -52,7 +53,7 @@ pub async fn start_server(target: Arc<LookupData>) {
 
                 let message: TelemetryMessage = decode_message(message);
                 // TODO: Batch these telemetry messages and send them to the server
-                let _ = client.post(&url).json(&message).send().await;
+                let _ = client.post(url.clone()).json(&message).send().await;
             }
         })
         .await;