use std::{sync::Arc, time::Duration}; use crate::{ pict::{Details, Extension, Images, Upload, Uploads}, Error, }; use actix_web::{body::BodyStream, http::StatusCode, web, HttpRequest, HttpResponse}; use awc::Client; use url::Url; pub(crate) static VALID_SIZES: &[u16] = &[80, 160, 320, 640, 1080, 2160]; pub(crate) struct Connection { upstream: Url, client: Client, } impl std::fmt::Debug for Connection { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Connection") .field("upstream", &self.upstream.as_str()) .field("client", &"Client") .finish() } } impl Connection { pub(crate) fn new(upstream: Url, client: Client) -> Self { Connection { upstream, client } } #[tracing::instrument(skip_all)] pub(crate) async fn claim(&self, upload: Upload) -> Result { let mut attempts = 0; const CLAIM_ATTEMPT_LIMIT: usize = 10; loop { match self.client.get(self.claim_url(&upload)).send().await { Ok(mut res) => { match res.status() { StatusCode::OK => { return res .json::() .await .map_err(|_| UploadError::Json.into()); } StatusCode::NO_CONTENT => { // continue } _ => { let images = res.json::().await.map_err(|_| UploadError::Json)?; let code = images.code().unwrap_or("uknown-error").to_string(); let msg = images.msg().to_string(); return Err(UploadError::UploadFailure(code, msg).into()); } } } Err(_) => { attempts += 1; if attempts > CLAIM_ATTEMPT_LIMIT { return Err(UploadError::Status.into()); } tokio::time::sleep(Duration::from_secs(1)).await; // continue } } } } #[tracing::instrument(skip_all)] pub(crate) async fn thumbnail( &self, size: u16, file: &str, extension: Extension, req: &HttpRequest, ) -> Result { if !VALID_SIZES.contains(&size) { return Err(UploadError::Size(size).into()); } self.proxy(self.thumbnail_url(size, file, extension), req) .await } #[tracing::instrument(skip_all)] pub(crate) async fn image(&self, file: &str, req: &HttpRequest) -> Result { self.proxy(self.image_url(file), req).await } #[tracing::instrument(skip_all)] pub(crate) async fn details(&self, file: &str) -> Result { let mut response = self .client .get(self.details_url(file)) .send() .await .map_err(|e| UploadError::Request(Arc::from(e.to_string())))?; if !response.status().is_success() { return Err(UploadError::Status.into()); } response.json().await.map_err(|_| UploadError::Json.into()) } #[tracing::instrument(skip_all)] pub(crate) async fn upload( &self, req: &HttpRequest, body: web::Payload, ) -> Result { let client_request = self.client.request_from(self.upload_url(), req.head()); let mut client_request = if let Some(addr) = req.head().peer_addr { client_request.append_header(("X-Forwarded-For", addr.to_string())) } else { client_request }; client_request.headers_mut().remove("Accept-Encoding"); let mut res = client_request .send_stream(body) .await .map_err(|e| UploadError::Request(Arc::from(e.to_string())))?; let uploads = res.json::().await.map_err(|_| UploadError::Json)?; Ok(uploads) } #[tracing::instrument(skip_all)] pub(crate) async fn delete(&self, file: &str, token: &str) -> Result<(), Error> { let res = self .client .delete(self.delete_url(file, token)) .send() .await .map_err(|e| UploadError::Request(Arc::from(e.to_string())))?; if !res.status().is_success() { return Err(UploadError::Status.into()); } Ok(()) } fn claim_url(&self, upload: &Upload) -> String { let mut url = self.upstream.clone(); url.set_path("/image/backgrounded/claim"); url.set_query(Some(&format!("upload_id={}", upload.id()))); url.to_string() } fn upload_url(&self) -> String { let mut url = self.upstream.clone(); url.set_path("/image/backgrounded"); url.to_string() } fn thumbnail_url(&self, size: u16, file: &str, extension: Extension) -> String { let mut url = self.upstream.clone(); url.set_path(&format!("/image/process.{extension}")); url.set_query(Some(&format!("src={file}&resize={size}"))); url.to_string() } fn image_url(&self, file: &str) -> String { let mut url = self.upstream.clone(); url.set_path(&format!("/image/original/{file}")); url.to_string() } fn details_url(&self, file: &str) -> String { let mut url = self.upstream.clone(); url.set_path(&format!("/image/details/original/{file}")); url.to_string() } fn delete_url(&self, file: &str, token: &str) -> String { let mut url = self.upstream.clone(); url.set_path(&format!("/image/delete/{token}/{file}")); url.to_string() } #[tracing::instrument(skip_all)] async fn proxy(&self, url: String, req: &HttpRequest) -> Result { let client_request = self.client.request_from(url, req.head()); let client_request = if let Some(addr) = req.head().peer_addr { client_request.append_header(("X-Forwarded-For", addr.to_string())) } else { client_request }; let res = client_request .no_decompress() .send() .await .map_err(|e| UploadError::Request(Arc::from(e.to_string())))?; let mut client_res = HttpResponse::build(res.status()); for (name, value) in res.headers().iter().filter(|(h, _)| *h != "connection") { client_res.append_header((name.clone(), value.clone())); } Ok(client_res.body(BodyStream::new(res))) } } #[derive(Clone, Debug, thiserror::Error)] pub(crate) enum UploadError { #[error("There was an error making the upstream request: {0}")] Request(Arc), #[error("There was an error parsing the upstream response")] Json, #[error("Request returned bad HTTP status")] Status, #[error("Request failed with {0}: {1}")] UploadFailure(String, String), #[error("Requested size {0} is invalid")] Size(u16), }