From 29cd271f8ad3035f2aa3e396d59e0864f37cf149 Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Thu, 18 May 2023 17:33:57 -0700
Subject: [PATCH] Only extract needed values in 'async_main'.

Previously, `async_main` would extract a full `Config`. This mean that values
like `address` were read and parsed even when they were unused. Should they
exist and be malformed, a configuration error would needlessly arise.

This commit fixes this by only extract values that are subsequently used.
---
 core/lib/src/config/config.rs | 13 ++++++++-----
 core/lib/src/config/mod.rs    |  3 ++-
 core/lib/src/lib.rs           | 21 ++++++++++++++++-----
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/core/lib/src/config/config.rs b/core/lib/src/config/config.rs
index 6fbd6e7e..993d7c00 100644
--- a/core/lib/src/config/config.rs
+++ b/core/lib/src/config/config.rs
@@ -306,10 +306,7 @@ impl Config {
     /// let config = Config::from(figment);
     /// ```
     pub fn from<T: Provider>(provider: T) -> Self {
-        Self::try_from(provider).unwrap_or_else(|e| {
-            pretty_print_error(e);
-            panic!("aborting due to configuration error(s)")
-        })
+        Self::try_from(provider).unwrap_or_else(bail_with_config_error)
     }
 
     /// Returns `true` if TLS is enabled.
@@ -588,12 +585,18 @@ impl<'r> FromRequest<'r> for &'r Config {
     }
 }
 
+#[doc(hidden)]
+pub fn bail_with_config_error<T>(error: figment::Error) -> T {
+    pretty_print_error(error);
+    panic!("aborting due to configuration error(s)")
+}
+
 #[doc(hidden)]
 pub fn pretty_print_error(error: figment::Error) {
     use figment::error::{Kind, OneOf};
 
     crate::log::init_default();
-    error!("Rocket configuration extraction from provider failed.");
+    error!("Failed to extract valid configuration.");
     for e in error {
         fn w<T: std::fmt::Display>(v: T) -> Paint<T> { Paint::white(v) }
 
diff --git a/core/lib/src/config/mod.rs b/core/lib/src/config/mod.rs
index 00a2613c..a811cf99 100644
--- a/core/lib/src/config/mod.rs
+++ b/core/lib/src/config/mod.rs
@@ -123,7 +123,8 @@ mod tls;
 mod secret_key;
 
 #[doc(hidden)]
-pub use config::pretty_print_error;
+pub use config::{pretty_print_error, bail_with_config_error};
+
 pub use config::Config;
 pub use crate::log::LogLevel;
 pub use shutdown::Shutdown;
diff --git a/core/lib/src/lib.rs b/core/lib/src/lib.rs
index e1b8a260..f3a11ddf 100644
--- a/core/lib/src/lib.rs
+++ b/core/lib/src/lib.rs
@@ -261,11 +261,22 @@ pub fn async_test<R>(fut: impl std::future::Future<Output = R>) -> R {
 /// WARNING: This is unstable! Do not use this method outside of Rocket!
 #[doc(hidden)]
 pub fn async_main<R>(fut: impl std::future::Future<Output = R> + Send) -> R {
-    // FIXME: These config values won't reflect swaps of `Rocket` in attach
-    // fairings with different config values, or values from non-Rocket configs.
-    // See tokio-rs/tokio#3329 for a necessary solution in `tokio`.
-    let c = Config::from(Config::figment());
-    async_run(fut, c.workers, c.max_blocking, c.shutdown.force, "rocket-worker-thread")
+    // FIXME: We need to run `fut` to get the user's `Figment` to properly set
+    // up the async env, but we need the async env to run `fut`. So we're stuck.
+    // Tokio doesn't let us take the state from one async env and migrate it to
+    // another, so we need to use one, making this impossible.
+    //
+    // So as a result, we only use values from Rocket's figment. These
+    // values won't reflect swaps of `Rocket` in attach fairings with different
+    // config values, or values from non-Rocket configs. See tokio-rs/tokio#3329
+    // for a necessary resolution in `tokio`.
+    use config::bail_with_config_error as bail;
+
+    let fig = Config::figment();
+    let workers = fig.extract_inner(Config::WORKERS).unwrap_or_else(bail);
+    let max_blocking = fig.extract_inner(Config::MAX_BLOCKING).unwrap_or_else(bail);
+    let force = fig.focus(Config::SHUTDOWN).extract_inner("force").unwrap_or_else(bail);
+    async_run(fut, workers, max_blocking, force, "rocket-worker-thread")
 }
 
 /// Executes a `future` to completion on a new tokio-based Rocket async runtime.