From 50dce3a80ba6052622e906ed83dd5816689be82f Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Sun, 5 Jul 2020 12:17:29 -0400 Subject: [PATCH] Implement config size checks --- src/cli.rs | 3 +++ src/config.rs | 16 +++++++++++++++- src/error.rs | 4 ++++ src/main.rs | 11 +++++++---- 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index ad21561..d559fab 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -13,4 +13,7 @@ pub struct Opts { /// Specify the location of the config file to read from. Needs read/write permissions. #[clap(short, long)] pub config: Option, + /// Allow config sizes larger than 100MB. + #[clap(long)] + pub large_config: bool, } diff --git a/src/config.rs b/src/config.rs index f490a8b..a4be7fa 100644 --- a/src/config.rs +++ b/src/config.rs @@ -263,8 +263,22 @@ pub fn load_custom_path_config( Ok(ConfigData { file, path }) } -pub fn read_config(mut config_file: File) -> Result { +pub fn read_config( + mut config_file: File, + large_config: bool, +) -> Result { trace!("Loading config file."); + let file_size = config_file.metadata()?.len(); + + // 100 MB + if file_size > 100_000_000 && !large_config { + return Err(BunBunError::ConfigTooLarge(file_size)); + } + + if file_size == 0 { + return Err(BunBunError::ZeroByteConfig); + } + let mut config_data = String::new(); config_file.read_to_string(&mut config_data)?; // Reading from memory is faster than reading directly from a reader for some diff --git a/src/error.rs b/src/error.rs index 50a8fd3..3c2dba3 100644 --- a/src/error.rs +++ b/src/error.rs @@ -10,6 +10,8 @@ pub enum BunBunError { CustomProgram(String), NoValidConfigPath, InvalidConfigPath(std::path::PathBuf, std::io::Error), + ConfigTooLarge(u64), + ZeroByteConfig, } impl Error for BunBunError {} @@ -26,6 +28,8 @@ impl fmt::Display for BunBunError { Self::InvalidConfigPath(path, reason) => { write!(f, "Failed to access {:?}: {}", path, reason) } + Self::ConfigTooLarge(size) => write!(f, "The config file was too large ({} bytes)! Pass in --large-config to bypass this check.", size), + Self::ZeroByteConfig => write!(f, "The config provided reported a size of 0 bytes. Please check your config path!") } } } diff --git a/src/main.rs b/src/main.rs index 8f818d8..14eebbc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -57,7 +57,7 @@ async fn run() -> Result<(), BunBunError> { None => get_config_data(), }?; - let conf = read_config(conf_data.file.try_clone()?)?; + let conf = read_config(conf_data.file.try_clone()?, opts.large_config)?; let state = Arc::from(RwLock::new(State { public_address: conf.public_address, default_route: conf.default_route, @@ -65,7 +65,7 @@ async fn run() -> Result<(), BunBunError> { groups: conf.groups, })); - let _watch = start_watch(Arc::clone(&state), conf_data)?; + let _watch = start_watch(Arc::clone(&state), conf_data, opts.large_config)?; HttpServer::new(move || { let templates = match compile_templates() { @@ -170,6 +170,7 @@ fn compile_templates() -> Result { fn start_watch( state: Arc>, config_data: ConfigData, + large_config: bool, ) -> Result { let mut watch = Hotwatch::new_with_custom_delay(Duration::from_millis(500))?; let ConfigData { path, file } = config_data; @@ -178,8 +179,10 @@ fn start_watch( trace!("Grabbing writer lock on state..."); let mut state = state.write().expect("Failed to get write lock on state"); trace!("Obtained writer lock on state!"); - match read_config(file.try_clone().expect("Failed to clone file handle")) - { + match read_config( + file.try_clone().expect("Failed to clone file handle"), + large_config, + ) { Ok(conf) => { state.public_address = conf.public_address; state.default_route = conf.default_route;