feat(reverse): add custom not found

src/http/reverse_proxy.rs

@@ -1,21 +1,111 @@
-use std::time::Duration;
+use std::{path::PathBuf, time::Duration};
 
+use anyhow::Context;
 use axum::{
     body::Body,
     extract::{Path, Request},
     response::{IntoResponse, Response},
 };
 use axum_extra::extract::Host;
-use http::{HeaderName, Uri};
+use http::{
+    HeaderName, HeaderValue, StatusCode, Uri,
+    header::{CONTENT_TYPE, ETAG, IF_NONE_MATCH},
+};
+use mime_guess::from_path;
 use reqwest::Client;
+use tokio::fs::File;
+use tokio_util::io::ReaderStream;
 
-use crate::{http::serve::resolve_parent_path, utils::parse_port_from_host};
+use crate::{
+    http::serve::{calculate_etag, resolve_parent_path},
+    utils::parse_port_from_host,
+};
 
 use super::{
     HOSTS,
     error::{RouteError, RouteResult},
 };
 
+macro_rules! custom_not_found {
+    ($host_config:expr, $request:expr) => {{
+        let page = $host_config
+            .not_found_page
+            .as_ref()
+            .ok_or(RouteError::RouteNotFound())?;
+        let root = $host_config
+            .root
+            .as_ref()
+            .ok_or(RouteError::InternalError())?;
+        let path = format!("{}/{}", root, page.page);
+        tracing::debug!("custom not found path: {:?}", path);
+
+        let file = File::open(path.clone())
+            .await
+            .with_context(|| "open file failed")?;
+
+        let etag = calculate_etag(&file, path.as_str()).await?;
+        let mut response = Response::builder();
+        let mut not_modified = false;
+
+        if let Some(if_none_match) = $request.headers().get(IF_NONE_MATCH)
+            && if_none_match
+                .to_str()
+                .with_context(|| "parse if-none-match failed")?
+                == etag
+        {
+            response = response.status(StatusCode::NOT_MODIFIED);
+            not_modified = true;
+        };
+
+        let stream = if not_modified {
+            let empty = File::open(PathBuf::from("/dev/null"))
+                .await
+                .with_context(|| "open /dev/null failed")?;
+            ReaderStream::new(empty)
+        } else {
+            ReaderStream::new(file)
+        };
+
+        let body = Body::from_stream(stream);
+
+        let mime = from_path(path).first_or_octet_stream();
+        response
+            .headers_mut()
+            .with_context(|| "insert header failed")?
+            .insert(
+                CONTENT_TYPE,
+                HeaderValue::from_str(mime.as_ref()).with_context(|| "insert header failed")?,
+            );
+        response
+            .headers_mut()
+            .with_context(|| "insert header failed")?
+            .insert(
+                ETAG,
+                HeaderValue::from_str(&etag).with_context(|| "insert header failed")?,
+            );
+
+        let response = response
+            .body(body)
+            .with_context(|| "Failed to build HTTP response with body")?;
+        Ok(response)
+    }};
+}
+
+/// Handles the reverse proxy logic for incoming requests.
+/// This function:
+/// 1. Extracts the request path, host, and other details.
+/// 2. Resolves the parent path and proxy configuration.
+/// 3. Forwards the request to the configured proxy server.
+/// 4. Returns the response from the proxy server to the client.
+///
+/// # Arguments
+/// * `req_uri` - The URI of the incoming request.
+/// * `path` - Optional path parameter extracted from the request.
+/// * `host` - The host header from the request.
+/// * `req` - The incoming HTTP request.
+///
+/// # Returns
+/// A `RouteResult` containing the response from the proxy server or an error.
 #[axum::debug_handler]
 pub async fn serve(
     req_uri: Uri,
@@ -42,8 +132,7 @@ pub async fn serve(
         .ok_or(RouteError::RouteNotFound())?;
     tracing::debug!("proxy pass: {:?}", proxy_config);
     let Some(ref proxy_pass) = proxy_config.proxy_pass else {
-        // return custom_not_found!(host_route, request).await;
-        return Err(RouteError::RouteNotFound());
+        return custom_not_found!(proxy_config, req);
     };
     let uri = format!("{proxy_pass}{path_query}");
     tracing::debug!("reverse proxy uri: {:?}", &uri);
@@ -99,6 +188,8 @@ pub async fn serve(
     Ok(res)
 }
 
+/// Checks if a given header should be excluded from being forwarded in the reverse proxy.
+/// Headers like "host", "connection", etc., are typically excluded to avoid conflicts or security issues.
 fn is_exclude_header(name: &HeaderName) -> bool {
     matches!(
         name.as_str(),
@@ -113,6 +204,8 @@ fn is_exclude_header(name: &HeaderName) -> bool {
     )
 }
 
+/// Copies headers from one `HeaderMap` to another, excluding headers specified in `is_exclude_header`.
+/// This ensures only relevant headers are forwarded, avoiding conflicts or security issues.
 fn copy_headers(from: &http::HeaderMap, to: &mut http::HeaderMap) {
     for (name, value) in from.iter() {
         if !is_exclude_header(name) {

src/http/serve.rs

@@ -2,18 +2,16 @@ use std::{path::PathBuf, str::FromStr, time::UNIX_EPOCH};
 
 use anyhow::{Context, anyhow};
 use axum::{
+    body::Body,
     extract::{Path, Request},
     response::{IntoResponse, Response},
 };
 use axum_extra::extract::Host;
 use dashmap::mapref::one::Ref;
-use futures_util::StreamExt;
 use http::{
     HeaderValue, StatusCode, Uri,
     header::{CONTENT_TYPE, ETAG, IF_NONE_MATCH},
 };
-use http_body_util::StreamBody;
-use hyper::body::Frame;
 use mime_guess::from_path;
 use tokio::fs::{self, File};
 use tokio_util::io::ReaderStream;
@@ -51,7 +49,7 @@ macro_rules! custom_not_found {
             let path = format!("{}/{}", root, page.page);
             let status = StatusCode::from_str(page.status.to_string().as_ref())
                 .map_err(|_| RouteError::BadRequest())?;
-            debug!("custom not found path: {:?}", path);
+            tracing::debug!("custom not found path: {:?}", path);
             match stream_file(path.into(), $request, Some(status)).await {
                 Ok(res) => RouteResult::Ok(res),
                 Err(e) => {
@@ -234,7 +232,7 @@ async fn stream_file(
         .await
         .with_context(|| "open file failed")?;
 
-    let path_str = path.to_str().ok_or(anyhow!(""))?;
+    let path_str = path.to_str().ok_or(anyhow!("convert path to str failed"))?;
     let etag = calculate_etag(&file, path_str).await?;
 
     let mut response = Response::builder();
@@ -260,8 +258,9 @@ async fn stream_file(
     } else {
         ReaderStream::new(file)
     };
-    let stream = stream.map(|res| res.map(Frame::data));
-    let body = StreamBody::new(stream);
+    // let stream = stream.map(|res| res.map(Frame::data));
+    // let body = StreamBody::new(stream);
+    let body = Body::from_stream(stream);
 
     let mime = from_path(path).first_or_octet_stream();
     response
@@ -281,11 +280,13 @@ async fn stream_file(
     if let Some(status) = status {
         response = response.status(status);
     }
-    let response = response.body(body).with_context(|| "")?;
+    let response = response
+        .body(body)
+        .with_context(|| "Failed to build HTTP response with body")?;
     Ok(response)
 }
 
-async fn calculate_etag(file: &File, path: &str) -> anyhow::Result<String> {
+pub async fn calculate_etag(file: &File, path: &str) -> anyhow::Result<String> {
     // calculate file metadata as etag
     let metadata = file
         .metadata()