From 16bb3d33c841b0a0112beb980886888b3423f0be Mon Sep 17 00:00:00 2001 From: xfy Date: Tue, 17 Jun 2025 17:23:34 +0800 Subject: [PATCH] feat(reverse): add custom not found --- src/http/reverse_proxy.rs | 103 ++++++++++++++++++++++++++++++++++++-- src/http/serve.rs | 19 +++---- 2 files changed, 108 insertions(+), 14 deletions(-) diff --git a/src/http/reverse_proxy.rs b/src/http/reverse_proxy.rs index e0e5f8a..dcd0cd1 100644 --- a/src/http/reverse_proxy.rs +++ b/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) { diff --git a/src/http/serve.rs b/src/http/serve.rs index bf70671..13fd5d5 100644 --- a/src/http/serve.rs +++ b/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 { +pub async fn calculate_etag(file: &File, path: &str) -> anyhow::Result { // calculate file metadata as etag let metadata = file .metadata()