From f775ad72d30a66620631df8c00a8a7e12af2651d Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Thu, 25 Mar 2021 21:06:54 -0400 Subject: [PATCH] Support CLI args --- Cargo.lock | 146 ++++++++++++++++++++++++++++++++++++++++++++++---- Cargo.toml | 1 + src/config.rs | 30 +++++++++++ src/main.rs | 60 ++++++++------------- src/ping.rs | 29 ++++++---- src/routes.rs | 18 +++++-- src/state.rs | 19 ++++--- 7 files changed, 233 insertions(+), 70 deletions(-) create mode 100644 src/config.rs diff --git a/Cargo.lock b/Cargo.lock index afebbf3..56eca57 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -606,6 +606,39 @@ dependencies = [ "winapi", ] +[[package]] +name = "clap" +version = "3.0.0-beta.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bd1061998a501ee7d4b6d449020df3266ca3124b941ec56cf2005c3779ca142" +dependencies = [ + "atty", + "bitflags", + "clap_derive", + "indexmap", + "lazy_static", + "os_str_bytes", + "strsim", + "termcolor", + "terminal_size", + "textwrap", + "unicode-width", + "vec_map", +] + +[[package]] +name = "clap_derive" +version = "3.0.0-beta.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "370f715b81112975b1b69db93e0b56ea4cd4e5002ac43b2da8474106a54096a1" +dependencies = [ + "heck", + "proc-macro-error", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "colored" version = "1.9.3" @@ -693,9 +726,9 @@ dependencies = [ [[package]] name = "ctor" -version = "0.1.19" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8f45d9ad417bcef4817d614a501ab55cdd96a6fdb24f49aab89a54acfd66b19" +checksum = "5e98e2ad1a782e33928b96fc3948e7c355e5af34ba4de7670fe8bac2a3b2006d" dependencies = [ "quote", "syn", @@ -984,9 +1017,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d832b01df74254fe364568d6ddc294443f61cbec82816b60904303af87efae78" +checksum = "fc018e188373e2777d0ef2467ebff62a08e66c3f5857b23c8fbec3018210dc00" dependencies = [ "bytes", "fnv", @@ -1010,6 +1043,15 @@ dependencies = [ "ahash 0.4.7", ] +[[package]] +name = "heck" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87cbf45460356b7deeb5e3415b5563308c0a9b057c85e12b06ad551f98d0a6ac" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "hermit-abi" version = "0.1.18" @@ -1176,9 +1218,9 @@ checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" [[package]] name = "libc" -version = "0.2.90" +version = "0.2.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba4aede83fc3617411dc6993bc8c70919750c1c257c6ca6a502aed6e0e2394ae" +checksum = "8916b1f6ca17130ec6568feccee27c156ad12037880833a3b842a823236502e7" [[package]] name = "libsodium-sys" @@ -1229,6 +1271,7 @@ dependencies = [ "bytes", "cacache", "chrono", + "clap", "ctrlc", "dotenv", "futures", @@ -1435,6 +1478,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "os_str_bytes" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "afb2e1c3ee07430c2cf76151675e583e0f19985fa6efae47d6848a3e2c824f85" + [[package]] name = "parking" version = "2.0.0" @@ -1474,18 +1523,18 @@ checksum = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" [[package]] name = "pin-project" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96fa8ebb90271c4477f144354485b8068bd8f6b78b428b01ba892ca26caf0b63" +checksum = "bc174859768806e91ae575187ada95c91a29e96a98dc5d2cd9a1fed039501ba6" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "758669ae3558c6f74bd2a18b41f7ac0b5a195aea6639d6a9b5e5d1ad5ba24c0b" +checksum = "a490329918e856ed1b083f244e3bfe2d8c4f336407e4ea9e1a9f479ff09049e5" dependencies = [ "proc-macro2", "quote", @@ -1529,6 +1578,30 @@ version = "0.2.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac74c624d6b2d21f425f752262f42188365d7b8ff1aff74c82e45136510a4857" +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "syn", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + [[package]] name = "proc-macro-hack" version = "0.5.19" @@ -2027,6 +2100,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "213701ba3370744dcd1a12960caa4843b3d68b4d1c0a5d575e0d65b2ee9d16c0" +[[package]] +name = "strsim" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" + [[package]] name = "syn" version = "1.0.64" @@ -2052,6 +2131,35 @@ dependencies = [ "winapi", ] +[[package]] +name = "termcolor" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dfed899f0eb03f32ee8c6a0aabdb8a7949659e3466561fc0adf54e26d88c5f4" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "terminal_size" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86ca8ced750734db02076f44132d802af0b33b09942331f4459dde8636fd2406" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "textwrap" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "203008d98caf094106cfaba70acfed15e18ed3ddb7d94e49baec153a2b462789" +dependencies = [ + "terminal_size", + "unicode-width", +] + [[package]] name = "thiserror" version = "1.0.24" @@ -2244,6 +2352,18 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-segmentation" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0d2e7be6ae3a5fa87eed5fb451aff96f2573d2694942e40543ae0bbe19c796" + +[[package]] +name = "unicode-width" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9337591893a19b88d8d87f2cec1e73fad5cdfd10e5a6f349f498ad6ea2ffb1e3" + [[package]] name = "unicode-xid" version = "0.2.1" @@ -2290,6 +2410,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34b2f665b594b07095e3ac3f718e13c2197143416fae4c5706cffb7b1af8d7f1" +[[package]] +name = "vec_map" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1bddf1187be692e79c5ffeab891132dfb0f236ed36a43c7ed39f1165ee20191" + [[package]] name = "version_check" version = "0.9.3" diff --git a/Cargo.toml b/Cargo.toml index 97fcbe5..88b9b35 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ bincode = "1" bytes = "1" cacache = "8" chrono = { version = "0.4", features = [ "serde" ] } +clap = { version = "3.0.0-beta.2", features = [ "wrap_help" ] } ctrlc = "3" dotenv = "0.15" futures = "0.3" diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..1252823 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,30 @@ +use std::num::{NonZeroU16, NonZeroUsize}; +use std::path::PathBuf; + +use clap::Clap; + +#[derive(Clap)] +pub struct CliArgs { + /// The port to listen on. + #[clap(short, long, default_value = "42069", env = "PORT")] + pub port: NonZeroU16, + /// How large, in bytes, the in-memory cache should be. Note that this does + /// not include runtime memory usage. + #[clap(long, env = "MEM_CACHE_QUOTA_BYTES")] + pub memory_quota: NonZeroUsize, + /// How large, in bytes, the on-disk cache should be. Note that actual + /// values may be larger for metadata information. + #[clap(long, env = "DISK_CACHE_QUOTA_BYTES")] + pub disk_quota: usize, + /// Sets the location of the disk cache. + #[clap(long, default_value = "./cache", env = "DISK_CACHE_PATH")] + pub cache_path: PathBuf, + /// The network speed to advertise to Mangadex@Home control server. + #[clap(long, env = "MAX_NETWORK_SPEED")] + pub network_speed: NonZeroUsize, + /// Whether or not to provide the Server HTTP header to clients. This is + /// useful for debugging, but is generally not recommended for security + /// reasons. + #[clap(long, env = "ENABLE_SERVER_STRING")] + pub enable_server_string: bool, +} diff --git a/src/main.rs b/src/main.rs index 1dcb1e2..5708a2e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,16 +2,20 @@ // We're end users, so these is ok #![allow(clippy::future_not_send, clippy::module_name_repetitions)] -use std::env::{self, VarError}; -use std::path::PathBuf; use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::time::Duration; +use std::{ + env::{self, VarError}, + process, +}; use std::{num::ParseIntError, sync::atomic::Ordering}; use actix_web::rt::{spawn, time, System}; use actix_web::web::{self, Data}; use actix_web::{App, HttpServer}; +use clap::Clap; +use config::CliArgs; use log::{debug, error, warn, LevelFilter}; use parking_lot::RwLock; use rustls::{NoClientAuth, ServerConfig}; @@ -21,6 +25,7 @@ use stop::send_stop; use thiserror::Error; mod cache; +mod config; mod ping; mod routes; mod state; @@ -44,18 +49,27 @@ enum ServerError { async fn main() -> Result<(), std::io::Error> { // It's ok to fail early here, it would imply we have a invalid config. dotenv::dotenv().ok(); + let cli_args = CliArgs::parse(); + SimpleLogger::new() .with_level(LevelFilter::Info) .init() .unwrap(); - let config = Config::new().unwrap(); - let port = config.port; - let server = ServerState::init(&config).await.unwrap(); + let port = cli_args.port; + + let client_secret = if let Ok(v) = env::var("CLIENT_SECRET") { + v + } else { + eprintln!("Client secret not found in ENV. Please set CLIENT_SECRET."); + process::exit(1); + }; + let client_secret_1 = client_secret.clone(); + + let server = ServerState::init(&client_secret, &cli_args).await.unwrap(); // Set ctrl+c to send a stop message let running = Arc::new(AtomicBool::new(true)); let r = running.clone(); - let client_secret = config.secret.clone(); ctrlc::set_handler(move || { let client_secret = client_secret.clone(); System::new().block_on(async move { @@ -75,7 +89,7 @@ async fn main() -> Result<(), std::io::Error> { loop { interval.tick().await; debug!("Sending ping!"); - ping::update_server_state(&config, &mut data).await; + ping::update_server_state(&client_secret_1, &cli_args, &mut data).await; } }); @@ -101,35 +115,3 @@ async fn main() -> Result<(), std::io::Error> { Ok(()) } - -pub struct Config { - secret: String, - port: u16, - memory_quota: usize, - disk_quota: usize, - disk_path: PathBuf, - network_speed: usize, -} - -impl Config { - fn new() -> Result { - let secret = env::var("CLIENT_SECRET")?; - let port = env::var("PORT")?.parse()?; - let disk_quota = env::var("DISK_CACHE_QUOTA_BYTES")?.parse()?; - let memory_quota = env::var("MEM_CACHE_QUOTA_BYTES")?.parse()?; - let network_speed = env::var("MAX_NETWORK_SPEED")?.parse()?; - let disk_path = env::var("DISK_CACHE_PATH") - .unwrap_or_else(|_| "./cache".to_string()) - .parse() - .unwrap(); - - Ok(Self { - secret, - port, - disk_quota, - memory_quota, - disk_path, - network_speed, - }) - } -} diff --git a/src/ping.rs b/src/ping.rs index a58a924..88fb3ed 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -1,4 +1,7 @@ -use std::sync::Arc; +use std::{ + num::{NonZeroU16, NonZeroUsize}, + sync::Arc, +}; use log::{error, info, warn}; use serde::{Deserialize, Serialize}; @@ -6,24 +9,28 @@ use sodiumoxide::crypto::box_::PrecomputedKey; use url::Url; use crate::state::RwLockServerState; -use crate::{client_api_version, Config}; +use crate::{client_api_version, config::CliArgs}; pub const CONTROL_CENTER_PING_URL: &str = "https://api.mangadex.network/ping"; #[derive(Serialize, Debug)] pub struct Request<'a> { secret: &'a str, - port: u16, + port: NonZeroU16, disk_space: usize, - network_speed: usize, + network_speed: NonZeroUsize, build_version: usize, tls_created_at: Option, } impl<'a> Request<'a> { - fn from_config_and_state(config: &'a Config, state: &'a Arc) -> Self { + fn from_config_and_state( + secret: &'a str, + config: &CliArgs, + state: &Arc, + ) -> Self { Self { - secret: &config.secret, + secret, port: config.port, disk_space: config.disk_quota, network_speed: config.network_speed, @@ -34,10 +41,10 @@ impl<'a> Request<'a> { } #[allow(clippy::fallible_impl_from)] -impl<'a> From<&'a Config> for Request<'a> { - fn from(config: &'a Config) -> Self { +impl<'a> From<(&'a str, &CliArgs)> for Request<'a> { + fn from((secret, config): (&'a str, &CliArgs)) -> Self { Self { - secret: &config.secret, + secret, port: config.port, disk_space: config.disk_quota, network_speed: config.network_speed, @@ -67,8 +74,8 @@ pub struct Tls { pub certificate: String, } -pub async fn update_server_state(req: &Config, data: &mut Arc) { - let req = Request::from_config_and_state(req, data); +pub async fn update_server_state(secret: &str, req: &CliArgs, data: &mut Arc) { + let req = Request::from_config_and_state(secret, req, data); let client = reqwest::Client::new(); let resp = client.post(CONTROL_CENTER_PING_URL).json(&req).send().await; match resp { diff --git a/src/routes.rs b/src/routes.rs index 7a2c477..ff153f7 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -1,4 +1,7 @@ -use std::convert::Infallible; +use std::{ + convert::Infallible, + sync::atomic::{AtomicBool, Ordering}, +}; use actix_web::dev::HttpResponseBuilder; use actix_web::http::header::{ @@ -22,6 +25,8 @@ use crate::state::RwLockServerState; pub const BASE64_CONFIG: base64::Config = base64::Config::new(base64::CharacterSet::UrlSafe, false); +pub static SEND_SERVER_VERSION: AtomicBool = AtomicBool::new(false); + const SERVER_ID_STRING: &str = concat!( env!("CARGO_CRATE_NAME"), " ", @@ -159,8 +164,13 @@ fn push_headers(builder: &mut HttpResponseBuilder) -> &mut HttpResponseBuilder { .insert_header((ACCESS_CONTROL_ALLOW_ORIGIN, "https://mangadex.org")) .insert_header((ACCESS_CONTROL_EXPOSE_HEADERS, "*")) .insert_header((CACHE_CONTROL, "public, max-age=1209600")) - .insert_header(("Timing-Allow-Origin", "https://mangadex.org")) - .insert_header(("Server", SERVER_ID_STRING)) + .insert_header(("Timing-Allow-Origin", "https://mangadex.org")); + + if SEND_SERVER_VERSION.load(Ordering::Acquire) { + builder.insert_header(("Server", SERVER_ID_STRING)); + } + + builder } async fn fetch_image( @@ -176,6 +186,7 @@ async fn fetch_image( } let mut state = state.0.write(); + let resp = if is_data_saver { reqwest::get(format!( "{}/data-saver/{}/{}", @@ -202,6 +213,7 @@ async fn fetch_image( if let Some(content_type) = content_type { resp_builder.insert_header((CONTENT_TYPE, content_type)); } + push_headers(&mut resp_builder); return ServerResponse::HttpResponse( diff --git a/src/state.rs b/src/state.rs index 79e616d..9bd2013 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,8 +1,11 @@ -use std::{io::BufReader, sync::Arc}; +use std::{ + io::BufReader, + sync::{atomic::Ordering, Arc}, +}; -use crate::cache::Cache; use crate::ping::{Request, Response, Tls, CONTROL_CENTER_PING_URL}; -use crate::Config; +use crate::routes::SEND_SERVER_VERSION; +use crate::{cache::Cache, config::CliArgs}; use log::{error, info, warn}; use parking_lot::RwLock; use rustls::internal::pemfile::{certs, rsa_private_keys}; @@ -26,13 +29,15 @@ pub struct LogState { } impl ServerState { - pub async fn init(config: &Config) -> Result { + pub async fn init(secret: &str, config: &CliArgs) -> Result { let resp = reqwest::Client::new() .post(CONTROL_CENTER_PING_URL) - .json(&Request::from(config)) + .json(&Request::from((secret, config))) .send() .await; + SEND_SERVER_VERSION.store(config.enable_server_string, Ordering::Release); + match resp { Ok(resp) => match resp.json::().await { Ok(resp) => { @@ -74,9 +79,9 @@ impl ServerState { force_tokens: resp.force_tokens, url: resp.url, cache: Cache::new( - config.memory_quota, + config.memory_quota.get(), config.disk_quota, - config.disk_path.clone(), + config.cache_path.clone(), ), log_state: LogState { was_paused_before: resp.paused,