feat: add nsfw checks

Cargo.toml

@@ -39,6 +39,9 @@ derive_more = "0.99.17"
 regress = "0.4.1"
 once_cell = "1.10.0"
 futures-util = "0.3.21"
+rayon = "1.5.2"
+http = "0.2.7"
+url = "2.2.2"
 
 [dependencies.tokio]
 version = "1.18.1"

README.md

@@ -29,6 +29,7 @@
 - `PORT` - the port that the application will run (optional, defaults to `3000`)
 - `REDIS_URL` - the address of your redis database (required)
 - `FULLSCREEN_SCREENSHOT` - if set, it will screenshot the whole website (optional)
+- `CHECK_IF_NSFW` - if set, it will check if the url is marked as NSFW (optional)
 
 ### Railway
 
@@ -175,7 +176,11 @@ Hello, world!
 
 Creates a screenshot.
 
-JSON payload with the `url` key.
+#### Payload
+
+- `url` - The url of the website. (string, required)
+- `fullscreen` - If you want to take a fullscreen screenshot. (boolean, optional, overrides the `FULLSCREEN_SCREENSHOT` environment variable)
+- `check_nsfw` - If you want to check if the url is marked as NSFW. (boolean, optional, overrides the `CHECK_IF_NSFW` environment variable)
 
 Example Payload:
 
