Support CLI args

feature/v32-tokens
Edward Shen 2021-03-25 21:06:54 -04:00
parent cf82aedb12
commit f775ad72d3
Signed by: edward
GPG Key ID: 19182661E818369F
7 changed files with 233 additions and 70 deletions

146
Cargo.lock generated
View File

@ -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"

View File

@ -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"

30
src/config.rs Normal file
View File

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

View File

@ -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<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,
})
}
}

View File

@ -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> {
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<RwLockServerState>) {
let req = Request::from_config_and_state(req, data);
pub async fn update_server_state(secret: &str, req: &CliArgs, data: &mut Arc<RwLockServerState>) {
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 {

View File

@ -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(

View File

@ -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,