From 5300afa20563dd9cb891a601711d390fc5f90e8e Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Fri, 9 Jul 2021 19:48:25 -0400 Subject: [PATCH] Clean up ClientSecret usage --- src/config.rs | 30 +++++++++++++++++++++--------- src/main.rs | 21 +++++++++++---------- src/ping.rs | 18 +++++++++++------- src/state.rs | 4 ++-- src/stop.rs | 6 ++++-- 5 files changed, 49 insertions(+), 30 deletions(-) diff --git a/src/config.rs b/src/config.rs index bdbe522..1d80230 100644 --- a/src/config.rs +++ b/src/config.rs @@ -11,6 +11,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use clap::{crate_authors, crate_description, crate_version, Clap}; use log::LevelFilter; use serde::{Deserialize, Serialize}; +use thiserror::Error; use url::Url; use crate::units::{KilobitsPerSecond, Mebibytes, Port}; @@ -19,7 +20,17 @@ use crate::units::{KilobitsPerSecond, Mebibytes, Port}; pub static VALIDATE_TOKENS: AtomicBool = AtomicBool::new(false); pub static OFFLINE_MODE: AtomicBool = AtomicBool::new(false); -pub fn load_config() -> Result { +#[derive(Error, Debug)] +pub enum ConfigError { + #[error("No config found. One has been created for you to modify.")] + NotInitialized, + #[error(transparent)] + Io(#[from] std::io::Error), + #[error(transparent)] + Parse(#[from] serde_yaml::Error), +} + +pub fn load_config() -> Result { // Load cli args first let cli_args: CliArgs = CliArgs::parse(); @@ -40,13 +51,10 @@ pub fn load_config() -> Result { let default_config = include_str!("../settings.sample.yaml"); file.write_all(default_config.as_bytes()).unwrap(); - serde_yaml::from_str(default_config) + + return Err(ConfigError::NotInitialized); } - e => panic!( - "Failed to open file at {}: {:?}", - config_path.to_string_lossy(), - e - ), + Err(e) => return Err(e.into()), } }; @@ -129,7 +137,11 @@ impl Config { .unwrap_or(unsafe { NonZeroU16::new_unchecked(60) }), log_level, // secret should never be in CLI - client_secret: file_args.server_settings.secret, + client_secret: if let Ok(v) = std::env::var("CLIENT_SECRET") { + ClientSecret(v) + } else { + file_args.server_settings.secret + }, port: cli_args.port.unwrap_or(file_args.server_settings.port), bind_address: SocketAddr::new( file_args @@ -188,7 +200,7 @@ struct YamlServerSettings { } // this intentionally does not implement display or debug -#[derive(Deserialize, Serialize)] +#[derive(Deserialize, Serialize, Clone)] pub struct ClientSecret(String); #[derive(Deserialize, Default)] diff --git a/src/main.rs b/src/main.rs index 0ca0bbc..683ac80 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,11 +2,10 @@ // We're end users, so these is ok #![allow(clippy::module_name_repetitions)] -use std::env::{self, VarError}; +use std::env::VarError; use std::error::Error; use std::fmt::Display; use std::num::ParseIntError; -use std::process; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -59,7 +58,14 @@ async fn main() -> Result<(), Box> { // Config loading // - let config = config::load_config()?; + let config = match config::load_config() { + Ok(c) => c, + Err(e) => { + eprintln!("{}", e); + return Err(Box::new(e) as Box<_>); + } + }; + let memory_quota = config.memory_quota; let disk_quota = config.disk_quota; let cache_type = config.cache_type; @@ -80,13 +86,8 @@ async fn main() -> Result<(), Box> { return Err(e); } - let client_secret = if let Ok(v) = env::var("CLIENT_SECRET") { - v - } else { - error!("Client secret not found in ENV. Please set CLIENT_SECRET."); - process::exit(1); - }; - let client_secret_1 = client_secret.clone(); + let client_secret = config.client_secret.clone(); + let client_secret_1 = config.client_secret.clone(); if config.ephemeral_disk_encryption { info!("Running with at-rest encryption!"); diff --git a/src/ping.rs b/src/ping.rs index 2ca2401..96c4ae0 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -11,7 +11,7 @@ use serde_repr::Deserialize_repr; use sodiumoxide::crypto::box_::PrecomputedKey; use url::Url; -use crate::config::{Config, UnstableOptions, VALIDATE_TOKENS}; +use crate::config::{ClientSecret, Config, UnstableOptions, VALIDATE_TOKENS}; use crate::state::{ RwLockServerState, PREVIOUSLY_COMPROMISED, PREVIOUSLY_PAUSED, TLS_CERTS, TLS_PREVIOUSLY_CREATED, TLS_SIGNING_KEY, @@ -21,9 +21,9 @@ use crate::CLIENT_API_VERSION; pub const CONTROL_CENTER_PING_URL: &str = "https://api.mangadex.network/ping"; -#[derive(Serialize, Debug)] +#[derive(Serialize)] pub struct Request<'a> { - secret: &'a str, + secret: &'a ClientSecret, port: Port, disk_space: Mebibytes, network_speed: BytesPerSecond, @@ -32,7 +32,7 @@ pub struct Request<'a> { } impl<'a> Request<'a> { - fn from_config_and_state(secret: &'a str, config: &Config) -> Self { + fn from_config_and_state(secret: &'a ClientSecret, config: &Config) -> Self { Self { secret, port: config.port, @@ -46,8 +46,8 @@ impl<'a> Request<'a> { } } -impl<'a> From<(&'a str, &Config)> for Request<'a> { - fn from((secret, config): (&'a str, &Config)) -> Self { +impl<'a> From<(&'a ClientSecret, &Config)> for Request<'a> { + fn from((secret, config): (&'a ClientSecret, &Config)) -> Self { Self { secret, port: config.port, @@ -161,7 +161,11 @@ impl std::fmt::Debug for Tls { } } -pub async fn update_server_state(secret: &str, cli: &Config, data: &mut Arc) { +pub async fn update_server_state( + secret: &ClientSecret, + cli: &Config, + data: &mut Arc, +) { let req = Request::from_config_and_state(secret, cli); let client = reqwest::Client::new(); let resp = client.post(CONTROL_CENTER_PING_URL).json(&req).send().await; diff --git a/src/state.rs b/src/state.rs index a175db4..8e3b542 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,7 +1,7 @@ use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use crate::config::{Config, UnstableOptions, OFFLINE_MODE, VALIDATE_TOKENS}; +use crate::config::{ClientSecret, Config, UnstableOptions, OFFLINE_MODE, VALIDATE_TOKENS}; use crate::ping::{Request, Response, CONTROL_CENTER_PING_URL}; use arc_swap::ArcSwap; use log::{error, info, warn}; @@ -45,7 +45,7 @@ pub enum ServerInitError { } impl ServerState { - pub async fn init(secret: &str, config: &Config) -> Result { + pub async fn init(secret: &ClientSecret, config: &Config) -> Result { let resp = reqwest::Client::new() .post(CONTROL_CENTER_PING_URL) .json(&Request::from((secret, config))) diff --git a/src/stop.rs b/src/stop.rs index 723f02e..5164ffa 100644 --- a/src/stop.rs +++ b/src/stop.rs @@ -2,14 +2,16 @@ use log::{info, warn}; use reqwest::StatusCode; use serde::Serialize; +use crate::config::ClientSecret; + const CONTROL_CENTER_STOP_URL: &str = "https://api.mangadex.network/ping"; #[derive(Serialize)] struct StopRequest<'a> { - secret: &'a str, + secret: &'a ClientSecret, } -pub async fn send_stop(secret: &str) { +pub async fn send_stop(secret: &ClientSecret) { let request = StopRequest { secret }; let client = reqwest::Client::new(); match client