From 285de170fb67c72021e3adff852ce0f681906a6e Mon Sep 17 00:00:00 2001 From: vicanso Date: Sat, 23 Mar 2024 14:34:12 +0800 Subject: [PATCH] refactor: adjust http cache mod --- src/cache/http_header.rs | 34 +++++++++++++++++++ .../http_response.rs} | 9 ++--- src/cache/mod.rs | 5 +++ src/main.rs | 1 + src/proxy/location.rs | 20 +++++------ src/proxy/mod.rs | 1 - src/proxy/server.rs | 8 ++--- src/utils/mod.rs | 32 ----------------- 8 files changed, 59 insertions(+), 51 deletions(-) create mode 100644 src/cache/http_header.rs rename src/{proxy/response_data.rs => cache/http_response.rs} (84%) create mode 100644 src/cache/mod.rs diff --git a/src/cache/http_header.rs b/src/cache/http_header.rs new file mode 100644 index 0000000..8b22efd --- /dev/null +++ b/src/cache/http_header.rs @@ -0,0 +1,34 @@ +use crate::utils::split_to_two_trim; +use http::{HeaderName, HeaderValue}; +use snafu::{ResultExt, Snafu}; +use std::str::FromStr; + +#[derive(Debug, Snafu)] +pub enum Error { + InvalidHeaderValue { + value: String, + source: http::header::InvalidHeaderValue, + }, + #[snafu(display("Invalid header name {source}, {value}"))] + InvalidHeaderName { + value: String, + source: http::header::InvalidHeaderName, + }, +} +type Result = std::result::Result; + +pub type HttpHeader = (HeaderName, HeaderValue); + +pub fn convert_headers(header_values: &[String]) -> Result> { + let mut arr = vec![]; + for item in header_values { + if let Some([k, v]) = split_to_two_trim(item, ":") { + let name = + HeaderName::from_str(&k).context(InvalidHeaderNameSnafu { value: k.clone() })?; + let value = + HeaderValue::from_str(&v).context(InvalidHeaderValueSnafu { value: v.clone() })?; + arr.push((name, value)); + } + } + Ok(arr) +} diff --git a/src/proxy/response_data.rs b/src/cache/http_response.rs similarity index 84% rename from src/proxy/response_data.rs rename to src/cache/http_response.rs index 321ec7a..312c335 100644 --- a/src/proxy/response_data.rs +++ b/src/cache/http_response.rs @@ -1,17 +1,18 @@ +use super::HttpHeader; use bytes::Bytes; -use http::{HeaderName, HeaderValue, StatusCode}; +use http::StatusCode; use pingora::{http::ResponseHeader, proxy::Session}; #[derive(Default)] -pub struct ResponseData { +pub struct HttpResponse { pub status: StatusCode, pub body: Bytes, pub max_age: Option, pub created_at: Option, - pub headers: Option>, + pub headers: Option>, } -impl ResponseData { +impl HttpResponse { pub async fn send(&self, session: &mut Session) -> pingora::Result { let mut resp = ResponseHeader::build(self.status, Some(4))?; resp.insert_header(http::header::CONTENT_LENGTH, self.body.len().to_string())?; diff --git a/src/cache/mod.rs b/src/cache/mod.rs new file mode 100644 index 0000000..7598b28 --- /dev/null +++ b/src/cache/mod.rs @@ -0,0 +1,5 @@ +mod http_header; +mod http_response; + +pub use http_header::*; +pub use http_response::*; diff --git a/src/main.rs b/src/main.rs index 69155f4..8754dde 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ use std::error::Error; use std::io::Write; use std::sync::Arc; +mod cache; mod config; mod proxy; mod utils; diff --git a/src/proxy/location.rs b/src/proxy/location.rs index 5bd3784..b9adb91 100644 --- a/src/proxy/location.rs +++ b/src/proxy/location.rs @@ -1,6 +1,6 @@ use super::Upstream; -use crate::{config::LocationConf, utils}; -use http::{HeaderName, HeaderValue}; +use crate::cache::{convert_headers, HttpHeader}; +use crate::config::LocationConf; use regex::Regex; use snafu::{ResultExt, Snafu}; use std::sync::Arc; @@ -58,14 +58,14 @@ pub struct Location { path_selector: PathSelector, host: String, reg_rewrite: Option<(Regex, String)>, - headers: Option>, - proxy_headers: Option>, + headers: Option>, + proxy_headers: Option>, pub upstream: Arc, } -fn convert_headers(values: &Option>) -> Result>> { +fn format_headers(values: &Option>) -> Result>> { if let Some(header_values) = values { - let arr = utils::convert_headers(header_values).map_err(|err| Error::Invalid { + let arr = convert_headers(header_values).map_err(|err| Error::Invalid { message: err.to_string(), })?; Ok(Some(arr)) @@ -103,8 +103,8 @@ impl Location { host: conf.host.clone().unwrap_or_default(), upstream: up.clone(), reg_rewrite, - headers: convert_headers(&conf.headers)?, - proxy_headers: convert_headers(&conf.proxy_headers)?, + headers: format_headers(&conf.headers)?, + proxy_headers: format_headers(&conf.proxy_headers)?, }) } #[inline] @@ -134,11 +134,11 @@ impl Location { None } #[inline] - pub fn get_proxy_headers(&self) -> Option> { + pub fn get_proxy_headers(&self) -> Option> { self.proxy_headers.clone() } #[inline] - pub fn get_header(&self) -> Option> { + pub fn get_header(&self) -> Option> { self.headers.clone() } } diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 443f6aa..a058826 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -1,6 +1,5 @@ mod location; mod logger; -mod response_data; mod server; mod state; mod upstream; diff --git a/src/proxy/server.rs b/src/proxy/server.rs index b5e480f..7d6a520 100644 --- a/src/proxy/server.rs +++ b/src/proxy/server.rs @@ -1,7 +1,7 @@ use super::logger::Parser; -use super::response_data::ResponseData; use super::state::State; use super::{Location, Upstream}; +use crate::cache::{convert_headers, HttpResponse}; use crate::config::{LocationConf, PingapConf, UpstreamConf}; use crate::utils; use async_trait::async_trait; @@ -272,20 +272,20 @@ impl Server { } Ok(ServerServices { lb, bg_services }) } - fn get_stats_response(&self) -> ResponseData { + fn get_stats_response(&self) -> HttpResponse { let buf = serde_json::to_vec(&ServerStats { accepted: self.accepted.load(Ordering::Relaxed), processing: self.processing.load(Ordering::Relaxed), hostname: HOST_NAME.to_string(), }) .unwrap_or_default(); - let headers = utils::convert_headers(&[ + let headers = convert_headers(&[ "Content-Type: application/json; charset=utf-8".to_string(), "Cache-Control: private, no-store".to_string(), ]) .unwrap_or_default(); - ResponseData { + HttpResponse { status: StatusCode::OK, body: Bytes::from(buf), headers: Some(headers), diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 17db13c..00c3398 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,21 +1,3 @@ -use http::{HeaderName, HeaderValue}; -use snafu::{ResultExt, Snafu}; -use std::str::FromStr; - -#[derive(Debug, Snafu)] -pub enum Error { - InvalidHeaderValue { - value: String, - source: http::header::InvalidHeaderValue, - }, - #[snafu(display("Invalid header name {source}, {value}"))] - InvalidHeaderName { - value: String, - source: http::header::InvalidHeaderName, - }, -} -type Result = std::result::Result; - pub fn split_to_two_trim(value: &str, pat: &str) -> Option<[String; 2]> { let arr: Vec<&str> = value.split(pat).collect(); if arr.len() < 2 { @@ -26,20 +8,6 @@ pub fn split_to_two_trim(value: &str, pat: &str) -> Option<[String; 2]> { Some([arr[0].trim().to_string(), value]) } -pub fn convert_headers(header_values: &[String]) -> Result> { - let mut arr = vec![]; - for item in header_values { - if let Some([k, v]) = split_to_two_trim(item, ":") { - let name = - HeaderName::from_str(&k).context(InvalidHeaderNameSnafu { value: k.clone() })?; - let value = - HeaderValue::from_str(&v).context(InvalidHeaderValueSnafu { value: v.clone() })?; - arr.push((name, value)); - } - } - Ok(arr) -} - const NAME: &str = env!("CARGO_PKG_NAME"); const VERSION: &str = env!("CARGO_PKG_VERSION");