use std::any::Any;
use std::fmt;
use std::future::Future;
use std::ops::{Deref, DerefMut};
use std::panic::Location;
use std::sync::Arc;
use std::time::Duration;
use either::Either;
use figment::{Figment, Provider};
use futures::TryFutureExt;
use crate::error::{Error, ErrorKind};
use crate::fairing::{Fairing, Fairings};
use crate::http::ext::IntoOwned;
use crate::http::uri::Origin;
use crate::listener::{Bind, DefaultListener, Endpoint, Listener};
use crate::phase::{Build, Building, Ignite, Igniting, Orbit, Orbiting, Phase};
use crate::phase::{State, StateRef, StateRefMut, Stateful};
use crate::router::Router;
use crate::shutdown::{Shutdown, Stages};
use crate::trace::{Trace, TraceAll};
use crate::{sentinel, shield::Shield, Catcher, Config, Route};
pub struct Rocket<P: Phase>(pub(crate) P::State);
impl Rocket<Build> {
#[must_use]
#[inline(always)]
pub fn build() -> Self {
Rocket::custom(Config::figment())
}
#[must_use]
pub fn custom<T: Provider>(provider: T) -> Self {
Rocket::<Build>(Building::default())
.reconfigure(provider)
.attach(Shield::default())
}
#[must_use]
pub fn reconfigure<T: Provider>(mut self, provider: T) -> Self {
self.figment = Figment::from(provider);
crate::trace::init(Config::try_from(&self.figment).ok().as_ref());
span_trace!("reconfigure" => self.figment().trace_trace());
self
}
#[track_caller]
fn load<'a, B, T, F, M>(mut self, kind: &str, base: B, items: Vec<T>, m: M, f: F) -> Self
where
B: TryInto<Origin<'a>> + Clone + fmt::Display,
B::Error: fmt::Display,
M: Fn(&Origin<'a>, T) -> T,
F: Fn(&mut Self, T),
T: Clone + Trace,
{
let mut base = match base.clone().try_into() {
Ok(origin) => origin.into_owned(),
Err(e) => {
error!(%base, location = %Location::caller(), "invalid {kind} base uri: {e}");
panic!("aborting due to {} base error", kind);
}
};
if base.query().is_some() {
warn!(%base, location = %Location::caller(), "query in {kind} base is ignored");
base.clear_query();
}
for unmounted_item in items {
f(&mut self, m(&base, unmounted_item.clone()))
}
self
}
#[must_use]
#[track_caller]
pub fn mount<'a, B, R>(self, base: B, routes: R) -> Self
where
B: TryInto<Origin<'a>> + Clone + fmt::Display,
B::Error: fmt::Display,
R: Into<Vec<Route>>,
{
self.load(
"route",
base,
routes.into(),
|base, route| route.rebase(base.clone()),
|r, route| r.0.routes.push(route),
)
}
#[must_use]
pub fn register<'a, B, C>(self, base: B, catchers: C) -> Self
where
B: TryInto<Origin<'a>> + Clone + fmt::Display,
B::Error: fmt::Display,
C: Into<Vec<Catcher>>,
{
self.load(
"catcher",
base,
catchers.into(),
|base, catcher| catcher.rebase(base.clone()),
|r, catcher| r.0.catchers.push(catcher),
)
}
#[must_use]
pub fn manage<T>(self, state: T) -> Self
where
T: Send + Sync + 'static,
{
let type_name = std::any::type_name::<T>();
if !self.state.set(state) {
error!("state for type '{}' is already being managed", type_name);
panic!("aborting due to duplicated managed state");
}
self
}
#[must_use]
pub fn attach<F: Fairing>(mut self, fairing: F) -> Self {
self.fairings.add(Box::new(fairing));
self
}
pub async fn ignite(mut self) -> Result<Rocket<Ignite>, Error> {
self = Fairings::handle_ignite(self).await;
self.fairings
.audit()
.map_err(|f| ErrorKind::FailedFairings(f.to_vec()))?;
#[allow(unused_mut)]
let mut config = Config::try_from(&self.figment).map_err(ErrorKind::Config)?;
crate::trace::init(&config);
#[cfg(feature = "secrets")]
if !config.secret_key.is_provided() {
if config.profile != Config::DEBUG_PROFILE {
return Err(Error::new(ErrorKind::InsecureSecretKey(
config.profile.clone(),
)));
}
if config.secret_key.is_zero() {
config.secret_key = crate::config::SecretKey::generate()
.unwrap_or_else(crate::config::SecretKey::zero);
}
}
let mut router = Router::new();
self.routes
.clone()
.into_iter()
.for_each(|r| router.routes.push(r));
self.catchers
.clone()
.into_iter()
.for_each(|c| router.catchers.push(c));
let router = router.finalize().map_err(|(r, c)| ErrorKind::Collisions {
routes: r,
catchers: c,
})?;
self.state.freeze();
let fairings = self.fairings.unique_set();
span_info!("config", profile = %self.figment().profile() => {
config.trace_info();
self.figment().trace_debug();
});
span_info!("routes", count = self.routes.len() => self.routes().trace_all_info());
span_info!("catchers", count = self.catchers.len() => self.catchers().trace_all_info());
span_info!("fairings", count = fairings.len() => fairings.trace_all_info());
let rocket: Rocket<Ignite> = Rocket(Igniting {
shutdown: Stages::new(),
figment: self.0.figment,
fairings: self.0.fairings,
state: self.0.state,
router,
config,
});
let sentinels = rocket.routes().flat_map(|r| r.sentinels.iter());
sentinel::query(sentinels, &rocket).map_err(ErrorKind::SentinelAborts)?;
Ok(rocket)
}
}
impl Rocket<Ignite> {
pub fn config(&self) -> &Config {
&self.config
}
pub fn shutdown(&self) -> Shutdown {
self.shutdown.start.clone()
}
pub(crate) fn into_orbit(self, endpoints: Vec<Endpoint>) -> Rocket<Orbit> {
Rocket(Orbiting {
endpoints,
router: self.0.router,
fairings: self.0.fairings,
figment: self.0.figment,
config: self.0.config,
state: self.0.state,
shutdown: self.0.shutdown,
})
}
async fn _local_launch(self, endpoint: Endpoint) -> Rocket<Orbit> {
let rocket = self.into_orbit(vec![endpoint]);
Rocket::liftoff(&rocket).await;
rocket
}
async fn _launch<L: Listener + 'static>(self, listener: L) -> Result<Rocket<Ignite>, Error> {
let rocket = self
.listen_and_serve(listener, |rocket| async move {
let rocket = Arc::new(rocket);
rocket.shutdown.spawn_listener(&rocket.config.shutdown);
if let Err(e) = tokio::spawn(Rocket::liftoff(rocket.clone())).await {
let rocket = rocket.try_wait_shutdown().await.map(Box::new);
return Err(ErrorKind::Liftoff(rocket, e).into());
}
Ok(rocket)
})
.await?;
Ok(rocket
.try_wait_shutdown()
.await
.map_err(ErrorKind::Shutdown)?)
}
}
impl Rocket<Orbit> {
async fn try_wait_shutdown(self: Arc<Self>) -> Result<Rocket<Ignite>, Arc<Self>> {
info!("Shutting down. Waiting for shutdown fairings and pending I/O...");
tokio::spawn({
let rocket = self.clone();
async move { rocket.fairings.handle_shutdown(&rocket).await }
});
let config = &self.config.shutdown;
let wait = Duration::from_micros(250);
for period in [wait, config.grace(), wait, config.mercy(), wait * 4] {
if Arc::strong_count(&self) == 1 {
break;
}
tokio::time::sleep(period).await;
}
match Arc::try_unwrap(self) {
Ok(rocket) => {
info!("Graceful shutdown completed successfully.");
Ok(rocket.deorbit())
}
Err(rocket) => {
warn!("Shutdown failed: outstanding background I/O.");
Err(rocket)
}
}
}
pub(crate) fn deorbit(self) -> Rocket<Ignite> {
Rocket(Igniting {
router: self.0.router,
fairings: self.0.fairings,
figment: self.0.figment,
config: self.0.config,
state: self.0.state,
shutdown: self.0.shutdown,
})
}
pub(crate) async fn liftoff<R: Deref<Target = Self>>(rocket: R) {
let rocket = rocket.deref();
rocket.fairings.handle_liftoff(rocket).await;
if !crate::running_within_rocket_async_rt().await {
warn!(
"Rocket is executing inside of a custom runtime.\n\
Rocket's runtime is enabled via `#[rkt::main]` or `#[launch]`\n\
Forced shutdown is disabled. Runtime settings may be suboptimal."
);
}
tracing::info!(name: "liftoff", endpoint = %rocket.endpoints[0]);
}
pub fn config(&self) -> &Config {
&self.config
}
pub fn endpoints(&self) -> impl Iterator<Item = &Endpoint> {
self.endpoints.iter()
}
pub fn shutdown(&self) -> Shutdown {
self.shutdown.start.clone()
}
}
impl<P: Phase> Rocket<P> {
pub fn routes(&self) -> impl Iterator<Item = &Route> {
match self.0.as_ref() {
StateRef::Build(p) => Either::Left(p.routes.iter()),
StateRef::Ignite(p) => Either::Right(p.router.routes.iter()),
StateRef::Orbit(p) => Either::Right(p.router.routes.iter()),
}
}
pub fn catchers(&self) -> impl Iterator<Item = &Catcher> {
match self.0.as_ref() {
StateRef::Build(p) => Either::Left(p.catchers.iter()),
StateRef::Ignite(p) => Either::Right(p.router.catchers.iter()),
StateRef::Orbit(p) => Either::Right(p.router.catchers.iter()),
}
}
pub fn state<T: Send + Sync + 'static>(&self) -> Option<&T> {
match self.0.as_ref() {
StateRef::Build(p) => p.state.try_get(),
StateRef::Ignite(p) => p.state.try_get(),
StateRef::Orbit(p) => p.state.try_get(),
}
}
pub fn fairing<F: Fairing>(&self) -> Option<&F> {
match self.0.as_ref() {
StateRef::Build(p) => p.fairings.filter::<F>().next(),
StateRef::Ignite(p) => p.fairings.filter::<F>().next(),
StateRef::Orbit(p) => p.fairings.filter::<F>().next(),
}
}
pub fn fairings<F: Fairing>(&self) -> impl Iterator<Item = &F> {
match self.0.as_ref() {
StateRef::Build(p) => Either::Left(p.fairings.filter::<F>()),
StateRef::Ignite(p) => Either::Right(p.fairings.filter::<F>()),
StateRef::Orbit(p) => Either::Right(p.fairings.filter::<F>()),
}
}
pub fn fairing_mut<F: Fairing>(&mut self) -> Option<&mut F> {
match self.0.as_mut() {
StateRefMut::Build(p) => p.fairings.filter_mut::<F>().next(),
StateRefMut::Ignite(p) => p.fairings.filter_mut::<F>().next(),
StateRefMut::Orbit(p) => p.fairings.filter_mut::<F>().next(),
}
}
pub fn fairings_mut<F: Fairing>(&mut self) -> impl Iterator<Item = &mut F> {
match self.0.as_mut() {
StateRefMut::Build(p) => Either::Left(p.fairings.filter_mut::<F>()),
StateRefMut::Ignite(p) => Either::Right(p.fairings.filter_mut::<F>()),
StateRefMut::Orbit(p) => Either::Right(p.fairings.filter_mut::<F>()),
}
}
pub fn figment(&self) -> &Figment {
match self.0.as_ref() {
StateRef::Build(p) => &p.figment,
StateRef::Ignite(p) => &p.figment,
StateRef::Orbit(p) => &p.figment,
}
}
async fn into_ignite(self) -> Result<Rocket<Ignite>, Error> {
match self.0.into_state() {
State::Build(s) => Rocket::from(s).ignite().await,
State::Ignite(s) => Ok(Rocket::from(s)),
State::Orbit(s) => Ok(Rocket::from(s).deorbit()),
}
}
pub(crate) async fn local_launch(self, e: Endpoint) -> Result<Rocket<Orbit>, Error> {
Ok(self.into_ignite().await?._local_launch(e).await)
}
pub async fn launch(self) -> Result<Rocket<Ignite>, Error> {
self.launch_with::<DefaultListener>().await
}
pub async fn launch_with<B: Bind>(self) -> Result<Rocket<Ignite>, Error> {
let rocket = self.into_ignite().await?;
let bind_endpoint = B::bind_endpoint(&rocket).ok();
let listener: B = B::bind(&rocket)
.await
.map_err(|e| ErrorKind::Bind(bind_endpoint, Box::new(e)))?;
let any: Box<dyn Any + Send + Sync> = Box::new(listener);
match any.downcast::<DefaultListener>() {
Ok(listener) => {
let listener = *listener;
crate::util::for_both!(listener, listener => {
crate::util::for_both!(listener, listener => {
rocket._launch(listener).await
})
})
}
Err(any) => {
let listener = *any.downcast::<B>().unwrap();
rocket._launch(listener).await
}
}
}
pub async fn try_launch_on<L, F, E>(self, listener: F) -> Result<Rocket<Ignite>, Error>
where
L: Listener + 'static,
F: Future<Output = Result<L, E>>,
E: std::error::Error + Send + 'static,
{
let listener = listener
.map_err(|e| ErrorKind::Bind(None, Box::new(e)))
.await?;
self.into_ignite().await?._launch(listener).await
}
pub async fn launch_on<L>(self, listener: L) -> Result<Rocket<Ignite>, Error>
where
L: Listener + 'static,
{
self.into_ignite().await?._launch(listener).await
}
}
#[doc(hidden)]
impl<P: Phase> Deref for Rocket<P> {
type Target = P::State;
fn deref(&self) -> &Self::Target {
&self.0
}
}
#[doc(hidden)]
impl<P: Phase> DerefMut for Rocket<P> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl<P: Phase> fmt::Debug for Rocket<P> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}