teapot/src/main.rs

mod api;
mod cache;
mod config;
mod error;
mod routes;
mod transcode;
mod translation;
mod types;
mod utils;
mod views;

use std::{
   env,
   net::{
      IpAddr,
      SocketAddr,
   },
   sync::Arc,
   time::Duration,
};

use axum::{
   Router,
   ServiceExt,
   extract::Request,
   http::{
      StatusCode,
      header::HeaderValue,
   },
   middleware,
};
use hyper::header;
use tokio::net::TcpListener;
use tower::{
   Layer as _,
   limit::ConcurrencyLimitLayer,
};
use tower_http::{
   normalize_path::NormalizePathLayer,
   services::{
      ServeDir,
      ServeFile,
   },
   set_header::SetResponseHeaderLayer,
   timeout::TimeoutLayer,
};
use tracing_subscriber::{
   fmt,
   layer::SubscriberExt as _,
   util::SubscriberInitExt as _,
};

use crate::{
   api::{
      auth::SessionPool,
      client::ApiClient,
      http::HttpClient,
   },
   cache::store::Cache,
   config::{
      Config,
      GifTranscodingMode,
   },
   transcode::GifTranscoder,
};

/// Application state shared across all routes.
#[derive(Clone)]
pub struct AppState {
   pub config:              Arc<Config>,
   pub cache:               Cache,
   pub api:                 ApiClient,
   pub http_client:         HttpClient,
   pub gif_transcoder:      Option<Arc<GifTranscoder>>,
   pub translation_limiter: Arc<translation::TranslationLimiter>,
   /// Peers permitted to name a different caller through forwarding headers.
   pub trusted_proxies:     Arc<[IpAddr]>,
}

/// The strict policy, and the one Cloudflare needs.
///
/// Cloudflare injects its JS-detection bootstrap inline into origin HTML, so
/// `script-src 'self'` stops it executing; the challenge then never records a
/// result and is re-served on every request. Allowing it costs an
/// `'unsafe-inline'`, so it stays opt-in rather than becoming the default.
const fn content_security_policy(behind_cloudflare: bool) -> &'static str {
   if behind_cloudflare {
      "default-src 'none'; script-src 'self' 'unsafe-inline' https://challenges.cloudflare.com \
       https://static.cloudflareinsights.com; style-src 'self' 'unsafe-inline'; img-src 'self' \
       data:; media-src 'self' blob:; font-src 'self'; connect-src 'self' \
       https://challenges.cloudflare.com https://cloudflareinsights.com; frame-src \
       https://challenges.cloudflare.com; form-action 'self'; base-uri 'self'; frame-ancestors \
       'none'"
   } else {
      "default-src 'none'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' \
       data:; media-src 'self' blob:; font-src 'self'; connect-src 'self'; form-action 'self'; \
       base-uri 'self'; frame-ancestors 'none'"
   }
}

