Replace axum with salvo

gamestream-webtransport-proxy/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2024"
 
 [dependencies]
 anyhow = "1.0.98"
-axum = "0.8.4"
+directories = "6.0.0"
 getrandom = { version = "0.3.3", features = ["std"] }
 hex = "0.4.3"
 libc = "0.2.174"
@@ -15,6 +15,7 @@ rand = "0.9.1"
 reqwest = { version = "0.12.20", features = [
   "rustls-tls",
 ], default-features = false }
+salvo = { version = "0.79.0", features = ["oapi"] }
 serde = { version = "1.0.219", features = ["serde_derive"] }
 serde-xml-rs = "0.8.1"
 tokio = { version = "1.45.1", features = ["full"] }

gamestream-webtransport-proxy/src/main.rs

@@ -1,4 +1,4 @@
-use axum::{Router, routing::get};
+use salvo::prelude::*;
 use std::ffi::CString;
 
 use moonlight_common_c_sys::{
@@ -101,12 +101,12 @@ fn barmain() {
 
 #[tokio::main]
 async fn main() {
-    // Create a router with a test endpoint
-    let app = Router::new().route("/pair/{url}/{port}", get(pair::get_pair));
-
-    // Bind to port 3000
-    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
-
-    // Start the server
-    axum::serve(listener, app).await.unwrap();
+    let mut router = Router::new().push(Router::with_path("pair").post(pair::post_pair));
+    let doc = OpenApi::new("test api", "0.0.1").merge_router(&router);
+    let router = router
+        .unshift(doc.into_router("/api-doc/openapi.json"))
+        .unshift(SwaggerUi::new("/api-doc/openapi.json").into_router("/swagger-ui"));
+    let listener = TcpListener::new("0.0.0.0:3001");
+    let acceptor = listener.join(TcpListener::new("0.0.0.0:3000")).bind().await;
+    salvo::Server::new(acceptor).serve(router).await;
 }

gamestream-webtransport-proxy/src/pair.rs

@@ -1,17 +1,26 @@
-use axum::extract::Path;
-use axum::http::StatusCode;
-use axum::response::{IntoResponse, Response};
+use anyhow::Result;
 use openssl::hash::MessageDigest;
-use openssl::sha::Sha256;
-use rand::Rng;
-use serde::{Deserialize, Serialize};
-
 use openssl::pkey::{PKey, Private};
 use openssl::rsa::Rsa;
-use openssl::x509::{self, X509};
+use openssl::sha::Sha256;
+use openssl::x509::X509;
+use rand::Rng;
+use salvo::prelude::*;
+use serde::{Deserialize, Serialize};
+use std::fs;
+use std::io::Write;
 
-use anyhow::Result;
-use url_constructor::UrlConstructor;
+#[derive(Debug, Deserialize, ToSchema)]
+struct PostPairParams {
+    host: String,
+    port: u16,
+    pair_endpoint: Option<String>,
+}
+
+#[derive(Debug, Serialize, ToSchema)]
+struct PostPairReturn {
+    paired: bool,
+}
 
 #[derive(Debug, Deserialize)]
 struct ServerCertResponse {
@@ -20,23 +29,28 @@ struct ServerCertResponse {
 }
 
 #[derive(Debug, Deserialize)]
-pub struct ClientChallengeResponse {
+struct ClientChallengeResponse {
     challengeresponse: String,
 }
 
 #[derive(Debug, Deserialize)]
-pub struct ServerChallengeResponseResponse {
+struct ServerChallengeResponseResponse {
     paired: i32,
     pairingsecret: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct ClientPairingSecretResponse {
+    paired: i32,
+}
+
 #[derive(Debug)]
 struct ServerCert {
     cert: Vec<u8>,
 }
 
 #[derive(Debug)]
-pub struct ServerPairingSecret {
+struct ServerPairingSecret {
     pairing_secret: Vec<u8>,
     signature: Vec<u8>,
 }
@@ -74,6 +88,49 @@ fn get_cert_and_private_key() -> Result<(X509, PKey<Private>)> {
     Ok((cert, key_pair))
 }
 
+fn save_cert_and_key_to_disk(
+    cert: X509,
+    key: PKey<Private>,
+    host: &String,
+    port: u16,
+) -> Result<()> {
+    let project_dirs =
+        directories::ProjectDirs::from("xyz", "ohea", "gamestream-webtransport-proxy")
+            .ok_or(anyhow::anyhow!("Could not get project dirs"))?;
+    let data_dir = project_dirs.data_dir();
+    let cert_dir = data_dir.join("certs");
+    fs::create_dir_all(&cert_dir)?;
+
+    let cert_filepath = cert_dir.join(format!("{host}_{port}_cert"));
+    let key_filepath = cert_dir.join(format!("{host}_{port}_key"));
+
+    let mut cert_file_builder = std::fs::OpenOptions::new();
+    cert_file_builder.create(true);
+    cert_file_builder.truncate(true);
+    cert_file_builder.write(true);
+
+    let mut key_file_builder = std::fs::OpenOptions::new();
+    key_file_builder.create(true);
+    key_file_builder.truncate(true);
+    key_file_builder.write(true);
+
+    #[cfg(target_family = "unix")]
+    {
+        use std::os::unix::fs::OpenOptionsExt;
+
+        key_file_builder.mode(0o600);
+        cert_file_builder.mode(0o600);
+    }
+
+    let mut cert_file = cert_file_builder.open(&cert_filepath)?;
+    let mut key_file = key_file_builder.open(&key_filepath)?;
+
+    cert_file.write_all(&cert.to_pem()?)?;
+    key_file.write_all(&key.private_key_to_pem_pkcs8()?)?;
+
+    Ok(())
+}
+
 async fn get_url(base_url: &mut url_constructor::UrlConstructor) -> Result<String> {
     let mut uuidv2 = [0u8; 16];
     openssl::rand::rand_bytes(&mut uuidv2)?;
@@ -218,7 +275,7 @@ async fn send_server_challenge_response(
     Ok(serde_xml_rs::from_str(&text)?)
 }
 
-async fn generate_aes_key(salt: [u8; 16], pin: [u8; 4]) -> Vec<u8> {
+fn generate_aes_key(salt: [u8; 16], pin: [u8; 4]) -> Vec<u8> {
     let mut salt_pin = Vec::with_capacity(salt.len() + pin.len());
     salt_pin.extend_from_slice(&salt);
     salt_pin.extend_from_slice(&pin);
@@ -235,7 +292,7 @@ async fn do_challenge(
     salt: [u8; 16],
     cert: &X509,
 ) -> Result<ServerPairingSecret> {
-    let aes_key = generate_aes_key(salt, pin).await;
+    let aes_key = generate_aes_key(salt, pin);
     let aes_hex = hex::encode(&aes_key);
     //println!("aes_hex: {aes_hex}");
 
@@ -259,15 +316,11 @@ async fn do_challenge(
     })
 }
 
-pub async fn get_base_url(
-    host: String,
-    port: u16,
-    unique_id: String,
-) -> url_constructor::UrlConstructor {
+fn get_base_url(host: &String, port: u16, unique_id: String) -> url_constructor::UrlConstructor {
     let mut base_url = url_constructor::UrlConstructor::new();
     base_url
         .scheme("http")
-        .host(&host)
+        .host(host)
         .port(port)
         .subdir("pair")
         .param("uniqueid", unique_id)
@@ -276,7 +329,7 @@ pub async fn get_base_url(
     base_url
 }
 
-pub async fn generate_pin() -> [u8; 4] {
+fn generate_pin() -> [u8; 4] {
     let mut pin = [0u8; 4];
     {
         print!("pairing pin: ");
@@ -284,16 +337,16 @@ pub async fn generate_pin() -> [u8; 4] {
         // TODO: reenable real RNG
         let mut rng = rand::rng();
         for i in 0..pin.len() {
-            pin[i] = rng.random_range(48..58); // Generate ascii number 0-9
+            // Generate ascii number 0-9
+            pin[i] = rng.random_range(48..58);
             print!("{}", pin[i] as char);
         }
-        // Print as a four-digit, zero-padded integer
         println!("");
     }
     pin
 }
 
-pub async fn verify_signature(secret: Vec<u8>, signature: Vec<u8>, cert: Vec<u8>) -> Result<()> {
+fn verify_signature(secret: Vec<u8>, signature: Vec<u8>, cert: Vec<u8>) -> Result<()> {
     let cert_x509 = openssl::x509::X509::from_pem(&cert)?;
     let cert_pubkey = cert_x509.public_key()?;
 
@@ -310,7 +363,7 @@ pub async fn verify_signature(secret: Vec<u8>, signature: Vec<u8>, cert: Vec<u8>
     }
 }
 
-pub async fn create_signature(data: &[u8; 16], private_key: &PKey<Private>) -> Result<Vec<u8>> {
+fn create_signature(data: &[u8; 16], private_key: &PKey<Private>) -> Result<Vec<u8>> {
     let mut signature = Vec::new();
 
     let md = openssl::md::Md::sha256();
@@ -324,12 +377,12 @@ pub async fn create_signature(data: &[u8; 16], private_key: &PKey<Private>) -> R
     Ok(signature)
 }
 
-pub async fn send_client_pairing_secret(
+async fn send_client_pairing_secret(
     mut base_url: url_constructor::UrlConstructor,
     client_secret_data: &[u8; 16],
     private_key: &PKey<Private>,
-) -> Result<()> {
-    let signature = create_signature(client_secret_data, private_key).await?;
+) -> Result<ClientPairingSecretResponse> {
+    let signature = create_signature(client_secret_data, private_key)?;
 
     let mut client_secret = Vec::with_capacity(client_secret_data.len() + signature.len());
     client_secret.extend_from_slice(client_secret_data);
@@ -340,32 +393,36 @@ pub async fn send_client_pairing_secret(
     let url = base_url.param("clientpairingsecret", client_secret_hex);
 
     let text = get_url(url).await?;
-    println!("{text:?}");
 
-    Ok(())
+    Ok(serde_xml_rs::from_str(&text)?)
 }
 
-async fn get_unique_id() -> Result<String> {
+fn get_unique_id() -> Result<String> {
     let mut bytes = [0u8; 8];
     openssl::rand::rand_bytes(&mut bytes)?;
     Ok(hex::encode(bytes))
 }
 
-// Do a pairing operation
-pub async fn get_pair(Path((host, port)): Path<(String, u16)>) -> Response {
-    //let unique_id = "0123456789ABCDEF".to_string();
-    //let uuidv2 = "23060a17e4fc40f5a83fee298995dd18".to_string();
+#[salvo::oapi::endpoint]
+pub async fn post_pair(body: salvo::oapi::extract::JsonBody<PostPairParams>) -> StatusCode {
+    let params = body.into_inner();
 
-    let unique_id = get_unique_id().await.unwrap();
+    let unique_id = match get_unique_id() {
+        Ok(u) => u,
+        Err(e) => {
+            println!("Could not generate unique id: {e}");
+            return StatusCode::INTERNAL_SERVER_ERROR;
+        }
+    };
 
-    let base_url = get_base_url(host, port, unique_id).await;
+    let base_url = get_base_url(&params.host, params.port, unique_id);
 
-    let pin = generate_pin().await;
+    let pin = generate_pin();
 
     let mut client_secret_data = [0u8; 16];
     if let Err(e) = openssl::rand::rand_bytes(&mut client_secret_data) {
         println!("Could not generate client secret data: {e}");
-        return StatusCode::INTERNAL_SERVER_ERROR.into_response();
+        return StatusCode::INTERNAL_SERVER_ERROR;
     }
 
     // Generate certificates
@@ -373,7 +430,7 @@ pub async fn get_pair(Path((host, port)): Path<(String, u16)>) -> Response {
         Ok(v) => v,
         Err(e) => {
             println!("Could not generate certs: {e}");
-            return StatusCode::INTERNAL_SERVER_ERROR.into_response();
+            return StatusCode::INTERNAL_SERVER_ERROR;
         }
     };
 
@@ -392,8 +449,8 @@ pub async fn get_pair(Path((host, port)): Path<(String, u16)>) -> Response {
     let server_cert = match get_server_cert(base_url.clone(), salt_hex, cert_hex).await {
         Ok(s) => s,
         Err(e) => {
-            println!("could not get server cert: {e}");
-            return StatusCode::INTERNAL_SERVER_ERROR.into_response();
+            println!("Could not get server cert: {e}");
+            return StatusCode::INTERNAL_SERVER_ERROR;
         }
     };
     //println!("{server_cert:?}");
@@ -404,8 +461,8 @@ pub async fn get_pair(Path((host, port)): Path<(String, u16)>) -> Response {
         match do_challenge(base_url.clone(), &client_secret_data, pin, salt, &cert).await {
             Ok(s) => s,
             Err(e) => {
-                println!("could not do challenge: {e}");
-                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
+                println!("Could not do challenge: {e}");
+                return StatusCode::INTERNAL_SERVER_ERROR;
             }
         };
     //println!("{server_pairing_secret:?}");
@@ -415,19 +472,35 @@ pub async fn get_pair(Path((host, port)): Path<(String, u16)>) -> Response {
         server_pairing_secret.pairing_secret,
         server_pairing_secret.signature,
         server_cert.cert,
-    )
-    .await
-    {
+    ) {
         println!("Could not verify signature: {e}");
-        return StatusCode::INTERNAL_SERVER_ERROR.into_response();
+        return StatusCode::INTERNAL_SERVER_ERROR;
     }
 
-    if let Err(e) =
-        send_client_pairing_secret(base_url.clone(), &client_secret_data, &private_key).await
-    {
-        println!("Could not send client pairing secret: {e}");
-        return StatusCode::INTERNAL_SERVER_ERROR.into_response();
+    let client_pairing_secret_response =
+        match send_client_pairing_secret(base_url.clone(), &client_secret_data, &private_key).await
+        {
+            Ok(p) => p,
+            Err(e) => {
+                println!("Could not send client pairing secret: {e}");
+                return StatusCode::INTERNAL_SERVER_ERROR;
+            }
+        };
+
+    if client_pairing_secret_response.paired != 1 {
+        println!("Failed to pair with server");
+        return StatusCode::INTERNAL_SERVER_ERROR;
+    } else {
+        let host = &params.host;
+        let port = params.port;
+        println!("Paired with server {host}:{port} successfully!");
+    }
+
+    // Save certificate to disk so it can be used for subsequent connections
+    if let Err(e) = save_cert_and_key_to_disk(cert, private_key, &params.host, params.port) {
+        println!("Could not save cert and key to disk");
+        return StatusCode::INTERNAL_SERVER_ERROR;
     };
 
-    "test".into_response()
+    StatusCode::OK
 }