use std::sync::Arc;
use std::time::Duration;
use rustlavel_client::Client;
use rustlavel_http::{BoxFuture, Plugin, Request, Response, Setup, Status};
use crate::forward;
use crate::health;
use crate::route::{Routes, Upstream};
pub struct Gateway {
routes: Routes,
client: Client,
health: Option<(String, String)>,
require_token: bool,
}
impl Default for Gateway {
fn default() -> Self {
Gateway::new()
}
}
impl Gateway {
pub fn new() -> Gateway {
Gateway {
routes: Routes::new(),
client: Client::new().timeout(Duration::from_secs(30)),
health: None,
require_token: false,
}
}
pub fn route(mut self, pattern: impl Into<String>, upstream: Upstream) -> Gateway {
self.routes.add(pattern, upstream);
self
}
pub fn client(mut self, client: Client) -> Gateway {
self.client = client;
self
}
pub fn require_token(mut self) -> Gateway {
self.require_token = true;
self
}
pub fn health(mut self, at: impl Into<String>, probe: impl Into<String>) -> Gateway {
self.health = Some((at.into(), probe.into()));
self
}
pub fn routes(&self) -> &Routes {
&self.routes
}
}
impl Plugin for Gateway {
fn name(&self) -> &'static str {
"gateway"
}
fn register(self: Box<Self>, setup: &mut Setup<'_>) {
let gateway = Arc::new(*self);
if let Some((at, probe)) = gateway.health.clone() {
let owner = Arc::clone(&gateway);
setup.router.get(&at, move |_req: Request| {
let owner = Arc::clone(&owner);
let probe = probe.clone();
Box::pin(async move { health::check(&owner.client, &owner.routes, &probe).await })
as BoxFuture<Response>
});
}
setup.router.fallback(move |request: Request| {
let gateway = Arc::clone(&gateway);
Box::pin(async move {
let Some(upstream) = gateway.routes.find(request.path()) else {
return Response::not_found();
};
if gateway.require_token
&& !upstream.is_open()
&& request.header("authorization").is_none()
{
return Response::new(Status::UNAUTHORIZED).with_json(
rustlavel_core::Json::object([(
"message",
rustlavel_core::Json::from("Unauthenticated."),
)]),
);
}
let Some(upstream) = upstream.resolved().await else {
rustlavel_core::warn!(
"gateway: {} has no instances",
upstream.service_name().unwrap_or("an upstream")
);
return Response::new(Status(503)).with_json(rustlavel_core::Json::object([(
"message",
rustlavel_core::Json::from("No instances are available."),
)]));
};
match forward::forward(&gateway.client, &upstream, request).await {
Ok(response) => response,
Err(error) => {
rustlavel_core::error!("gateway failed to forward: {error}");
Response::new(Status::INTERNAL_ERROR)
}
}
}) as BoxFuture<Response>
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn routes_are_matched_by_specificity_through_the_builder() {
let gateway = Gateway::new()
.route("/api/*", Upstream::at("http://api"))
.route("/api/billing/*", Upstream::at("http://billing"));
assert_eq!(
gateway.routes().find("/api/billing/x").map(Upstream::base),
Some("http://billing")
);
assert_eq!(gateway.routes().find("/api/x").map(Upstream::base), Some("http://api"));
}
#[test]
fn a_health_endpoint_is_mounted_only_when_asked_for() {
assert!(Gateway::new().health.is_none());
let gateway = Gateway::new().health("/health", "/up");
assert_eq!(gateway.health.as_ref().map(|(at, probe)| (at.as_str(), probe.as_str())),
Some(("/health", "/up")));
}
#[test]
fn a_token_is_not_required_unless_it_is_asked_for() {
assert!(!Gateway::new().require_token);
assert!(Gateway::new().require_token().require_token);
}
#[test]
fn an_authorization_server_can_be_excepted_from_the_token_requirement() {
let gateway = Gateway::new()
.require_token()
.route("/oauth/*", Upstream::at("http://auth").open())
.route("/api/*", Upstream::at("http://api"));
assert!(
gateway.routes().find("/oauth/token").expect("a route").is_open(),
"the token endpoint is behind the credential it issues"
);
assert!(
!gateway.routes().find("/api/orders").expect("a route").is_open(),
"opening one route opened another"
);
}
#[test]
fn an_upstream_is_closed_until_it_is_opened() {
assert!(!Upstream::at("http://api").is_open());
assert!(Upstream::at("http://api").open().is_open());
assert!(Upstream::at("http://api").open().strip("/x").host("h").is_open());
}
}