#[tokio::main]
async fn main() -> eyre::Result<()> {
   tracing_subscriber::registry()
      .with(
         tracing_subscriber::EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| "teapot=debug,tower_http=debug".into()),
      )
      .with(fmt::layer())
      .init();

   let config_path =
      env::var("TEAPOT_CONF_FILE").unwrap_or_else(|_| "config/teapot.toml".to_owned());
   let config = Config::load(&config_path)?;
   let config = Arc::new(config);

   tracing::info!(
      "Starting {} on {}:{}",
      config.server.title,
      config.server.address,
      config.server.port
   );

   error::set_messages(error::Messages {
      client_budget:  config.config.client_budget_message.clone(),
      rate_limited:   config.config.rate_limited_message.clone(),
      internal_error: config.config.internal_error_message.clone(),
   });

   let cache = Cache::new(config.cache.max_entries);

   let sessions_path =
      env::var("TEAPOT_SESSIONS_FILE").unwrap_or_else(|_| "sessions.jsonl".to_owned());
   let sessions = SessionPool::load(&sessions_path, config.config.max_concurrent_reqs).await?;

   let api = ApiClient::new(&config, sessions);

   let http_client = HttpClient::new(&config.config.proxy, &config.config.proxy_auth);
   let gif_transcoder = if config.gif_transcoding.mode == GifTranscodingMode::Local {
      match GifTranscoder::new(http_client.clone(), config.gif_transcoding.clone()).await {
         Ok(transcoder) => {
            tracing::info!("GIF transcoder enabled (local mode)");
            Some(Arc::new(transcoder))
         },
         Err(err) => {
            tracing::error!("Failed to initialize GIF transcoder: {err}");
            None
         },
      }
   } else {
      if config.gif_transcoding.mode == GifTranscodingMode::External {
         tracing::info!(
            "GIF transcoding enabled (external mode, domain: {})",
            config.gif_transcoding.external_domain
         );
      }
      None
   };

   // Validated at load, so anything here parses.
   let trusted_proxies = config
      .config
      .trusted_proxies
      .iter()
      .filter_map(|proxy| proxy.parse::<IpAddr>().ok())
      .collect::<Vec<_>>();
   if trusted_proxies.is_empty() {
      tracing::info!("No trustedProxies configured, billing clients by socket address");
   }

   let state = AppState {
      config: Arc::clone(&config),
      cache,
      api,
      http_client,
      gif_transcoder,
      translation_limiter: Arc::new(translation::TranslationLimiter::new()),
      trusted_proxies: trusted_proxies.into(),
   };

   // Build the router with specific routes before static files
   let static_dir = config.server.static_dir.clone();
   let app = Router::new()
      .merge(routes::app::router())
      .nest_service("/public", ServeDir::new(&static_dir))
      .nest_service("/css", ServeDir::new(format!("{static_dir}/css")))
      .nest_service("/js", ServeDir::new(format!("{static_dir}/js")))
      .nest_service("/fonts", ServeDir::new(format!("{static_dir}/fonts")))
      .nest_service("/md", ServeDir::new(format!("{static_dir}/md")))
      .route_service(
         "/logo.svg",
         ServeFile::new(format!("{static_dir}/logo.svg")),
      )
      .route_service(
         "/logo.png",
         ServeFile::new(format!("{static_dir}/logo.png")),
      )
      .route_service(
         "/favicon.ico",
         ServeFile::new(format!("{static_dir}/favicon.ico")),
      )
      .route_service(
         "/favicon-16x16.png",
         ServeFile::new(format!("{static_dir}/favicon-16x16.png")),
      )
      .route_service(
         "/favicon-32x32.png",
         ServeFile::new(format!("{static_dir}/favicon-32x32.png")),
      )
      .route_service(
         "/apple-touch-icon.png",
         ServeFile::new(format!("{static_dir}/apple-touch-icon.png")),
      )
      .route_service(
         "/site.webmanifest",
         ServeFile::new(format!("{static_dir}/site.webmanifest")),
      )
      .route_service(
         "/robots.txt",
         ServeFile::new(format!("{static_dir}/robots.txt")),
      )
      .route_service(
         "/opensearch.xml",
         ServeFile::new(format!("{static_dir}/opensearch.xml")),
      )
      .layer(middleware::from_fn(routes::middleware::prefs_middleware))
      .layer(middleware::from_fn(routes::middleware::snowflake_guard))
      .layer(middleware::from_fn_with_state(
         state.clone(),
         routes::middleware::client_middleware,
      ))
      .layer(SetResponseHeaderLayer::overriding(
         header::REFERRER_POLICY,
         HeaderValue::from_static("no-referrer"),
      ))
      .layer(SetResponseHeaderLayer::if_not_present(
         header::CONTENT_SECURITY_POLICY,
         HeaderValue::from_static(content_security_policy(config.config.behind_cloudflare)),
      ))
      .layer(SetResponseHeaderLayer::if_not_present(
         header::X_CONTENT_TYPE_OPTIONS,
         HeaderValue::from_static("nosniff"),
      ))
      .layer(SetResponseHeaderLayer::if_not_present(
         header::STRICT_TRANSPORT_SECURITY,
         HeaderValue::from_static("max-age=31536000; includeSubDomains"),
      ))
      .layer(TimeoutLayer::with_status_code(
         StatusCode::REQUEST_TIMEOUT,
         Duration::from_secs(45),
      ))
      .layer(ConcurrencyLimitLayer::new(
         usize::try_from(config.server.http_max_connections).unwrap_or(usize::MAX),
      ))
      .with_state(state);

   let addr = SocketAddr::new(config.server.address.parse()?, config.server.port);
   let listener = TcpListener::bind(addr).await?;

   tracing::info!("Listening on {addr}");
   axum::serve(
      listener,
      ServiceExt::<Request>::into_make_service_with_connect_info::<SocketAddr>(
         NormalizePathLayer::trim_trailing_slash().layer(app),
      ),
   )
   .await?;

   Ok(())
}

#[cfg(test)]
mod tests {
   use axum::http::HeaderValue;

   use super::content_security_policy;

   #[test]
   fn both_policies_are_valid_header_values() {
      for behind_cloudflare in [false, true] {
         let policy = content_security_policy(behind_cloudflare);
         assert!(
            HeaderValue::from_static(policy).to_str().is_ok(),
            "{policy}"
         );
         // Line continuations in the source must not leak into the header.
         assert!(!policy.contains('\n'), "{policy}");
         assert!(!policy.contains("  "), "{policy}");
      }
   }

   #[test]
   fn the_default_policy_still_forbids_inline_scripts() {
      let policy = content_security_policy(false);
      assert!(policy.contains("script-src 'self';"));
      assert!(!policy.contains("script-src 'self' 'unsafe-inline'"));
      assert!(!policy.contains("cloudflare"));
   }

   /// Cloudflare's bootstrap is inline and loads from its own host, so both
   /// have to be allowed or the challenge loops.
   #[test]
   fn the_cloudflare_policy_admits_the_challenge_bootstrap() {
      let policy = content_security_policy(true);
      assert!(
         policy.contains("script-src 'self' 'unsafe-inline' https://challenges.cloudflare.com")
      );
      assert!(policy.contains("frame-src https://challenges.cloudflare.com"));
      // Web Analytics loads its beacon from a second host and reports to a
      // third, so allowing only the challenge host still breaks the page.
      assert!(policy.contains("https://static.cloudflareinsights.com"));
      assert!(policy.contains(
         "connect-src 'self' https://challenges.cloudflare.com https://cloudflareinsights.com"
      ));
      // Everything else stays as strict as the default.
      assert!(policy.contains("default-src 'none'"));
      assert!(policy.contains("frame-ancestors 'none'"));
      assert!(policy.contains("base-uri 'self'"));
   }
}