@@ -192,7 +197,9 @@ Example Response
   "slug": "abcdefghijk",
   "path": "/s/abcdefghijk",
   "metadata": {
-    "url": "https://rust-lang.org"
+    "url": "https://rust-lang.org",
+    "fullscreen": false,
+    "check_nsfw": false
   }
 }
 ```

src/cdp.rs

@@ -0,0 +1,111 @@
+/// MIT License
+///
+/// Copyright (c) 2019-2021 Stephen Pryde and the thirtyfour contributors
+///
+/// Permission is hereby granted, free of charge, to any person obtaining a copy
+/// of this software and associated documentation files (the "Software"), to deal
+/// in the Software without restriction, including without limitation the rights
+/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+/// copies of the Software, and to permit persons to whom the Software is
+/// furnished to do so, subject to the following conditions:
+///
+/// The above copyright notice and this permission notice shall be included in all
+/// copies or substantial portions of the Software.
+///
+/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+/// SOFTWARE.
+use fantoccini::wd::WebDriverCompatibleCommand;
+use http::Method;
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+use url::{ParseError, Url};
+
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(rename = "lowercase")]
+pub enum ConnectionType {
+    None,
+    Cellular2G,
+    Cellular3G,
+    Cellular4G,
+    Bluetooth,
+    Ethernet,
+    Wifi,
+    Wimax,
+    Other,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(rename = "camelCase")]
+pub struct NetworkConditions {
+    pub offline: bool,
+    pub latency: u32,
+    pub download_throughput: i32,
+    pub upload_throughput: i32,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub connection_type: Option<ConnectionType>,
+}
+
+#[derive(Debug)]
+pub enum ChromeCommand {
+    LaunchApp(String),
+    GetNetworkConditions,
+    SetNetworkConditions(NetworkConditions),
+    ExecuteCdpCommand(String, Value),
+    GetSinks,
+    GetIssueMessage,
+    SetSinkToUse(String),
+    StartTabMirroring(String),
+    StopCasting(String),
+}
+
+impl WebDriverCompatibleCommand for ChromeCommand {
+    fn endpoint(&self, base_url: &Url, session_id: Option<&str>) -> Result<Url, ParseError> {
+        let base = { base_url.join(&format!("session/{}/", session_id.as_ref().unwrap()))? };
+        match &self {
+            ChromeCommand::LaunchApp(_) => base.join("chromium/launch_app"),
+            ChromeCommand::GetNetworkConditions | ChromeCommand::SetNetworkConditions(_) => {
+                base.join("chromium/network_conditions")
+            },
+            ChromeCommand::ExecuteCdpCommand(..) => base.join("goog/cdp/execute"),
+            ChromeCommand::GetSinks => base.join("goog/cast/get_sinks"),
+            ChromeCommand::GetIssueMessage => base.join("goog/cast/get_issue_message"),
+            ChromeCommand::SetSinkToUse(_) => base.join("goog/cast/set_sink_to_use"),
+            ChromeCommand::StartTabMirroring(_) => base.join("goog/cast/start_tab_mirroring"),
+            ChromeCommand::StopCasting(_) => base.join("goog/cast/stop_casting"),
+        }
+    }
+
+    fn method_and_body(&self, _request_url: &Url) -> (Method, Option<String>) {
+        let mut method = Method::GET;
+        let mut body = None;
+
+        match &self {
+            ChromeCommand::LaunchApp(app_id) => {
+                method = Method::POST;
+                body = Some(json!({ "id": app_id }).to_string())
+            },
+            ChromeCommand::SetNetworkConditions(conditions) => {
+                method = Method::POST;
+                body = Some(json!({ "network_conditions": conditions }).to_string())
+            },
+            ChromeCommand::ExecuteCdpCommand(command, params) => {
+                method = Method::POST;
+                body = Some(json!({"cmd": command, "params": params }).to_string())
+            },
+            ChromeCommand::SetSinkToUse(sink_name)
+            | ChromeCommand::StartTabMirroring(sink_name)
+            | ChromeCommand::StopCasting(sink_name) => {
+                method = Method::POST;
+                body = Some(json!({ "sinkName": sink_name }).to_string())
+            },
+            _ => {},
+        }
+
+        (method, body)
+    }
+}

src/error.rs

@@ -18,6 +18,8 @@ pub enum Error {
     Unauthorized,
     #[display(fmt = "The screenshot with that slug can't be found.")]
     ScreenshotNotFound,
+    #[display(fmt = "The url provided is marked as NSFW.")]
+    UrlNotSafeForWork,
 }
 
 impl ResponseError for Error {
@@ -32,6 +34,7 @@ impl ResponseError for Error {
             Error::InvalidUrl | Error::MissingAuthToken => StatusCode::BAD_REQUEST,
             Error::Unauthorized => StatusCode::UNAUTHORIZED,
             Error::ScreenshotNotFound => StatusCode::NOT_FOUND,
+            Error::UrlNotSafeForWork => StatusCode::FORBIDDEN,
         }
     }
 }

src/main.rs

@@ -1,4 +1,4 @@
-#![feature(let_chains)]
+#![feature(let_chains, pattern)]
 
 #[macro_use]
 extern crate tracing;
@@ -13,13 +13,15 @@ use actix_web::{web, App, Error, HttpServer};
 use fantoccini::{Client, ClientBuilder};
 use portpicker::pick_unused_port;
 use providers::{Provider, Storage};
+use reqwest::Client as ReqwestClient;
 use serde_json::Map;
 use tokio::process::Command;
 use tokio_process_stream::ProcessLineStream;
 use tokio_stream::StreamExt;
 use tracing_actix_web::TracingLogger;
 use util::{initialize_tracing, load_env};
 
+pub mod cdp;
 pub mod error;
 pub mod middlewares;
 pub mod providers;
@@ -32,6 +34,7 @@ pub type Result<T, E = Error> = anyhow::Result<T, E>;
 pub struct State {
     pub browser: Arc<Client>,
     pub storage: Arc<Storage>,
+    pub reqwest: ReqwestClient,
 }
 
 #[actix_web::main]
@@ -84,6 +87,7 @@ async fn main() -> anyhow::Result<()> {
     let state = web::Data::new(State {
         browser: Arc::new(client.clone()),
         storage: Arc::new(Storage::new()),
+        reqwest: ReqwestClient::new(),
     });
 
     let port =

src/middlewares/auth.rs

@@ -9,6 +9,7 @@ use futures_util::future::LocalBoxFuture;
 
 use crate::error::Error as Errors;
 
+#[derive(Debug)]
 pub struct Auth;
 
 impl<S, B> Transform<S, ServiceRequest> for Auth
@@ -30,6 +31,7 @@ where
     }
 }
 
+#[derive(Debug)]
 pub struct AuthMiddleware<S> {
     service: S,
 }

src/routes/screenshot.rs

@@ -8,19 +8,26 @@ use serde_json::json;
 
 use crate::error::Error;
 use crate::providers::Provider;
-use crate::util::check_if_url;
+use crate::util::{check_if_nsfw, check_if_url};
 use crate::{Result, State};
 
 #[inline]
 fn default_fullscreen() -> bool {
     env::var("FULLSCREEN_SCREENSHOT").is_ok()
 }
 
+#[inline]
+fn default_check_nsfw() -> bool {
+    env::var("CHECK_IF_NSFW").is_ok()
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct RequestData {
     url: String,
     #[serde(default = "default_fullscreen")]
     fullscreen: bool,
+    #[serde(default = "default_check_nsfw")]
+    check_nsfw: bool,
 }
 
 #[post("/screenshot")]
@@ -30,9 +37,21 @@ pub async fn screenshot(
 ) -> Result<HttpResponse, Error> {
     check_if_url(&payload.url).map_err(|_| Error::InvalidUrl)?;
 
+    let req =
+        data.reqwest.head(&payload.url).send().await.expect("Failed sending head request to url");
+    let url = req.url();
+
+    if payload.check_nsfw
+        && check_if_nsfw(url.host_str().expect("Failed getting url host"))
+            .await
+            .expect("Failed checking if nsfw")
+    {
+        return Err(Error::UrlNotSafeForWork);
+    }
+
     let client = &data.browser;
 
-    client.goto(&payload.url).await.expect("Failed navigating to site");
+    client.goto(url.as_str()).await.expect("Failed navigating to site");
     client.set_window_size(1980, 1080).await.expect("Failed setting window size");
     client
         .execute("document.body.style.overflow = 'hidden'", vec![])
@@ -87,7 +106,9 @@ pub async fn screenshot(
         "slug": &slug,
         "path": format!("/s/{}", &slug),
         "metadata": {
-            "url": payload.url
+            "url": &payload.url,
+            "fullscreen": payload.fullscreen,
+            "check_nsfw": payload.check_nsfw
         }
     })))
 }