aigis/
lib.rs

1//! Crate for Aigis, a simple and configurable content proxy.
2
3#[cfg(feature = "rustls-tls")]
4#[cfg(feature = "native-tls")]
5compile_error!("You can only enable one TLS backend");
6
7pub extern crate mime;
8pub extern crate url;
9
10mod http_client;
11mod middleware;
12mod mime_util;
13mod routes;
14
15use crate::http_client::{build_http_client, BuildHttpClientArgs};
16use anyhow::Result;
17use axum::{middleware as axum_middleware, routing::get, Router};
18use http_client::HttpClient;
19use mime::Mime;
20use routes::{HEALTH_ENDPOINT, INDEX_ENDPOINT, PROXY_ENDPOINT};
21use std::{net::SocketAddr, time::Duration};
22use tokio::net::TcpListener;
23use tower_http::{
24    catch_panic::CatchPanicLayer,
25    normalize_path::NormalizePathLayer,
26    timeout::TimeoutLayer,
27    trace::{self, TraceLayer},
28};
29use tracing::{info, Level};
30use url::Url;
31
32/// The Aigis server itself.
33/// # Example
34/// ```rust,no_run
35/// use std::net::{SocketAddr, IpAddr, Ipv4Addr};
36/// use aigis::{AigisServer, AigisServerSettings};
37///
38/// # #[tokio::main]
39/// # async fn main() {
40/// let server = AigisServer::new(AigisServerSettings::default()).unwrap();
41/// server.start(&SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080)).await.unwrap();
42/// # }
43/// ```
44#[derive(Debug, Clone)]
45pub struct AigisServer {
46    router_inner: Router,
47}
48
49/// Settings to run the Aigis server with.
50#[derive(Debug, Clone)]
51pub struct AigisServerSettings {
52    /// How many seconds that can elapse before a request is abandoned for taking too long.
53    pub request_timeout: u64,
54
55    /// See [`UpstreamSettings`].
56    pub upstream_settings: UpstreamSettings,
57
58    /// See [`ProxySettings`].
59    pub proxy_settings: ProxySettings,
60}
61
62/// Configuration options used for the `proxy` route.
63#[derive(Debug, Clone)]
64pub struct ProxySettings {
65    /// [`Mime`]s that are allowed to be proxied, checked against the Content-Type header
66    /// received from the upstream server.
67    ///
68    /// Supports type wildcards such as 'image/*'.
69    pub allowed_mimetypes: Vec<Mime>,
70
71    /// The maximum Content-Lenth that can be proxied.
72    /// Anything larger than this value will not be sent and an error will shown instead.
73    pub max_size: u64,
74
75    /// [`Url`]s that are allowed to be proxied.
76    ///
77    /// Does not support subdomain wildcards, each domain must be added seperately.
78    pub allowed_domains: Option<Vec<Url>>,
79
80    /// The maximum resolution that can be requested for content that supports resizing.
81    pub max_content_rescale_resolution: u32,
82}
83
84/// Configuration options used when making any call to an upstream service regardless of route.
85#[derive(Debug, Clone)]
86pub struct UpstreamSettings {
87    /// Headers that will be passed on from the client to the upstream server verbatim.
88    pub pass_headers: Option<Vec<String>>,
89
90    /// Whether or not to allow invalid/expired/forged TLS certificates when making upstream requests.
91    ///
92    /// **Enabling this is dangerous and is usually not necessary.**
93    pub allow_invalid_certs: bool,
94
95    /// How many seconds that can elapse after sending a request to an upstream server before it's abandoned
96    /// and considered failed.
97    pub request_timeout: u64,
98
99    /// The maximum amount of redirects to follow when making a request to an upstream server before abandoning the request.
100    pub max_redirects: usize,
101
102    /// Whether or not to send the client the `Cache-Control` header value that was received when making the
103    /// request to the upstream server if one is available.
104    ///
105    /// If one of the `cache-*` crate features are enabled the request will already be cached server-side for that requested duration,
106    /// so sending the `Cache-Control` header to the client is favourable behaviour as it can sometimes lighten server load.
107    pub use_received_cache_headers: bool,
108}
109
110impl Default for AigisServerSettings {
111    fn default() -> Self {
112        Self {
113            request_timeout: 10,
114            proxy_settings: ProxySettings::default(),
115            upstream_settings: UpstreamSettings::default(),
116        }
117    }
118}
119
120impl Default for ProxySettings {
121    fn default() -> Self {
122        Self {
123            allowed_mimetypes: vec![mime::IMAGE_STAR],
124            allowed_domains: None,
125            max_size: 100000000,
126            max_content_rescale_resolution: 1024,
127        }
128    }
129}
130
131impl Default for UpstreamSettings {
132    fn default() -> Self {
133        Self {
134            allow_invalid_certs: false,
135            max_redirects: 10,
136            pass_headers: None,
137            request_timeout: 30,
138            use_received_cache_headers: true,
139        }
140    }
141}
142
143#[derive(Debug, Clone)]
144struct AppState {
145    client: HttpClient,
146    settings: AigisServerSettings,
147}
148
149impl AigisServer {
150    /// Create a new server with the provided settings.
151    pub fn new(settings: AigisServerSettings) -> Result<Self> {
152        let router = Router::new()
153            .route(PROXY_ENDPOINT, get(routes::proxy_handler))
154            .route(INDEX_ENDPOINT, get(routes::index_handler))
155            .route(HEALTH_ENDPOINT, get(routes::health_handler))
156            .layer(
157                TraceLayer::new_for_http()
158                    .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
159                    .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
160            )
161            .layer(TimeoutLayer::new(Duration::from_secs(
162                settings.request_timeout,
163            )))
164            .layer(NormalizePathLayer::trim_trailing_slash())
165            .layer(CatchPanicLayer::new())
166            .layer(axum_middleware::from_fn(middleware::header_middleware))
167            .with_state(AppState {
168                client: build_http_client(BuildHttpClientArgs {
169                    allow_invalid_certs: settings.upstream_settings.allow_invalid_certs,
170                    max_redirects: settings.upstream_settings.max_redirects,
171                    request_timeout: Duration::from_secs(
172                        settings.upstream_settings.request_timeout,
173                    ),
174                })?,
175                settings,
176            });
177
178        Ok(Self {
179            router_inner: router,
180        })
181    }
182
183    /// Start the server and expose it locally on the provided [`SocketAddr`].
184    pub async fn start(self, address: &SocketAddr) -> Result<()> {
185        let tcp_listener = TcpListener::bind(&address).await?;
186        info!("Listening on http://{}", tcp_listener.local_addr()?);
187        axum::serve(tcp_listener, self.router_inner)
188            .with_graceful_shutdown(async {
189                tokio::signal::ctrl_c()
190                    .await
191                    .expect("failed to listen for ctrl-c");
192            })
193            .await?;
194
195        Ok(())
196    }
197}