From 7155ad229aaf43acdd301084fc25590eb5833800 Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Sun, 26 Mar 2023 17:57:31 -0700
Subject: [PATCH] Add HTTP to HTTPs redirector to TLS example.

---
 examples/tls/src/main.rs       |  8 +++--
 examples/tls/src/redirector.rs | 65 ++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 examples/tls/src/redirector.rs

diff --git a/examples/tls/src/main.rs b/examples/tls/src/main.rs
index 7cd59828..8ecb685c 100644
--- a/examples/tls/src/main.rs
+++ b/examples/tls/src/main.rs
@@ -1,6 +1,8 @@
 #[macro_use] extern crate rocket;
 
-#[cfg(test)] mod tests;
+#[cfg(test)]
+mod tests;
+mod redirector;
 
 use rocket::mtls::Certificate;
 
@@ -18,5 +20,7 @@ fn hello() -> &'static str {
 fn rocket() -> _ {
     // See `Rocket.toml` and `Cargo.toml` for TLS configuration.
     // Run `./private/gen_certs.sh` to generate a CA and key pairs.
-    rocket::build().mount("/", routes![hello, mutual])
+    rocket::build()
+        .mount("/", routes![hello, mutual])
+        .attach(redirector::Redirector { port: 3000 })
 }
diff --git a/examples/tls/src/redirector.rs b/examples/tls/src/redirector.rs
new file mode 100644
index 00000000..0aafddf9
--- /dev/null
+++ b/examples/tls/src/redirector.rs
@@ -0,0 +1,65 @@
+//! Redirect all HTTP requests to HTTPs.
+
+use rocket::http::Status;
+use rocket::log::LogLevel;
+use rocket::{route, Error, Request, Data, Route, Orbit, Rocket, Ignite, Config};
+use rocket::fairing::{Fairing, Info, Kind};
+use rocket::response::Redirect;
+
+#[derive(Debug, Copy, Clone)]
+pub struct Redirector {
+    pub port: u16
+}
+
+impl Redirector {
+    // Route function that gets call on every single request.
+    fn redirect<'r>(req: &'r Request, _: Data<'r>) -> route::BoxFuture<'r> {
+        // FIXME: Check the host against a whitelist!
+        if let Some(host) = req.host() {
+            let https_uri = format!("https://{}{}", host, req.uri());
+            route::Outcome::from(req, Redirect::permanent(https_uri)).pin()
+        } else {
+            route::Outcome::from(req, Status::BadRequest).pin()
+        }
+    }
+
+    // Launch an instance of Rocket than handles redirection on `self.port`.
+    pub async fn try_launch(self, mut config: Config) -> Result<Rocket<Ignite>, Error> {
+        use rocket::http::Method::*;
+
+        // Adjust config for redirector: disable TLS, set port, disable logging.
+        config.tls = None;
+        config.port = self.port;
+        config.log_level = LogLevel::Critical;
+
+        // Build a vector of routes to `redirect` on `<path..>` for each method.
+        let redirects = [Get, Put, Post, Delete, Options, Head, Trace, Connect, Patch]
+            .into_iter()
+            .map(|m| Route::new(m, "/<path..>", Self::redirect))
+            .collect::<Vec<_>>();
+
+        rocket::custom(config)
+            .mount("/", redirects)
+            .launch()
+            .await
+    }
+}
+
+#[rocket::async_trait]
+impl Fairing for Redirector {
+    fn info(&self) -> Info {
+        Info { name: "HTTP -> HTTPS Redirector", kind: Kind::Liftoff }
+    }
+
+    async fn on_liftoff(&self, rkt: &Rocket<Orbit>) {
+        let (this, shutdown, config) = (*self, rkt.shutdown(), rkt.config().clone());
+        let _ = rocket::tokio::spawn(async move {
+            if let Err(e) = this.try_launch(config).await {
+                error!("Failed to start HTTP -> HTTPS redirector.");
+                info_!("Error: {}", e);
+                error_!("Shutting down main instance.");
+                shutdown.notify();
+            }
+        });
+    }
+}