Support CLI args

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"

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"

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,
+}

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 = cli_args.port;
-    let port = config.port;
+
-    let server = ServerState::init(&config).await.unwrap();
+    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<Self, ServerError> {
-        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,
-        })
-    }
-}

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<String>,
 }
 
 impl<'a> Request<'a> {
-    fn from_config_and_state(config: &'a Config, state: &'a Arc<RwLockServerState>) -> Self {
+    fn from_config_and_state(
+        secret: &'a str,
+        config: &CliArgs,
+        state: &Arc<RwLockServerState>,
+    ) -> 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> {
+impl<'a> From<(&'a str, &CliArgs)> for Request<'a> {
-    fn from(config: &'a Config) -> Self {
+    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<RwLockServerState>) {
+pub async fn update_server_state(secret: &str, req: &CliArgs, data: &mut Arc<RwLockServerState>) {
-    let req = Request::from_config_and_state(req, data);
+    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 {

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(("Timing-Allow-Origin", "https://mangadex.org"));
-        .insert_header(("Server", SERVER_ID_STRING))
+
+    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(

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<Self, ()> {
+    pub async fn init(secret: &str, config: &CliArgs) -> Result<Self, ()> {
         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::<Response>().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,