feat: allow setting max bytes to proxy/embed

Cargo.toml

@@ -20,6 +20,7 @@ env_logger = "0.8"
 dotenv = "0.15"
 tempfile = "3.2.0"
 ffprobe = "0.3.0"
+encoding_rs = "0.8"
 moka = { version = "0.10", features = ["future"]}
 
 validator = { version = "0.15.0", features = ["derive"] }

src/routes/proxy.rs

@@ -4,7 +4,7 @@ use actix_web::web::Bytes;
 use actix_web::{web::Query, HttpResponse, Responder};
 use serde::Deserialize;
 
-use crate::util::request::fetch;
+use crate::util::request::{fetch, get_bytes};
 use crate::util::result::Error;
 
 lazy_static! {
@@ -24,14 +24,11 @@ pub struct Parameters {
 }
 
 async fn proxy(url: String) -> Result<Bytes, Error> {
-    let (resp, mime) = fetch(&url).await?;
+    let (mut resp, mime) = fetch(&url).await?;
 
     if matches!(mime.type_(), mime::IMAGE | mime::VIDEO) {
-        let body = resp
-            .bytes()
-            .await
-            .map_err(|_| Error::FailedToConsumeBytes)?;
-        Ok(body)
+        let bytes = get_bytes(&mut resp).await?;
+        Ok(bytes)
     } else {
         Err(Error::NotAllowedToProxy)
     }

src/util/request.rs

@@ -1,12 +1,17 @@
 use std::time::Duration;
 
+use actix_web::web::Bytes;
+use encoding_rs::{Encoding, UTF_8_INIT};
 use mime::Mime;
-use reqwest::{header::CONTENT_TYPE, Client, Response};
+use reqwest::{
+    header::{self, CONTENT_TYPE},
+    Client, Response,
+};
 use scraper::Html;
 use std::io::Write;
 use tempfile::NamedTempFile;
 
-use super::result::Error;
+use super::{result::Error, variables::MAX_BYTES};
 
 lazy_static! {
     static ref CLIENT: Client = reqwest::Client::builder()
@@ -42,9 +47,41 @@ pub async fn fetch(url: &str) -> Result<(Response, Mime), Error> {
     Ok((resp, mime))
 }
 
-pub async fn consume_fragment(resp: Response) -> Result<Html, Error> {
-    let body = resp.text().await.map_err(|_| Error::FailedToConsumeText)?;
-    Ok(Html::parse_document(&body))
+pub async fn get_bytes(resp: &mut Response) -> Result<Bytes, Error> {
+    let content_length = resp.content_length().unwrap_or(0) as usize;
+    if content_length > *MAX_BYTES {
+        return Err(Error::ExceedsMaxBytes);
+    }
+    let mut bytes = Vec::with_capacity(content_length);
+    while let Some(chunk) = resp
+        .chunk()
+        .await
+        .map_err(|_| Error::FailedToConsumeBytes)?
+    {
+        if bytes.len() + chunk.len() > *MAX_BYTES {
+            return Err(Error::ExceedsMaxBytes);
+        }
+        bytes.extend(chunk)
+    }
+    Ok(Bytes::from(bytes))
+}
+
+pub async fn consume_fragment(mut resp: Response) -> Result<Html, Error> {
+    let bytes = get_bytes(&mut resp).await?;
+
+    let content_type = resp
+        .headers()
+        .get(header::CONTENT_TYPE)
+        .and_then(|value| value.to_str().ok())
+        .and_then(|value| value.parse::<Mime>().ok());
+    let encoding_name = content_type
+        .as_ref()
+        .and_then(|mime| mime.get_param("charset").map(|charset| charset.as_str()))
+        .unwrap_or("utf-8");
+    let encoding = Encoding::for_label(encoding_name.as_bytes()).unwrap_or(&UTF_8_INIT);
+
+    let (text, _, _) = encoding.decode(&bytes);
+    Ok(Html::parse_document(&text))
 }
 
 pub fn determine_video_size(path: &std::path::Path) -> Result<(isize, isize), Error> {
@@ -62,11 +99,8 @@ pub fn determine_video_size(path: &std::path::Path) -> Result<(isize, isize), Er
     Err(Error::ProbeError)
 }
 
-pub async fn consume_size(resp: Response, mime: Mime) -> Result<(isize, isize), Error> {
-    let bytes = resp
-        .bytes()
-        .await
-        .map_err(|_| Error::FailedToConsumeBytes)?;
+pub async fn consume_size(mut resp: Response, mime: Mime) -> Result<(isize, isize), Error> {
+    let bytes = get_bytes(&mut resp).await?;
 
     match mime.type_() {
         mime::IMAGE => {

src/util/result.rs

@@ -17,6 +17,7 @@ pub enum Error {
     MissingContentType,
     NotAllowedToProxy,
     ConversionFailed,
+    ExceedsMaxBytes,
     ReqwestFailed,
     RequestFailed,
     ProbeError,
@@ -45,6 +46,7 @@ impl ResponseError for Error {
             Error::MissingContentType => StatusCode::BAD_REQUEST,
             Error::NotAllowedToProxy => StatusCode::BAD_REQUEST,
             Error::ConversionFailed => StatusCode::INTERNAL_SERVER_ERROR,
+            Error::ExceedsMaxBytes => StatusCode::BAD_REQUEST,
             Error::ReqwestFailed => StatusCode::INTERNAL_SERVER_ERROR,
             Error::RequestFailed => StatusCode::BAD_REQUEST,
             Error::ProbeError => StatusCode::INTERNAL_SERVER_ERROR,

src/util/variables.rs

@@ -4,4 +4,6 @@ lazy_static! {
     // Application Settings
     pub static ref HOST: String =
         env::var("JANUARY_HOST").expect("Missing JANUARY_HOST environment variable.");
+    pub static ref MAX_BYTES: usize =
+        env::var("JANUARY_MAX_BYTES").unwrap_or("104857600".to_string()).parse().expect("Invalid JANUARY_MAX_BYTES environment variable.");
 }