Limit the size of served files

benches/compression.rs

@@ -1,17 +1,19 @@
-use cratesfyi::storage::{compress, decompress};
+use cratesfyi::storage::{compress, decompress, CompressionAlgorithm};
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 
+const ALGORITHM: CompressionAlgorithm = CompressionAlgorithm::Zstd;
+
 pub fn criterion_benchmark(c: &mut Criterion) {
     // this isn't a great benchmark because it only tests on one file
     // ideally we would build a whole crate and compress each file, taking the average
     let html = std::fs::read_to_string("benches/struct.CaptureMatches.html").unwrap();
     let html_slice = html.as_bytes();
     c.bench_function("compress regex html", |b| {
-        b.iter(|| compress(black_box(html_slice)))
+        b.iter(|| compress(black_box(html_slice, ALGORITHM)))
     });
-    let (compressed, alg) = compress(html_slice).unwrap();
+    let compressed = compress(html_slice, ALGORITHM).unwrap();
     c.bench_function("decompress regex html", |b| {
-        b.iter(|| decompress(black_box(compressed.as_slice()), alg))
+        b.iter(|| decompress(black_box(compressed.as_slice()), ALGORITHM))
     });
 }
 

src/bin/cratesfyi.rs

@@ -1,17 +1,19 @@
 use std::env;
 use std::path::PathBuf;
+use std::sync::Arc;
 
 use cratesfyi::db::{self, add_path_into_database, connect_db};
 use cratesfyi::utils::{add_crate_to_queue, remove_crate_priority, set_crate_priority};
 use cratesfyi::Server;
 use cratesfyi::{DocBuilder, DocBuilderOptions, RustwideBuilder};
+use failure::Error;
 use structopt::StructOpt;
 
