feat: add support for clamav virus scanning

closes #27

Author: insertish

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "autumn"
-version = "1.1.8"
+version = "1.1.9"
 authors = ["Paul Makles <[email protected]>"]
 edition = "2018"
 
@@ -38,3 +38,6 @@ actix-web = "4.0.0-beta.9"
 actix-cors = "0.6.0-beta.2"
 actix-files = "0.6.0-beta.7"
 actix-multipart = "0.4.0-beta.6"
+
+# virus scanning
+revolt_clamav-client = { version = "0.1.5" }

Cargo.toml

@@ -3,6 +3,7 @@ pub mod db;
 pub mod routes;
 pub mod util;
 pub mod version;
+pub mod virus_scan;
 
 use futures::StreamExt;
 use util::variables::{CONFIG, HOST, LOCAL_STORAGE_PATH, USE_S3};
@@ -38,6 +39,8 @@ async fn main() -> std::io::Result<()> {
 
     info!("Starting Autumn server.");
 
+    virus_scan::init();
+
     db::connect().await;
 
     if !*USE_S3 {

src/main.rs

@@ -1,7 +1,7 @@
 use crate::config::{get_tag, Config, ContentType};
 use crate::db::*;
 use crate::util::result::Error;
-use crate::util::variables::{get_s3_bucket, LOCAL_STORAGE_PATH, USE_S3};
+use crate::util::variables::{get_s3_bucket, LOCAL_STORAGE_PATH, USE_S3, USE_CLAMD, CLAMD_HOST};
 
 use actix_multipart::Multipart;
 use actix_web::{web, HttpRequest, HttpResponse};
@@ -203,6 +203,17 @@ pub async fn post(req: HttpRequest, mut payload: Multipart) -> Result<HttpRespon
                 if inspect(&buf).is_text() {
                     Metadata::Text
                 } else {
+                    // Scan the file for malware
+                    if *USE_CLAMD {
+                        let scan_response =
+                            revolt_clamav_client::scan_buffer_tcp(&buf, CLAMD_HOST.to_string(), None).unwrap();
+
+                        let file_clean = revolt_clamav_client::clean(&scan_response).unwrap();
+                        if !file_clean {
+                            return Err(Error::Malware)
+                        }
+                    }
+
                     Metadata::File
                 }
             }

src/routes/upload.rs

@@ -16,6 +16,7 @@ pub enum Error {
     UnknownTag,
     ProbeError,
     NotFound,
+    Malware,
     IOError,
     S3Error,
     LabelMe,
@@ -42,6 +43,7 @@ impl ResponseError for Error {
             Error::IOError => StatusCode::INTERNAL_SERVER_ERROR,
             Error::S3Error => StatusCode::INTERNAL_SERVER_ERROR,
             Error::LabelMe => StatusCode::INTERNAL_SERVER_ERROR,
+            Error::Malware => StatusCode::FORBIDDEN,
         }
     }
 

src/util/result.rs

@@ -15,6 +15,8 @@ lazy_static! {
         env::var("AUTUMN_MONGO_DATABASE").unwrap_or_else(|_| "revolt".to_string());
     pub static ref CORS_ALLOWED_ORIGIN: String =
         env::var("AUTUMN_CORS_ALLOWED_ORIGIN").expect("Missing AUTUMN_CORS_ALLOWED_ORIGIN environment variable.");
+    pub static ref CLAMD_HOST: String =
+        env::var("CLAMD_HOST").expect("Missing CLAMD_HOST environment variable.");
 
     // Storage Settings
     pub static ref LOCAL_STORAGE_PATH: String =
@@ -27,6 +29,7 @@ lazy_static! {
 
     // Application Flags
     pub static ref USE_S3: bool = env::var("AUTUMN_S3_REGION").is_ok() && env::var("AUTUMN_S3_ENDPOINT").is_ok();
+    pub static ref USE_CLAMD: bool = env::var("CLAMD_HOST").is_ok();
 }
 
 pub fn get_s3_bucket(bucket: &str) -> Result<s3::Bucket, Error> {

src/util/variables.rs

@@ -0,0 +1,30 @@
+use std::time::Duration;
+
+use log::{error, info};
+
+use crate::util::variables::{CLAMD_HOST, USE_CLAMD};
+
+pub fn init() {
+    if *USE_CLAMD {
+        info!("Waiting for clamd to be ready...");
+
+        loop {
+            let clamd_available = match revolt_clamav_client::ping_tcp(CLAMD_HOST.to_string()) {
+                Ok(ping_response) => ping_response == b"PONG\0",
+                Err(_) => false,
+            };
+
+            if clamd_available {
+                info!("clamd is ready, virus protection enabled!");
+                break;
+            } else {
+                error!(
+                    "Could not ping clamd host at {}, retrying in 10 seconds...",
+                    CLAMD_HOST.to_string()
+                );
+
+                std::thread::sleep(Duration::from_secs(10));
+            }
+        }
+    }
+}