1#[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#[derive(Debug, Clone)]
45pub struct AigisServer {
46 router_inner: Router,
47}
48
49#[derive(Debug, Clone)]
51pub struct AigisServerSettings {
52 pub request_timeout: u64,
54
55 pub upstream_settings: UpstreamSettings,
57
58 pub proxy_settings: ProxySettings,
60}
61
62#[derive(Debug, Clone)]
64pub struct ProxySettings {
65 pub allowed_mimetypes: Vec<Mime>,
70
71 pub max_size: u64,
74
75 pub allowed_domains: Option<Vec<Url>>,
79
80 pub max_content_rescale_resolution: u32,
82}
83
84#[derive(Debug, Clone)]
86pub struct UpstreamSettings {
87 pub pass_headers: Option<Vec<String>>,
89
90 pub allow_invalid_certs: bool,
94
95 pub request_timeout: u64,
98
99 pub max_redirects: usize,
101
102 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 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 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}