-pub fn main() {
+pub fn main() -> Result<(), Error> {
     let _ = dotenv::dotenv();
     logger_init();
 
-    CommandLine::from_args().handle_args();
+    CommandLine::from_args().handle_args()
 }
 
 fn logger_init() {
@@ -79,19 +81,23 @@ enum CommandLine {
 }
 
 impl CommandLine {
-    pub fn handle_args(self) {
+    pub fn handle_args(self) -> Result<(), Error> {
+        let config = Arc::new(cratesfyi::Config::from_env()?);
+
         match self {
             Self::Build(build) => build.handle_args(),
             Self::StartWebServer {
                 socket_addr,
                 reload_templates,
             } => {
-                Server::start(Some(&socket_addr), reload_templates);
+                Server::start(Some(&socket_addr), reload_templates, config);
             }
-            Self::Daemon { foreground } => cratesfyi::utils::start_daemon(!foreground),
+            Self::Daemon { foreground } => cratesfyi::utils::start_daemon(!foreground, config),
             Self::Database { subcommand } => subcommand.handle_args(),
             Self::Queue { subcommand } => subcommand.handle_args(),
         }
+
+        Ok(())
     }
 }
 

src/config.rs

@@ -0,0 +1,37 @@
+use failure::{bail, Error, Fail, ResultExt};
+use std::env::VarError;
+use std::str::FromStr;
+use std::sync::Arc;
+
+#[derive(Debug)]
+pub struct Config {
+    pub(crate) max_file_size: usize,
+    pub(crate) max_file_size_html: usize,
+}
+
+impl Config {
+    pub fn from_env() -> Result<Self, Error> {
+        Ok(Self {
+            max_file_size: env("DOCSRS_MAX_FILE_SIZE", 50 * 1024 * 1024)?,
+            max_file_size_html: env("DOCSRS_MAX_FILE_SIZE_HTML", 5 * 1024 * 1024)?,
+        })
+    }
+}
+
+impl iron::typemap::Key for Config {
+    type Value = Arc<Config>;
+}
+
+fn env<T>(var: &str, default: T) -> Result<T, Error>
+where
+    T: FromStr,
+    T::Err: Fail,
+{
+    match std::env::var(var) {
+        Ok(content) => Ok(content
+            .parse::<T>()
+            .with_context(|_| format!("failed to parse configuration variable {}", var))?),
+        Err(VarError::NotPresent) => Ok(default),
+        Err(VarError::NotUnicode(_)) => bail!("configuration variable {} is not UTF-8", var),
+    }
+}

src/db/file.rs

@@ -13,8 +13,8 @@ use std::path::{Path, PathBuf};
 
 pub(crate) use crate::storage::Blob;
 
-pub(crate) fn get_path(conn: &Connection, path: &str) -> Result<Blob> {
-    Storage::new(conn).get(path)
+pub(crate) fn get_path(conn: &Connection, path: &str, max_size: usize) -> Result<Blob> {
+    Storage::new(conn).get(path, max_size)
 }
 
 /// Store all files in a directory and return [[mimetype, filename]] as Json

src/error.rs

@@ -5,3 +5,14 @@ use std::result::Result as StdResult;
 pub(crate) use failure::Error;
 
 pub type Result<T> = StdResult<T, Error>;
+
+#[derive(Debug, Copy, Clone)]
+pub(crate) struct SizeLimitReached;
+
+impl std::error::Error for SizeLimitReached {}
+
+impl std::fmt::Display for SizeLimitReached {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "the size limit for the buffer was reached")
+    }
+}

src/lib.rs

@@ -2,11 +2,13 @@
 //! documentation of crates for the Rust Programming Language.
 #![allow(clippy::cognitive_complexity)]
 
+pub use self::config::Config;
 pub use self::docbuilder::options::DocBuilderOptions;
 pub use self::docbuilder::DocBuilder;
 pub use self::docbuilder::RustwideBuilder;
 pub use self::web::Server;
 
+mod config;
 pub mod db;
 mod docbuilder;
 mod error;

src/storage/database.rs

@@ -16,20 +16,38 @@ impl<'a> DatabaseBackend<'a> {
         Self { conn }
     }
 
-    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
+    pub(super) fn get(&self, path: &str, max_size: usize) -> Result<Blob, Error> {
         use std::convert::TryInto;
 
+        // The maximum size for a BYTEA (the type used for `content`) is 1GB, so this cast is safe:
+        // https://www.postgresql.org/message-id/162867790712200946i7ba8eb92v908ac595c0c35aee%40mail.gmail.com
+        let max_size = max_size.min(std::i32::MAX as usize) as i32;
+
+        // The size limit is checked at the database level, to avoid receiving data altogether if
+        // the limit is exceeded.
         let rows = self.conn.query(
-            "SELECT path, mime, date_updated, content, compression
+            "SELECT
+                 path, mime, date_updated, compression,
+                 (CASE WHEN LENGTH(content) <= $2 THEN content ELSE NULL END) AS content,
+                 (LENGTH(content) > $2) AS is_too_big
              FROM files
              WHERE path = $1;",
-            &[&path],
+            &[&path, &(max_size)],
         )?;
 
         if rows.is_empty() {
             Err(PathNotFoundError.into())
         } else {
             let row = rows.get(0);
+
+            if row.get("is_too_big") {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::Other,
+                    crate::error::SizeLimitReached,
+                )
+                .into());
+            }
+
             let compression = row.get::<_, Option<i32>>("compression").map(|i| {
                 i.try_into()
                     .expect("invalid compression algorithm stored in database")
@@ -91,22 +109,69 @@ mod tests {
                     content: "Hello world!".bytes().collect(),
                     compression: None,
                 },
-                backend.get("dir/foo.txt")?
+                backend.get("dir/foo.txt", std::usize::MAX)?
             );
 
             // Test that other files are not returned
             assert!(backend
-                .get("dir/bar.txt")
+                .get("dir/bar.txt", std::usize::MAX)
                 .unwrap_err()
                 .downcast_ref::<PathNotFoundError>()
                 .is_some());
             assert!(backend
-                .get("foo.txt")
+                .get("foo.txt", std::usize::MAX)
                 .unwrap_err()
                 .downcast_ref::<PathNotFoundError>()
                 .is_some());
 
             Ok(())
         });
     }
+
+    #[test]
+    fn test_get_too_big() {
+        const MAX_SIZE: usize = 1024;
+
+        crate::test::wrapper(|env| {
+            let conn = env.db().conn();
+            let backend = DatabaseBackend::new(&conn);
+
+            let small_blob = Blob {
+                path: "small-blob.bin".into(),
+                mime: "text/plain".into(),
+                date_updated: Utc::now(),
+                content: vec![0; MAX_SIZE],
+                compression: None,
+            };
+            let big_blob = Blob {
+                path: "big-blob.bin".into(),
+                mime: "text/plain".into(),
+                date_updated: Utc::now(),
+                content: vec![0; MAX_SIZE * 2],
+                compression: None,
+            };
+
+            let transaction = conn.transaction()?;
+            backend
+                .store_batch(std::slice::from_ref(&small_blob), &transaction)
+                .unwrap();
+            backend
+                .store_batch(std::slice::from_ref(&big_blob), &transaction)
+                .unwrap();
+            transaction.commit()?;
+
+            let blob = backend.get("small-blob.bin", MAX_SIZE).unwrap();
+            assert_eq!(blob.content.len(), small_blob.content.len());
+
+            assert!(backend
+                .get("big-blob.bin", MAX_SIZE)
+                .unwrap_err()
+                .downcast_ref::<std::io::Error>()
+                .and_then(|io| io.get_ref())
+                .and_then(|err| err.downcast_ref::<crate::error::SizeLimitReached>())
+                .is_some());
+
+            Ok(())
+        });
+    }
 }