From 5fd8e86d9739c63ae0db46e99b8ef6fa1ddadba7 Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Sat, 22 May 2021 23:06:05 -0400 Subject: [PATCH] Add support for offline mode --- src/config.rs | 16 +++++++-- src/main.rs | 96 ++++++++++++++++++++++++++++++++++++-------------- src/metrics.rs | 21 +++++++---- src/routes.rs | 21 +++++++++-- src/state.rs | 15 ++++++-- 5 files changed, 131 insertions(+), 38 deletions(-) diff --git a/src/config.rs b/src/config.rs index 87df0e2..73466b9 100644 --- a/src/config.rs +++ b/src/config.rs @@ -13,6 +13,8 @@ pub static VALIDATE_TOKENS: AtomicBool = AtomicBool::new(false); // everywhere. pub static SEND_SERVER_VERSION: AtomicBool = AtomicBool::new(false); +pub static OFFLINE_MODE: AtomicBool = AtomicBool::new(false); + #[derive(Clap, Clone)] #[clap(version = crate_version!(), author = crate_authors!(), about = crate_description!())] pub struct CliArgs { @@ -82,17 +84,25 @@ pub enum UnstableOptions { /// Disables token validation. Don't use this unless you know the /// ramifications of this command. DisableTokenValidation, + + /// Tries to run without communication to MangaDex. + OfflineMode, + + /// Serves HTTP in plaintext + DisableTls, } impl FromStr for UnstableOptions { - type Err = &'static str; + type Err = String; fn from_str(s: &str) -> Result { match s { "override-upstream" => Ok(Self::OverrideUpstream), "use-lfu" => Ok(Self::UseLfu), "disable-token-validation" => Ok(Self::DisableTokenValidation), - _ => Err("Unknown unstable option"), + "offline-mode" => Ok(Self::OfflineMode), + "disable-tls" => Ok(Self::DisableTls), + _ => Err(format!("Unknown unstable option '{}'", s)), } } } @@ -103,6 +113,8 @@ impl Display for UnstableOptions { Self::OverrideUpstream => write!(f, "override-upstream"), Self::UseLfu => write!(f, "use-lfu"), Self::DisableTokenValidation => write!(f, "disable-token-validation"), + Self::OfflineMode => write!(f, "offline-mode"), + Self::DisableTls => write!(f, "disable-tls"), } } } diff --git a/src/main.rs b/src/main.rs index 50dc3d6..d15086d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -29,7 +29,7 @@ use thiserror::Error; use crate::cache::mem::{Lfu, Lru}; use crate::cache::{MemoryCache, ENCRYPTION_KEY}; -use crate::config::UnstableOptions; +use crate::config::{UnstableOptions, OFFLINE_MODE}; use crate::state::DynamicServerCert; mod cache; @@ -60,8 +60,12 @@ async fn main() -> Result<(), Box> { sodiumoxide::init().expect("Failed to initialize crypto"); // It's ok to fail early here, it would imply we have a invalid config. dotenv::dotenv().ok(); - let cli_args = CliArgs::parse(); + // + // Config loading + // + + let cli_args = CliArgs::parse(); let port = cli_args.port; let memory_max_size = cli_args .memory_quota @@ -71,6 +75,19 @@ async fn main() -> Result<(), Box> { let cache_path = cli_args.cache_path.clone(); let low_mem_mode = cli_args.low_memory; let use_lfu = cli_args.unstable_options.contains(&UnstableOptions::UseLfu); + let disable_tls = cli_args + .unstable_options + .contains(&UnstableOptions::DisableTls); + OFFLINE_MODE.store( + cli_args + .unstable_options + .contains(&UnstableOptions::OfflineMode), + Ordering::Release, + ); + + // + // Logging and warnings + // let log_level = match (cli_args.quiet, cli_args.verbose) { (n, _) if n > 2 => LevelFilter::Off, @@ -103,9 +120,15 @@ async fn main() -> Result<(), Box> { ENCRYPTION_KEY.set(gen_key()).unwrap(); } + metrics::init(); + // HTTP Server init - let server = ServerState::init(&client_secret, &cli_args).await?; + let server = if OFFLINE_MODE.load(Ordering::Acquire) { + ServerState::init_offline() + } else { + ServerState::init(&client_secret, &cli_args).await? + }; let data_0 = Arc::new(RwLockServerState(RwLock::new(server))); let data_1 = Arc::clone(&data_0); @@ -126,28 +149,32 @@ async fn main() -> Result<(), Box> { let system = &system; let client_secret = client_secret.clone(); let running_2 = Arc::clone(&running_1); - System::new().block_on(async move { - if running_2.load(Ordering::SeqCst) { - send_stop(&client_secret).await; - } else { - warn!("Got second Ctrl-C, forcefully exiting"); - system.stop() - } - }); + if !OFFLINE_MODE.load(Ordering::Acquire) { + System::new().block_on(async move { + if running_2.load(Ordering::SeqCst) { + send_stop(&client_secret).await; + } else { + warn!("Got second Ctrl-C, forcefully exiting"); + system.stop() + } + }); + } running_1.store(false, Ordering::SeqCst); }) .expect("Error setting Ctrl-C handler"); // Spawn ping task - spawn(async move { - let mut interval = time::interval(Duration::from_secs(90)); - let mut data = Arc::clone(&data_0); - loop { - interval.tick().await; - debug!("Sending ping!"); - ping::update_server_state(&client_secret_1, &cli_args, &mut data).await; - } - }); + if !OFFLINE_MODE.load(Ordering::Acquire) { + spawn(async move { + let mut interval = time::interval(Duration::from_secs(90)); + let mut data = Arc::clone(&data_0); + loop { + interval.tick().await; + debug!("Sending ping!"); + ping::update_server_state(&client_secret_1, &cli_args, &mut data).await; + } + }); + } let cache = DiskCache::new(disk_quota, cache_path.clone()).await; let cache: Arc = if low_mem_mode { @@ -161,19 +188,25 @@ async fn main() -> Result<(), Box> { let cache_0 = Arc::clone(&cache); // Start HTTPS server - HttpServer::new(move || { + let server = HttpServer::new(move || { App::new() .service(routes::token_data) - .service(routes::metrics) .service(routes::token_data_saver) + .service(routes::metrics) .route("{tail:.*}", web::get().to(routes::default)) .app_data(Data::from(Arc::clone(&data_1))) .app_data(Data::from(Arc::clone(&cache_0))) }) - .shutdown_timeout(60) - .bind_rustls(format!("0.0.0.0:{}", port), tls_config)? - .run() - .await?; + .shutdown_timeout(60); + + if disable_tls { + server.bind(format!("0.0.0.0:{}", port))?.run().await?; + } else { + server + .bind_rustls(format!("0.0.0.0:{}", port), tls_config)? + .run() + .await?; + } // Waiting for us to finish sending stop message while running.load(Ordering::SeqCst) { @@ -232,6 +265,17 @@ fn print_preamble_and_warnings(args: &CliArgs) -> Result<(), Box> { warn!("Unstable options are enabled. These options should not be used in production!"); } + if args + .unstable_options + .contains(&UnstableOptions::OfflineMode) + { + warn!("Running in offline mode. No communication to MangaDex will be made!"); + } + + if args.unstable_options.contains(&UnstableOptions::DisableTls) { + warn!("Serving insecure traffic! You better be running this for development only."); + } + if args.override_upstream.is_some() && !args .unstable_options diff --git a/src/metrics.rs b/src/metrics.rs index e07d2ee..f0051f8 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -2,18 +2,18 @@ use once_cell::sync::Lazy; use prometheus::{register_int_counter, IntCounter}; pub static CACHE_HIT_COUNTER: Lazy = - Lazy::new(|| register_int_counter!("cache.hit", "The number of cache hits").unwrap()); + Lazy::new(|| register_int_counter!("cache_hit", "The number of cache hits").unwrap()); pub static CACHE_MISS_COUNTER: Lazy = - Lazy::new(|| register_int_counter!("cache.miss", "The number of cache misses").unwrap()); + Lazy::new(|| register_int_counter!("cache_miss", "The number of cache misses").unwrap()); pub static REQUESTS_TOTAL_COUNTER: Lazy = Lazy::new(|| { - register_int_counter!("requests.total", "The total number of requests served.").unwrap() + register_int_counter!("requests_total", "The total number of requests served.").unwrap() }); pub static REQUESTS_DATA_COUNTER: Lazy = Lazy::new(|| { register_int_counter!( - "requests.data", + "requests_data", "The number of requests served from the /data endpoint." ) .unwrap() @@ -21,7 +21,7 @@ pub static REQUESTS_DATA_COUNTER: Lazy = Lazy::new(|| { pub static REQUESTS_DATA_SAVER_COUNTER: Lazy = Lazy::new(|| { register_int_counter!( - "requests.data-saver", + "requests_data_saver", "The number of requests served from the /data-saver endpoint." ) .unwrap() @@ -29,8 +29,17 @@ pub static REQUESTS_DATA_SAVER_COUNTER: Lazy = Lazy::new(|| { pub static REQUESTS_OTHER_COUNTER: Lazy = Lazy::new(|| { register_int_counter!( - "requests.other", + "requests_other", "The total number of request not served by primary endpoints." ) .unwrap() }); + +pub fn init() { + let _a = CACHE_HIT_COUNTER.get(); + let _a = CACHE_MISS_COUNTER.get(); + let _a = REQUESTS_TOTAL_COUNTER.get(); + let _a = REQUESTS_DATA_COUNTER.get(); + let _a = REQUESTS_DATA_SAVER_COUNTER.get(); + let _a = REQUESTS_OTHER_COUNTER.get(); +} diff --git a/src/routes.rs b/src/routes.rs index 46926a2..f8bf219 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -1,5 +1,6 @@ use std::sync::atomic::Ordering; +use actix_web::error::ErrorNotFound; use actix_web::http::header::{ ACCESS_CONTROL_ALLOW_ORIGIN, ACCESS_CONTROL_EXPOSE_HEADERS, CACHE_CONTROL, CONTENT_LENGTH, CONTENT_TYPE, LAST_MODIFIED, X_CONTENT_TYPE_OPTIONS, @@ -21,7 +22,7 @@ use thiserror::Error; use crate::cache::{Cache, CacheKey, ImageMetadata, UpstreamError}; use crate::client_api_version; -use crate::config::{SEND_SERVER_VERSION, VALIDATE_TOKENS}; +use crate::config::{OFFLINE_MODE, SEND_SERVER_VERSION, VALIDATE_TOKENS}; use crate::metrics::{ CACHE_HIT_COUNTER, CACHE_MISS_COUNTER, REQUESTS_DATA_COUNTER, REQUESTS_DATA_SAVER_COUNTER, REQUESTS_OTHER_COUNTER, REQUESTS_TOTAL_COUNTER, @@ -102,7 +103,16 @@ pub async fn default(state: Data, req: HttpRequest) -> impl R state.0.read().image_server, req.path().chars().skip(1).collect::() ); - info!("Got unknown path, just proxying: {}", path); + + if OFFLINE_MODE.load(Ordering::Acquire) { + info!("Got unknown path in offline mode, returning 404: {}", path); + return ServerResponse::HttpResponse( + ErrorNotFound("Path is not valid in offline mode").into(), + ); + } else { + info!("Got unknown path, just proxying: {}", path); + } + let resp = match HTTP_CLIENT.get(path).send().await { Ok(resp) => resp, Err(e) => { @@ -231,6 +241,13 @@ async fn fetch_image( CACHE_MISS_COUNTER.inc(); + // If in offline mode, return early since there's nothing else we can do + if OFFLINE_MODE.load(Ordering::Acquire) { + return ServerResponse::HttpResponse( + ErrorNotFound("Offline mode enabled and image not in cache").into(), + ); + } + // It's important to not get a write lock before this request, else we're // holding the read lock until the await resolves. diff --git a/src/state.rs b/src/state.rs index f113d3a..796bda3 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,6 +1,7 @@ +use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use crate::config::{CliArgs, UnstableOptions, SEND_SERVER_VERSION, VALIDATE_TOKENS}; +use crate::config::{CliArgs, UnstableOptions, OFFLINE_MODE, SEND_SERVER_VERSION, VALIDATE_TOKENS}; use crate::ping::{Request, Response, CONTROL_CENTER_PING_URL}; use arc_swap::ArcSwap; use log::{error, info, warn}; @@ -9,7 +10,7 @@ use parking_lot::RwLock; use rustls::sign::{CertifiedKey, SigningKey}; use rustls::Certificate; use rustls::{ClientHello, ResolvesServerCert}; -use sodiumoxide::crypto::box_::PrecomputedKey; +use sodiumoxide::crypto::box_::{PrecomputedKey, PRECOMPUTEDKEYBYTES}; use thiserror::Error; use url::Url; @@ -144,6 +145,16 @@ impl ServerState { }, } } + + pub fn init_offline() -> Self { + assert!(OFFLINE_MODE.load(Ordering::Acquire)); + Self { + precomputed_key: PrecomputedKey::from_slice(&[41; PRECOMPUTEDKEYBYTES]).unwrap(), + image_server: Url::from_file_path("/dev/null").unwrap(), + url: Url::from_str("http://localhost").unwrap(), + url_overridden: false, + } + } } pub struct RwLockServerState(pub RwLock);