minimal viable product

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<u8>>,
+    cache: LruCache<(String, String, bool), CachedImage>,
+}
+
+struct CachedImage {
+    data: Vec<u8>,
+    content_type: Option<Vec<u8>>,
+    content_length: Option<Vec<u8>>,
+    last_modified: Option<Vec<u8>>,
 }
 
 impl ServerState {
     async fn init(config: &Config) -> Result<Self, ()> {
         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::<PingResponse>().await {
+            Ok(mut resp) => match resp.json::<Response>().await {
                 Ok(resp) => {
                     let key = resp
                         .token_key
-                        .and_then(|key| match PrecomputedKey::from_slice(key.as_bytes()) {
+                        .and_then(|key| {
-                            Some(key) => Some(key),
+                            if let Some(key) = PrecomputedKey::from_slice(key.as_bytes()) {
-                            None => {
+                                Some(key)
+                            } else {
                                 error!("Failed to parse token key: got {}", key);
                                 None
                             }
@@ -106,7 +115,7 @@ impl ServerState {
     }
 }
 
-struct RwLockServerState(RwLock<ServerState>);
+pub struct RwLockServerState(RwLock<ServerState>);
 
 impl ResolvesServerCert for RwLockServerState {
     fn resolve(&self, _: rustls::ClientHello) -> Option<CertifiedKey> {
@@ -166,7 +175,7 @@ async fn main() -> Result<(), std::io::Error> {
     .await
 }
 
-struct Config {
+pub struct Config {
     secret: String,
     port: u16,
     disk_quota: usize,

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<String>,
 }
 
-impl<'a> PingRequest<'a> {
+impl<'a> Request<'a> {
     fn from_config_and_state(config: &'a Config, state: &'a Arc<RwLockServerState>) -> 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 struct Response {
-    pub(crate) image_server: Url,
+    pub image_server: Url,
-    pub(crate) latest_build: usize,
+    pub latest_build: usize,
-    pub(crate) url: String,
+    pub url: String,
-    pub(crate) token_key: Option<String>,
+    pub token_key: Option<String>,
-    pub(crate) compromised: bool,
+    pub compromised: bool,
-    pub(crate) paused: bool,
+    pub paused: bool,
-    pub(crate) disabled_tokens: bool,
+    pub disabled_tokens: bool,
-    pub(crate) tls: Option<Tls>,
+    pub tls: Option<Tls>,
 }
 
 #[derive(Deserialize)]
-pub(crate) struct Tls {
+pub struct Tls {
     pub created_at: String,
     pub private_key: Vec<u8>,
     pub certificate: Vec<u8>,
 }
 
-pub(crate) async fn update_server_state(req: &Config, data: &mut Arc<RwLockServerState>) {
+pub async fn update_server_state(req: &Config, data: &mut Arc<RwLockServerState>) {
-    let req = PingRequest::from_config_and_state(req, data);
+    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::<PingResponse>().await {
+        Ok(mut resp) => match resp.json::<Response>().await {
             Ok(resp) => {
                 let mut write_guard = data.0.write();
 

src/routes.rs

@@ -1,7 +1,12 @@
 use std::convert::Infallible;
 
-use actix_web::{dev::HttpResponseBuilder, get, web::Data, HttpRequest, Responder};
+use actix_web::dev::HttpResponseBuilder;
-use actix_web::{web::Path, HttpResponse};
+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),
+            Self::TokenValidationError(e) => e.respond_to(req),
-            ServerResponse::HttpResponse(resp) => resp.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<Utc>,
+        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<Utc>,
-        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((X_CONTENT_TYPE_OPTIONS, "nosniff"))
-        .insert_header(("Access-Control-Allow-Origin", "https://mangadex.org"))
+        .insert_header((ACCESS_CONTROL_ALLOW_ORIGIN, "https://mangadex.org"))
-        .insert_header(("Access-Control-Expose-Headers", "*"))
+        .insert_header((ACCESS_CONTROL_EXPOSE_HEADERS, "*"))
-        .insert_header(("Cache-Control", "public, max-age=1209600"))
+        .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();
+        return construct_response(cached);
-        let data: Vec<Result<Bytes, Infallible>> = data
-            .chunks(1024)
-            .map(|v| Ok(Bytes::from(v.to_vec())))
-            .collect();
-        return ServerResponse::HttpResponse(HttpResponse::Ok().streaming(stream::iter(data)));
     }
 
     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(mut resp) => {
-            Ok(bytes) => {
+            let headers = resp.headers();
-                state.cache.put(key, bytes.to_vec());
+            let content_type = headers.get(CONTENT_TYPE).map(AsRef::as_ref).map(Vec::from);
-                let bytes: Vec<Result<Bytes, Infallible>> = bytes
+            let content_length = headers
-                    .chunks(1024)
+                .get(CONTENT_LENGTH)
-                    .map(|v| Ok(Bytes::from(v.to_vec())))
+                .map(AsRef::as_ref)
-                    .collect();
+                .map(Vec::from);
-                return ServerResponse::HttpResponse(
+            let last_modified = headers.get(LAST_MODIFIED).map(AsRef::as_ref).map(Vec::from);
-                    HttpResponse::Ok().streaming(stream::iter(bytes)),
+            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<Result<Bytes, Infallible>> = 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)));
+}