feat: add auth

BREAKING CHANGE: msrv is now nightly 1.62
jzeuzs · May 4, 2022 · 39514ad · 39514ad
1 parent 22283d2
commit 39514ad
Show file tree

Hide file tree

Showing 9 changed files with 116 additions and 11 deletions.
diff --git a/.github/workflows/continuous-delivery.yml b/.github/workflows/continuous-delivery.yml
@@ -73,7 +73,7 @@ jobs:
       - name: Setup | Rust
         uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # tag=v1.0.7
         with:
-          toolchain: stable
+          toolchain: nightly
           override: true
           profile: minimal
           target: ${{ matrix.target }}
@@ -141,7 +141,7 @@ jobs:
       - name: Setup | Rust
         uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # tag=v1.0.7
         with:
-          toolchain: stable
+          toolchain: nightly
           profile: minimal
           override: true
 

diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml
@@ -18,7 +18,7 @@ jobs:
         uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # tag=v1
         with:
           profile: minimal
-          toolchain: stable
+          toolchain: nightly
           override: true
 
       - name: Setup | Install cargo-msrv
@@ -61,7 +61,7 @@ jobs:
         uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # tag=v1
         with:
           profile: minimal
-          toolchain: stable
+          toolchain: nightly
           override: true
 
       - name: Setup | Install clippy

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "website-screenshot"
 version = "0.2.0"
-rust-version = "1.60"
+rust-version = "1.62"
 authors = ["Tomio <mail@tomio.fun>"]
 license = "MIT/Apache-2.0"
 edition = "2021"
@@ -39,6 +39,7 @@ portpicker = "0.1.1"
 derive_more = "0.99.17"
 regress = "0.4.1"
 once_cell = "1.10.0"
+futures-util = "0.3.21"
 
 [dependencies.tokio]
 version = "1.18.0"

diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # website-screenshot
 
