From 50dce3a80ba6052622e906ed83dd5816689be82f Mon Sep 17 00:00:00 2001
From: Edward Shen <code@eddie.sh>
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<PathBuf>,
+  /// 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<Config, BunBunError> {
+pub fn read_config(
+  mut config_file: File,
+  large_config: bool,
+) -> Result<Config, BunBunError> {
   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<Handlebars, TemplateError> {
 fn start_watch(
   state: Arc<RwLock<State>>,
   config_data: ConfigData,
+  large_config: bool,
 ) -> Result<Hotwatch, BunBunError> {
   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;