From 9242b9d1962db1ba3448d4a967bcde2bac55cfe3 Mon Sep 17 00:00:00 2001
From: jamesqh <18192979+jamesqh@users.noreply.github.com>
Date: Wed, 19 Jun 2024 21:15:28 +0100
Subject: [PATCH] Update TLS examples to use better HTTP->HTTPS redirect

---
 examples/tls-graceful-shutdown/src/main.rs | 20 +++++++++++++++-----
 examples/tls-rustls/src/main.rs            | 20 +++++++++++++++-----
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/examples/tls-graceful-shutdown/src/main.rs b/examples/tls-graceful-shutdown/src/main.rs
index cc5b6ecb..cd82655b 100644
--- a/examples/tls-graceful-shutdown/src/main.rs
+++ b/examples/tls-graceful-shutdown/src/main.rs
@@ -7,7 +7,7 @@
 use axum::{
     extract::Host,
     handler::HandlerWithoutStateExt,
-    http::{StatusCode, Uri},
+    http::{uri::Authority, StatusCode, Uri},
     response::Redirect,
     routing::get,
     BoxError, Router,
@@ -106,7 +106,7 @@ async fn redirect_http_to_https<F>(ports: Ports, signal: F)
 where
     F: Future<Output = ()> + Send + 'static,
 {
-    fn make_https(host: String, uri: Uri, ports: Ports) -> Result<Uri, BoxError> {
+    fn make_https(host: &str, uri: Uri, https_port: u16) -> Result<Uri, BoxError> {
         let mut parts = uri.into_parts();
 
         parts.scheme = Some(axum::http::uri::Scheme::HTTPS);
@@ -115,14 +115,24 @@ where
             parts.path_and_query = Some("/".parse().unwrap());
         }
 
-        let https_host = host.replace(&ports.http.to_string(), &ports.https.to_string());
-        parts.authority = Some(https_host.parse()?);
+        let authority: Authority = host.parse()?;
+        let bare_host = match authority.port() {
+            Some(port_struct) => authority
+                .as_str()
+                .strip_suffix(port_struct.as_str())
+                .unwrap()
+                .strip_suffix(':')
+                .unwrap(), // if authority.port() is Some(port) then we can be sure authority ends with :{port}
+            None => authority.as_str(),
+        };
+
+        parts.authority = Some(format!("{bare_host}:{https_port}").parse()?);
 
         Ok(Uri::from_parts(parts)?)
     }
 
     let redirect = move |Host(host): Host, uri: Uri| async move {
-        match make_https(host, uri, ports) {
+        match make_https(&host, uri, ports.https) {
             Ok(uri) => Ok(Redirect::permanent(&uri.to_string())),
             Err(error) => {
                 tracing::warn!(%error, "failed to convert URI to HTTPS");
diff --git a/examples/tls-rustls/src/main.rs b/examples/tls-rustls/src/main.rs
index 3649c75c..48b3c987 100644
--- a/examples/tls-rustls/src/main.rs
+++ b/examples/tls-rustls/src/main.rs
@@ -9,7 +9,7 @@
 use axum::{
     extract::Host,
     handler::HandlerWithoutStateExt,
-    http::{StatusCode, Uri},
+    http::{uri::Authority, StatusCode, Uri},
     response::Redirect,
     routing::get,
     BoxError, Router,
@@ -72,7 +72,7 @@ async fn handler() -> &'static str {
 
 #[allow(dead_code)]
 async fn redirect_http_to_https(ports: Ports) {
-    fn make_https(host: String, uri: Uri, ports: Ports) -> Result<Uri, BoxError> {
+    fn make_https(host: &str, uri: Uri, https_port: u16) -> Result<Uri, BoxError> {
         let mut parts = uri.into_parts();
 
         parts.scheme = Some(axum::http::uri::Scheme::HTTPS);
@@ -81,14 +81,24 @@ async fn redirect_http_to_https(ports: Ports) {
             parts.path_and_query = Some("/".parse().unwrap());
         }
 
-        let https_host = host.replace(&ports.http.to_string(), &ports.https.to_string());
-        parts.authority = Some(https_host.parse()?);
+        let authority: Authority = host.parse()?;
+        let bare_host = match authority.port() {
+            Some(port_struct) => authority
+                .as_str()
+                .strip_suffix(port_struct.as_str())
+                .unwrap()
+                .strip_suffix(':')
+                .unwrap(), // if authority.port() is Some(port) then we can be sure authority ends with :{port}
+            None => authority.as_str(),
+        };
+
+        parts.authority = Some(format!("{bare_host}:{https_port}").parse()?);
 
         Ok(Uri::from_parts(parts)?)
     }
 
     let redirect = move |Host(host): Host, uri: Uri| async move {
-        match make_https(host, uri, ports) {
+        match make_https(&host, uri, ports.https) {
             Ok(uri) => Ok(Redirect::permanent(&uri.to_string())),
             Err(error) => {
                 tracing::warn!(%error, "failed to convert URI to HTTPS");