-[![Rust: 1.60+](https://img.shields.io/badge/rust-1.60+-93450a)](https://blog.rust-lang.org/2022/04/07/Rust-1.60.0.html)
+[![Rust: 1.62+](https://img.shields.io/badge/rust-1.62+-93450a)](https://github.com/rust-lang/rust/milestone/93)
 [![Continuous Delivery](https://github.com/devtomio/website-screenshot/actions/workflows/continuous-delivery.yml/badge.svg)](https://github.com/devtomio/website-screenshot/actions/workflows/continuous-delivery.yml)
 [![Continuous Integration](https://github.com/devtomio/website-screenshot/actions/workflows/continuous-integration.yml/badge.svg)](https://github.com/devtomio/website-screenshot/actions/workflows/continuous-integration.yml)
 
@@ -12,13 +12,14 @@
 - 🚀 blazing fast
 - 👮‍‍ built-in [ratelimiter](https://github.com/antifuchs/governor)
 - 👜 built-in [storage providers](#storage-providers)
+- 🛡️ built-in [authentication](#authentication)
 - 🗼 configurable
 
 ## Deployment
 
 ### Prerequisites
 
-- [Rust] 1.60+ or greater
+- [Rust] nightly 1.62+ or greater
 - [Redis] 6 or greater
 - [Chrome] browser
 - [Chromedriver] (must match with the version your [Chrome] browser)
@@ -66,6 +67,12 @@ cargo install website-screenshot
 1. Clone this repository. e.g. `git clone https://github.com/devtomio/website-screenshot`
 2. Build the binary `cargo build --release`
 
+### Authentication
+
+Authentication will be enabled if the `AUTH_TOKEN` variable is set.
+
+It will check if the `Authorization` header sent by the user is equal to the `AUTH_TOKEN` that you set.
+
 ## Storage Providers
 
 ### Fs (Filesystem) Provider

diff --git a/src/error.rs b/src/error.rs
@@ -10,18 +10,25 @@ use derive_more::{Display, Error as DeriveError};
 pub enum Error {
     #[display(fmt = "The url that you provided was invalid.")]
     InvalidUrl,
+    #[display(
+        fmt = "Authentication was enabled but the \"Authorization\" header was not present."
+    )]
+    MissingAuthToken,
+    #[display(fmt = "Invalid token provided.")]
+    Unauthorized,
 }
 
 impl ResponseError for Error {
     fn error_response(&self) -> HttpResponse {
         HttpResponse::build(self.status_code())
-            .insert_header(ContentType::html())
+            .insert_header(ContentType::plaintext())
             .body(self.to_string())
     }
 
     fn status_code(&self) -> StatusCode {
         match self.deref() {
-            Error::InvalidUrl => StatusCode::BAD_REQUEST,
+            Error::InvalidUrl | Error::MissingAuthToken => StatusCode::BAD_REQUEST,
+            Error::Unauthorized => StatusCode::UNAUTHORIZED,
         }
     }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -1,3 +1,5 @@
+#![feature(let_chains)]
+
 #[macro_use]
 extern crate tracing;
 
@@ -6,6 +8,7 @@ use std::process::Stdio;
 use std::sync::Arc;
 
 use actix_governor::{Governor, GovernorConfigBuilder};
+use actix_web::middleware::Compress;
 use actix_web::{web, App, Error, HttpServer};
 use fantoccini::{Client, ClientBuilder};
 use portpicker::pick_unused_port;
@@ -18,8 +21,9 @@ use tracing_actix_web::TracingLogger;
 use util::{initialize_tracing, load_env};
 
 pub mod error;
-mod providers;
-mod routes;
+pub mod middlewares;
+pub mod providers;
+pub mod routes;
 pub mod util;
 
 pub type Result<T, E = Error> = anyhow::Result<T, E>;
@@ -89,6 +93,8 @@ async fn main() -> anyhow::Result<()> {
 
     HttpServer::new(move || {
         App::new()
+            .wrap(Compress::default())
+            .wrap(middlewares::Auth)
             .wrap(TracingLogger::default())
             .wrap(Governor::new(&governor_config))
             .app_data(state.clone())

diff --git a/src/middlewares/auth.rs b/src/middlewares/auth.rs
@@ -0,0 +1,80 @@
+use std::env;
+use std::future::{ready, Ready};
+
+use actix_web::body::EitherBody;
+use actix_web::dev::{self, Service, ServiceRequest, ServiceResponse, Transform};
+use actix_web::http::header;
+use actix_web::{Error, HttpResponse};
+use futures_util::future::LocalBoxFuture;
+
+use crate::error::Error as Errors;
+
+pub struct Auth;
+
+impl<S, B> Transform<S, ServiceRequest> for Auth
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type InitError = ();
+    type Transform = AuthMiddleware<S>;
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        ready(Ok(AuthMiddleware {
+            service,
+        }))
+    }
+}
+
+pub struct AuthMiddleware<S> {
+    service: S,
+}
+
+impl<S, B> Service<ServiceRequest> for AuthMiddleware<S>
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    dev::forward_ready!(service);
+
+    fn call(&self, req: ServiceRequest) -> Self::Future {
+        let auth_token = env::var("AUTH_TOKEN").ok();
+        let headers = req.headers().clone();
+        let auth_header = headers.get(header::AUTHORIZATION);
+        let (req, pl) = req.into_parts();
+
+        if let Some(auth_token) = auth_token && req.path() == "/screenshot" {
+            match auth_header {
+                Some(auth) => {
+                    let auth = auth.to_str().expect("Failed converting to str").to_owned();
+
+                    if auth_token != auth {
+                        let res = HttpResponse::from_error(Errors::Unauthorized)
+                            .map_into_right_body::<B>();
+
+                        return Box::pin(async { Ok(ServiceResponse::new(req, res)) });
+                    }
+                },
+                None => {
+                    let res = HttpResponse::from_error(Errors::MissingAuthToken)
+                        .map_into_right_body::<B>();
+
+                    return Box::pin(async { Ok(ServiceResponse::new(req, res)) });
+                },
+            };
+        }
+
+        let res = self.service.call(ServiceRequest::from_parts(req, pl));
+
+        Box::pin(async move { res.await.map(ServiceResponse::map_into_left_body) })
+    }
+}
diff --git a/src/middlewares/mod.rs b/src/middlewares/mod.rs
@@ -0,0 +1,3 @@
+mod auth;
+
+pub use auth::*;