tuxmain/rust-rpxy
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor

Browse source
This commit is contained in:
Jun Kurihara 2023-08-07 22:27:18 +09:00
commit 265cc025b0
No known key found for this signature in database
GPG key ID: D992B3E3DE1DED23
2 changed files with 13 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

14
rpxy-lib/src/handler/utils_headers.rs
View file

@ -23,7 +23,7 @@ pub(super) fn takeout_sticky_cookie_lb_context(
) -> Result<Option<LbContext>> { ) -> Result<Option<LbContext>> {
let mut headers_clone = headers.clone(); let mut headers_clone = headers.clone();
match headers_clone.entry(hyper::header::COOKIE) { match headers_clone.entry(header::COOKIE) {
header::Entry::Vacant(_) => Ok(None), header::Entry::Vacant(_) => Ok(None),
header::Entry::Occupied(entry) => { header::Entry::Occupied(entry) => {
let cookies_iter = entry let cookies_iter = entry
@ -43,8 +43,8 @@ pub(super) fn takeout_sticky_cookie_lb_context(
} }
let cookies_passed_to_upstream = without_sticky_cookies.join("; "); let cookies_passed_to_upstream = without_sticky_cookies.join("; ");
let cookie_passed_to_lb = sticky_cookies.first().unwrap(); let cookie_passed_to_lb = sticky_cookies.first().unwrap();
headers.remove(hyper::header::COOKIE); headers.remove(header::COOKIE);
headers.insert(hyper::header::COOKIE, cookies_passed_to_upstream.parse()?); headers.insert(header::COOKIE, cookies_passed_to_upstream.parse()?);
let sticky_cookie = StickyCookie { let sticky_cookie = StickyCookie {
value: StickyCookieValue::try_from(cookie_passed_to_lb, expected_cookie_name)?, value: StickyCookieValue::try_from(cookie_passed_to_lb, expected_cookie_name)?,
@ -63,7 +63,7 @@ pub(super) fn set_sticky_cookie_lb_context(headers: &mut HeaderMap, context_from
let sticky_cookie_string: String = context_from_lb.sticky_cookie.clone().try_into()?; let sticky_cookie_string: String = context_from_lb.sticky_cookie.clone().try_into()?;
let new_header_val: HeaderValue = sticky_cookie_string.parse()?; let new_header_val: HeaderValue = sticky_cookie_string.parse()?;
let expected_cookie_name = &context_from_lb.sticky_cookie.value.name; let expected_cookie_name = &context_from_lb.sticky_cookie.value.name;
match headers.entry(hyper::header::SET_COOKIE) { match headers.entry(header::SET_COOKIE) {
header::Entry::Vacant(entry) => { header::Entry::Vacant(entry) => {
entry.insert(new_header_val); entry.insert(new_header_val);
} }
@ -173,13 +173,13 @@ pub(super) fn add_header_entry_overwrite_if_exist(
pub(super) fn make_cookie_single_line(headers: &mut HeaderMap) -> Result<()> { pub(super) fn make_cookie_single_line(headers: &mut HeaderMap) -> Result<()> {
let cookies = headers let cookies = headers
.iter() .iter()
.filter(|(k, _)| **k == hyper::header::COOKIE) .filter(|(k, _)| **k == header::COOKIE)
.map(|(_, v)| v.to_str().unwrap_or("")) .map(|(_, v)| v.to_str().unwrap_or(""))
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join("; "); .join("; ");
if !cookies.is_empty() { if !cookies.is_empty() {
headers.remove(hyper::header::COOKIE); headers.remove(header::COOKIE);
headers.insert(hyper::header::COOKIE, HeaderValue::from_bytes(cookies.as_bytes())?); headers.insert(header::COOKIE, HeaderValue::from_bytes(cookies.as_bytes())?);
} }
Ok(()) Ok(())
} }

11
rpxy-lib/src/log.rs
View file

@ -1,4 +1,5 @@
use crate::utils::ToCanonical; use crate::utils::ToCanonical;
use hyper::header;
use std::net::SocketAddr; use std::net::SocketAddr;
pub use tracing::{debug, error, info, warn}; pub use tracing::{debug, error, info, warn};
@ -20,7 +21,7 @@ pub struct MessageLog {
impl<T> From<&hyper::Request<T>> for MessageLog { impl<T> From<&hyper::Request<T>> for MessageLog {
fn from(req: &hyper::Request<T>) -> Self { fn from(req: &hyper::Request<T>) -> Self {
let header_mapper = |v: hyper::header::HeaderName| { let header_mapper = |v: header::HeaderName| {
req req
.headers() .headers()
.get(v) .get(v)
@ -31,7 +32,7 @@ impl<T> From<&hyper::Request<T>> for MessageLog {
// tls_server_name: "".to_string(), // tls_server_name: "".to_string(),
client_addr: "".to_string(), client_addr: "".to_string(),
method: req.method().to_string(), method: req.method().to_string(),
host: header_mapper(hyper::header::HOST), host: header_mapper(header::HOST),
p_and_q: req p_and_q: req
.uri() .uri()
.path_and_query() .path_and_query()
@ -40,8 +41,8 @@ impl<T> From<&hyper::Request<T>> for MessageLog {
version: req.version(), version: req.version(),
uri_scheme: req.uri().scheme_str().unwrap_or("").to_string(), uri_scheme: req.uri().scheme_str().unwrap_or("").to_string(),
uri_host: req.uri().host().unwrap_or("").to_string(), uri_host: req.uri().host().unwrap_or("").to_string(),
ua: header_mapper(hyper::header::USER_AGENT), ua: header_mapper(header::USER_AGENT),
xff: header_mapper(hyper::header::HeaderName::from_static("x-forwarded-for")), xff: header_mapper(header::HeaderName::from_static("x-forwarded-for")),
status: "".to_string(), status: "".to_string(),
upstream: "".to_string(), upstream: "".to_string(),
} }
@ -61,7 +62,7 @@ impl MessageLog {
self.status = status_code.to_string(); self.status = status_code.to_string();
self self
} }
pub fn xff(&mut self, xff: &Option<&hyper::header::HeaderValue>) -> &mut Self { pub fn xff(&mut self, xff: &Option<&header::HeaderValue>) -> &mut Self {
self.xff = xff.map_or_else(|| "", |v| v.to_str().unwrap_or("")).to_string(); self.xff = xff.map_or_else(|| "", |v| v.to_str().unwrap_or("")).to_string();
self self
} }