Implement config size checks

This commit is contained in:
Edward Shen 2020-07-05 12:17:29 -04:00
parent 7995babb64
commit 50dce3a80b
Signed by: edward
GPG key ID: 19182661E818369F
4 changed files with 29 additions and 5 deletions

View file

@ -13,4 +13,7 @@ pub struct Opts {
/// Specify the location of the config file to read from. Needs read/write permissions. /// Specify the location of the config file to read from. Needs read/write permissions.
#[clap(short, long)] #[clap(short, long)]
pub config: Option<PathBuf>, pub config: Option<PathBuf>,
/// Allow config sizes larger than 100MB.
#[clap(long)]
pub large_config: bool,
} }

View file

@ -263,8 +263,22 @@ pub fn load_custom_path_config(
Ok(ConfigData { file, path }) Ok(ConfigData { file, path })
} }
pub fn read_config(mut config_file: File) -> Result<Config, BunBunError> { pub fn read_config(
mut config_file: File,
large_config: bool,
) -> Result<Config, BunBunError> {
trace!("Loading config file."); 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(); let mut config_data = String::new();
config_file.read_to_string(&mut config_data)?; config_file.read_to_string(&mut config_data)?;
// Reading from memory is faster than reading directly from a reader for some // Reading from memory is faster than reading directly from a reader for some

View file

@ -10,6 +10,8 @@ pub enum BunBunError {
CustomProgram(String), CustomProgram(String),
NoValidConfigPath, NoValidConfigPath,
InvalidConfigPath(std::path::PathBuf, std::io::Error), InvalidConfigPath(std::path::PathBuf, std::io::Error),
ConfigTooLarge(u64),
ZeroByteConfig,
} }
impl Error for BunBunError {} impl Error for BunBunError {}
@ -26,6 +28,8 @@ impl fmt::Display for BunBunError {
Self::InvalidConfigPath(path, reason) => { Self::InvalidConfigPath(path, reason) => {
write!(f, "Failed to access {:?}: {}", 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!")
} }
} }
} }

View file

@ -57,7 +57,7 @@ async fn run() -> Result<(), BunBunError> {
None => get_config_data(), 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 { let state = Arc::from(RwLock::new(State {
public_address: conf.public_address, public_address: conf.public_address,
default_route: conf.default_route, default_route: conf.default_route,
@ -65,7 +65,7 @@ async fn run() -> Result<(), BunBunError> {
groups: conf.groups, 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 || { HttpServer::new(move || {
let templates = match compile_templates() { let templates = match compile_templates() {
@ -170,6 +170,7 @@ fn compile_templates() -> Result<Handlebars, TemplateError> {
fn start_watch( fn start_watch(
state: Arc<RwLock<State>>, state: Arc<RwLock<State>>,
config_data: ConfigData, config_data: ConfigData,
large_config: bool,
) -> Result<Hotwatch, BunBunError> { ) -> Result<Hotwatch, BunBunError> {
let mut watch = Hotwatch::new_with_custom_delay(Duration::from_millis(500))?; let mut watch = Hotwatch::new_with_custom_delay(Duration::from_millis(500))?;
let ConfigData { path, file } = config_data; let ConfigData { path, file } = config_data;
@ -178,8 +179,10 @@ fn start_watch(
trace!("Grabbing writer lock on state..."); trace!("Grabbing writer lock on state...");
let mut state = state.write().expect("Failed to get write lock on state"); let mut state = state.write().expect("Failed to get write lock on state");
trace!("Obtained writer 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) => { Ok(conf) => {
state.public_address = conf.public_address; state.public_address = conf.public_address;
state.default_route = conf.default_route; state.default_route = conf.default_route;