diff --git a/src/main.rs b/src/main.rs index 9ea9596..fac2745 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ #![warn(clippy::pedantic, clippy::nursery)] +#![allow(clippy::future_not_send)] // We're end users, so this is ok use std::env::{self, VarError}; use std::time::Duration; @@ -12,7 +13,7 @@ use awc::{error::SendRequestError, Client}; use log::{debug, error, info, warn}; use lru::LruCache; use parking_lot::RwLock; -use ping::{PingRequest, PingResponse}; +use ping::{Request, Response}; use rustls::sign::{CertifiedKey, RSASigningKey}; use rustls::PrivateKey; use rustls::{Certificate, NoClientAuth, ResolvesServerCert, ServerConfig}; @@ -40,24 +41,32 @@ struct ServerState { tls_config: Tls, disabled_tokens: bool, url: String, - cache: LruCache<(String, String, bool), Vec>, + cache: LruCache<(String, String, bool), CachedImage>, +} + +struct CachedImage { + data: Vec, + content_type: Option>, + content_length: Option>, + last_modified: Option>, } impl ServerState { async fn init(config: &Config) -> Result { let resp = Client::new() .post(CONTROL_CENTER_PING_URL) - .send_json(&PingRequest::from(config)) + .send_json(&Request::from(config)) .await; match resp { - Ok(mut resp) => match resp.json::().await { + Ok(mut resp) => match resp.json::().await { Ok(resp) => { let key = resp .token_key - .and_then(|key| match PrecomputedKey::from_slice(key.as_bytes()) { - Some(key) => Some(key), - None => { + .and_then(|key| { + if let Some(key) = PrecomputedKey::from_slice(key.as_bytes()) { + Some(key) + } else { error!("Failed to parse token key: got {}", key); None } @@ -106,7 +115,7 @@ impl ServerState { } } -struct RwLockServerState(RwLock); +pub struct RwLockServerState(RwLock); impl ResolvesServerCert for RwLockServerState { fn resolve(&self, _: rustls::ClientHello) -> Option { @@ -166,7 +175,7 @@ async fn main() -> Result<(), std::io::Error> { .await } -struct Config { +pub struct Config { secret: String, port: u16, disk_quota: usize, diff --git a/src/ping.rs b/src/ping.rs index f3e52d4..073862c 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -9,7 +9,7 @@ use url::Url; use crate::{client_api_version, Config, RwLockServerState, CONTROL_CENTER_PING_URL}; #[derive(Serialize)] -pub struct PingRequest<'a> { +pub struct Request<'a> { secret: &'a str, port: u16, disk_space: usize, @@ -18,7 +18,7 @@ pub struct PingRequest<'a> { tls_created_at: Option, } -impl<'a> PingRequest<'a> { +impl<'a> Request<'a> { fn from_config_and_state(config: &'a Config, state: &'a Arc) -> Self { Self { secret: &config.secret, @@ -31,7 +31,7 @@ impl<'a> PingRequest<'a> { } } -impl<'a> From<&'a Config> for PingRequest<'a> { +impl<'a> From<&'a Config> for Request<'a> { fn from(config: &'a Config) -> Self { Self { secret: &config.secret, @@ -45,32 +45,32 @@ impl<'a> From<&'a Config> for PingRequest<'a> { } #[derive(Deserialize)] -pub(crate) struct PingResponse { - pub(crate) image_server: Url, - pub(crate) latest_build: usize, - pub(crate) url: String, - pub(crate) token_key: Option, - pub(crate) compromised: bool, - pub(crate) paused: bool, - pub(crate) disabled_tokens: bool, - pub(crate) tls: Option, +pub struct Response { + pub image_server: Url, + pub latest_build: usize, + pub url: String, + pub token_key: Option, + pub compromised: bool, + pub paused: bool, + pub disabled_tokens: bool, + pub tls: Option, } #[derive(Deserialize)] -pub(crate) struct Tls { +pub struct Tls { pub created_at: String, pub private_key: Vec, pub certificate: Vec, } -pub(crate) async fn update_server_state(req: &Config, data: &mut Arc) { - let req = PingRequest::from_config_and_state(req, data); +pub async fn update_server_state(req: &Config, data: &mut Arc) { + let req = Request::from_config_and_state(req, data); let resp = Client::new() .post(CONTROL_CENTER_PING_URL) .send_json(&req) .await; match resp { - Ok(mut resp) => match resp.json::().await { + Ok(mut resp) => match resp.json::().await { Ok(resp) => { let mut write_guard = data.0.write(); diff --git a/src/routes.rs b/src/routes.rs index ac22e6a..d55ba42 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -1,7 +1,12 @@ use std::convert::Infallible; -use actix_web::{dev::HttpResponseBuilder, get, web::Data, HttpRequest, Responder}; -use actix_web::{web::Path, HttpResponse}; +use actix_web::dev::HttpResponseBuilder; +use actix_web::http::header::{ + ACCESS_CONTROL_ALLOW_ORIGIN, ACCESS_CONTROL_EXPOSE_HEADERS, CACHE_CONTROL, CONTENT_LENGTH, + CONTENT_TYPE, LAST_MODIFIED, X_CONTENT_TYPE_OPTIONS, +}; +use actix_web::web::Path; +use actix_web::{get, web::Data, HttpRequest, HttpResponse, Responder}; use awc::Client; use base64::DecodeError; use bytes::Bytes; @@ -12,7 +17,7 @@ use serde::Deserialize; use sodiumoxide::crypto::box_::{open_precomputed, Nonce, PrecomputedKey, NONCEBYTES}; use thiserror::Error; -use crate::{client_api_version, RwLockServerState}; +use crate::{client_api_version, CachedImage, RwLockServerState}; const BASE64_CONFIG: base64::Config = base64::Config::new(base64::CharacterSet::UrlSafe, false); @@ -33,8 +38,8 @@ enum ServerResponse { impl Responder for ServerResponse { fn respond_to(self, req: &HttpRequest) -> HttpResponse { match self { - ServerResponse::TokenValidationError(e) => e.respond_to(req), - ServerResponse::HttpResponse(resp) => resp.respond_to(req), + Self::TokenValidationError(e) => e.respond_to(req), + Self::HttpResponse(resp) => resp.respond_to(req), } } } @@ -115,6 +120,12 @@ fn validate_token( token: String, chapter_hash: &str, ) -> Result<(), TokenValidationError> { + #[derive(Deserialize)] + struct Token<'a> { + expires: DateTime, + hash: &'a str, + } + let data = base64::decode_config(token, BASE64_CONFIG)?; if data.len() < NONCEBYTES { return Err(TokenValidationError::IncompleteNonce); @@ -124,12 +135,6 @@ fn validate_token( let decrypted = open_precomputed(&data[NONCEBYTES..], &nonce, precomputed_key) .map_err(|_| TokenValidationError::DecryptionFailure)?; - #[derive(Deserialize)] - struct Token<'a> { - expires: DateTime, - hash: &'a str, - } - let parsed_token: Token = serde_json::from_slice(&decrypted).map_err(|_| TokenValidationError::InvalidToken)?; @@ -146,10 +151,10 @@ fn validate_token( fn push_headers(builder: &mut HttpResponseBuilder) -> &mut HttpResponseBuilder { builder - .insert_header(("X-Content-Type-Options", "nosniff")) - .insert_header(("Access-Control-Allow-Origin", "https://mangadex.org")) - .insert_header(("Access-Control-Expose-Headers", "*")) - .insert_header(("Cache-Control", "public, max-age=1209600")) + .insert_header((X_CONTENT_TYPE_OPTIONS, "nosniff")) + .insert_header((ACCESS_CONTROL_ALLOW_ORIGIN, "https://mangadex.org")) + .insert_header((ACCESS_CONTROL_EXPOSE_HEADERS, "*")) + .insert_header((CACHE_CONTROL, "public, max-age=1209600")) .insert_header(("Timing-Allow-Origin", "https://mangadex.org")) .insert_header(("Server", SERVER_ID_STRING)) } @@ -163,12 +168,7 @@ async fn fetch_image( let key = (chapter_hash, file_name, is_data_saver); if let Some(cached) = state.0.write().cache.get(&key) { - let data = cached.to_vec(); - let data: Vec> = data - .chunks(1024) - .map(|v| Ok(Bytes::from(v.to_vec()))) - .collect(); - return ServerResponse::HttpResponse(HttpResponse::Ok().streaming(stream::iter(data))); + return construct_response(cached); } let mut state = state.0.write(); @@ -184,25 +184,61 @@ async fn fetch_image( .await; match resp { - Ok(mut resp) => match resp.body().await { - Ok(bytes) => { - state.cache.put(key, bytes.to_vec()); - let bytes: Vec> = bytes - .chunks(1024) - .map(|v| Ok(Bytes::from(v.to_vec()))) - .collect(); - return ServerResponse::HttpResponse( - HttpResponse::Ok().streaming(stream::iter(bytes)), - ); + Ok(mut resp) => { + let headers = resp.headers(); + let content_type = headers.get(CONTENT_TYPE).map(AsRef::as_ref).map(Vec::from); + let content_length = headers + .get(CONTENT_LENGTH) + .map(AsRef::as_ref) + .map(Vec::from); + let last_modified = headers.get(LAST_MODIFIED).map(AsRef::as_ref).map(Vec::from); + let body = resp.body().await; + match body { + Ok(bytes) => { + let cached = CachedImage { + data: bytes.to_vec(), + content_type, + content_length, + last_modified, + }; + let resp = construct_response(&cached); + state.cache.put(key, cached); + return resp; + } + Err(e) => { + warn!("Got payload error from image server: {}", e); + ServerResponse::HttpResponse( + push_headers(&mut HttpResponse::ServiceUnavailable()).finish(), + ) + } } - Err(e) => { - warn!("Got payload error from image server: {}", e); - todo!() - } - }, + } Err(e) => { error!("Failed to fetch image from server: {}", e); - todo!() + ServerResponse::HttpResponse( + push_headers(&mut HttpResponse::ServiceUnavailable()).finish(), + ) } } } + +fn construct_response(cached: &CachedImage) -> ServerResponse { + let data: Vec> = cached + .data + .to_vec() + .chunks(1024) + .map(|v| Ok(Bytes::from(v.to_vec()))) + .collect(); + let mut resp = HttpResponse::Ok(); + if let Some(content_type) = &cached.content_type { + resp.append_header((CONTENT_TYPE, &**content_type)); + } + if let Some(content_length) = &cached.content_length { + resp.append_header((CONTENT_LENGTH, &**content_length)); + } + if let Some(last_modified) = &cached.last_modified { + resp.append_header((LAST_MODIFIED, &**last_modified)); + } + + return ServerResponse::HttpResponse(push_headers(&mut resp).streaming(stream::iter(data))); +}