mirror of
https://github.com/sadoyan/aralez.git
synced 2026-04-30 23:08:40 +08:00
some more type changes, performance improvements
This commit is contained in:
@@ -14,6 +14,7 @@ use pingora_core::listeners::ALPN;
|
|||||||
use pingora_core::prelude::HttpPeer;
|
use pingora_core::prelude::HttpPeer;
|
||||||
use pingora_limits::rate::Rate;
|
use pingora_limits::rate::Rate;
|
||||||
use pingora_proxy::{ProxyHttp, Session};
|
use pingora_proxy::{ProxyHttp, Session};
|
||||||
|
// use std::net::{IpAddr, Ipv4Addr};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use tokio::time::Instant;
|
use tokio::time::Instant;
|
||||||
@@ -32,11 +33,12 @@ pub struct LB {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub struct Context {
|
pub struct Context {
|
||||||
backend_id: String,
|
backend_id: Arc<str>,
|
||||||
|
// backend_id: Arc<(IpAddr, u16, bool)>,
|
||||||
to_https: bool,
|
to_https: bool,
|
||||||
redirect_to: String,
|
redirect_to: Arc<str>,
|
||||||
start_time: Instant,
|
start_time: Instant,
|
||||||
hostname: Option<String>,
|
hostname: Option<Arc<str>>,
|
||||||
upstream_peer: Option<InnerMap>,
|
upstream_peer: Option<InnerMap>,
|
||||||
extraparams: arc_swap::Guard<Arc<Extraparams>>,
|
extraparams: arc_swap::Guard<Arc<Extraparams>>,
|
||||||
client_headers: Arc<Vec<(Arc<str>, Arc<str>)>>,
|
client_headers: Arc<Vec<(Arc<str>, Arc<str>)>>,
|
||||||
@@ -47,9 +49,10 @@ impl ProxyHttp for LB {
|
|||||||
type CTX = Context;
|
type CTX = Context;
|
||||||
fn new_ctx(&self) -> Self::CTX {
|
fn new_ctx(&self) -> Self::CTX {
|
||||||
Context {
|
Context {
|
||||||
backend_id: String::new(),
|
backend_id: Arc::from(""),
|
||||||
|
// backend_id: Arc::new((IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0, false)),
|
||||||
to_https: false,
|
to_https: false,
|
||||||
redirect_to: String::new(),
|
redirect_to: Arc::from(""),
|
||||||
start_time: Instant::now(),
|
start_time: Instant::now(),
|
||||||
hostname: None,
|
hostname: None,
|
||||||
upstream_peer: None,
|
upstream_peer: None,
|
||||||
@@ -92,7 +95,7 @@ impl ProxyHttp for LB {
|
|||||||
None => return Ok(false),
|
None => return Ok(false),
|
||||||
Some(host) => {
|
Some(host) => {
|
||||||
// let optioninnermap = self.get_host(host.as_str(), host.as_str(), backend_id);
|
// let optioninnermap = self.get_host(host.as_str(), host.as_str(), backend_id);
|
||||||
let optioninnermap = self.get_host(host.as_str(), session.req_header().uri.path(), backend_id);
|
let optioninnermap = self.get_host(host, session.req_header().uri.path(), backend_id);
|
||||||
match optioninnermap {
|
match optioninnermap {
|
||||||
None => return Ok(false),
|
None => return Ok(false),
|
||||||
Some(ref innermap) => {
|
Some(ref innermap) => {
|
||||||
@@ -120,17 +123,14 @@ impl ProxyHttp for LB {
|
|||||||
}
|
}
|
||||||
async fn upstream_peer(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<Box<HttpPeer>> {
|
async fn upstream_peer(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<Box<HttpPeer>> {
|
||||||
match ctx.hostname.as_ref() {
|
match ctx.hostname.as_ref() {
|
||||||
Some(hostname) => {
|
Some(hostname) => match ctx.upstream_peer.as_ref() {
|
||||||
match ctx.upstream_peer.as_ref() {
|
|
||||||
// Some((address, port, ssl, is_h2, to_https)) => {
|
|
||||||
Some(innermap) => {
|
Some(innermap) => {
|
||||||
let mut peer = Box::new(HttpPeer::new((innermap.address.clone(), innermap.port.clone()), innermap.is_ssl, String::new()));
|
let mut peer = Box::new(HttpPeer::new((innermap.address.clone(), innermap.port.clone()), innermap.is_ssl, String::new()));
|
||||||
// if session.is_http2() {
|
|
||||||
if innermap.is_http2 {
|
if innermap.is_http2 {
|
||||||
peer.options.alpn = ALPN::H2;
|
peer.options.alpn = ALPN::H2;
|
||||||
}
|
}
|
||||||
if innermap.is_ssl {
|
if innermap.is_ssl {
|
||||||
peer.sni = hostname.clone();
|
peer.sni = hostname.to_string();
|
||||||
peer.options.verify_cert = false;
|
peer.options.verify_cert = false;
|
||||||
peer.options.verify_hostname = false;
|
peer.options.verify_hostname = false;
|
||||||
}
|
}
|
||||||
@@ -142,14 +142,13 @@ impl ProxyHttp for LB {
|
|||||||
let uri = session.req_header().uri.path_and_query().map_or("/", |pq| pq.as_str());
|
let uri = session.req_header().uri.path_and_query().map_or("/", |pq| pq.as_str());
|
||||||
let port = self.config.proxy_port_tls.unwrap_or(403);
|
let port = self.config.proxy_port_tls.unwrap_or(403);
|
||||||
ctx.to_https = true;
|
ctx.to_https = true;
|
||||||
ctx.redirect_to = format!("https://{}:{}{}", host, port, uri);
|
ctx.redirect_to = Arc::from(format!("https://{}:{}{}", host, port, uri));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
ctx.backend_id = Arc::from(format!("{}:{}:{}", innermap.address, innermap.port, innermap.is_ssl));
|
||||||
ctx.backend_id = format!("{}:{}:{}", innermap.address.clone(), innermap.port.clone(), innermap.is_ssl);
|
|
||||||
Ok(peer)
|
Ok(peer)
|
||||||
}
|
}
|
||||||
None => {
|
None => {
|
||||||
@@ -164,8 +163,7 @@ impl ProxyHttp for LB {
|
|||||||
context: Option::from(ImmutStr::Static("Upstream not found")),
|
context: Option::from(ImmutStr::Static("Upstream not found")),
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
}
|
},
|
||||||
}
|
|
||||||
None => {
|
None => {
|
||||||
// session.respond_error_with_body(502, Bytes::from("502 Bad Gateway\n")).await.expect("Failed to send error");
|
// session.respond_error_with_body(502, Bytes::from("502 Bad Gateway\n")).await.expect("Failed to send error");
|
||||||
if let Err(e) = session.respond_error_with_body(502, Bytes::from("502 Bad Gateway\n")).await {
|
if let Err(e) = session.respond_error_with_body(502, Bytes::from("502 Bad Gateway\n")).await {
|
||||||
@@ -184,13 +182,13 @@ impl ProxyHttp for LB {
|
|||||||
|
|
||||||
async fn upstream_request_filter(&self, session: &mut Session, upstream_request: &mut RequestHeader, ctx: &mut Self::CTX) -> Result<()> {
|
async fn upstream_request_filter(&self, session: &mut Session, upstream_request: &mut RequestHeader, ctx: &mut Self::CTX) -> Result<()> {
|
||||||
if let Some(hostname) = ctx.hostname.as_ref() {
|
if let Some(hostname) = ctx.hostname.as_ref() {
|
||||||
upstream_request.insert_header("Host", hostname)?;
|
upstream_request.insert_header("Host", hostname.as_ref())?;
|
||||||
}
|
}
|
||||||
if let Some(peer) = ctx.upstream_peer.as_ref() {
|
if let Some(peer) = ctx.upstream_peer.as_ref() {
|
||||||
upstream_request.insert_header("X-Forwarded-For", peer.address.to_string())?;
|
upstream_request.insert_header("X-Forwarded-For", peer.address.to_string())?;
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(headers) = self.get_header(ctx.hostname.as_ref().unwrap_or(&"localhost".to_string()), session.req_header().uri.path()) {
|
if let Some(headers) = self.get_header(ctx.hostname.as_ref().unwrap_or(&Arc::from("localhost")), session.req_header().uri.path()) {
|
||||||
if let Some(server_headers) = headers.server_headers {
|
if let Some(server_headers) = headers.server_headers {
|
||||||
for k in server_headers {
|
for k in server_headers {
|
||||||
upstream_request.insert_header(k.0, k.1)?;
|
upstream_request.insert_header(k.0, k.1)?;
|
||||||
@@ -208,13 +206,13 @@ impl ProxyHttp for LB {
|
|||||||
async fn response_filter(&self, session: &mut Session, _upstream_response: &mut ResponseHeader, ctx: &mut Self::CTX) -> Result<()> {
|
async fn response_filter(&self, session: &mut Session, _upstream_response: &mut ResponseHeader, ctx: &mut Self::CTX) -> Result<()> {
|
||||||
if ctx.extraparams.sticky_sessions {
|
if ctx.extraparams.sticky_sessions {
|
||||||
let backend_id = ctx.backend_id.clone();
|
let backend_id = ctx.backend_id.clone();
|
||||||
if let Some(bid) = self.ump_byid.get(&backend_id) {
|
if let Some(bid) = self.ump_byid.get(&backend_id.to_string()) {
|
||||||
let _ = _upstream_response.insert_header("set-cookie", format!("backend_id={}; Path=/; Max-Age=600; HttpOnly; SameSite=Lax", bid.address));
|
let _ = _upstream_response.insert_header("set-cookie", format!("backend_id={}; Path=/; Max-Age=600; HttpOnly; SameSite=Lax", bid.address));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if ctx.to_https {
|
if ctx.to_https {
|
||||||
let mut redirect_response = ResponseHeader::build(StatusCode::MOVED_PERMANENTLY, None)?;
|
let mut redirect_response = ResponseHeader::build(StatusCode::MOVED_PERMANENTLY, None)?;
|
||||||
redirect_response.insert_header("Location", ctx.redirect_to.clone())?;
|
redirect_response.insert_header("Location", ctx.redirect_to.as_ref())?;
|
||||||
redirect_response.insert_header("Content-Length", "0")?;
|
redirect_response.insert_header("Content-Length", "0")?;
|
||||||
session.write_response_header(Box::new(redirect_response), false).await?;
|
session.write_response_header(Box::new(redirect_response), false).await?;
|
||||||
}
|
}
|
||||||
@@ -240,17 +238,17 @@ impl ProxyHttp for LB {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn return_header_host(session: &Session) -> Option<String> {
|
fn return_header_host(session: &Session) -> Option<Arc<str>> {
|
||||||
if session.is_http2() {
|
if session.is_http2() {
|
||||||
match session.req_header().uri.host() {
|
match session.req_header().uri.host() {
|
||||||
Some(host) => Option::from(host.to_string()),
|
Some(host) => Option::from(Arc::from(host)),
|
||||||
None => None,
|
None => None,
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
match session.req_header().headers.get("host") {
|
match session.req_header().headers.get("host") {
|
||||||
Some(host) => {
|
Some(host) => {
|
||||||
let header_host = host.to_str().unwrap().splitn(2, ':').collect::<Vec<&str>>();
|
let header_host = host.to_str().unwrap().splitn(2, ':').collect::<Vec<&str>>();
|
||||||
Option::from(header_host[0].to_string())
|
Option::from(Arc::from(header_host[0].to_string()))
|
||||||
}
|
}
|
||||||
None => None,
|
None => None,
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user