Github/aralez
1
0
Fork 0
mirror of https://github.com/sadoyan/aralez.git synced 2026-04-30 23:08:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improvements and cleanup

Browse Source
This commit is contained in:
Ara Sadoyan
2025-05-09 17:43:03 +02:00
parent 1be09d21fc
commit c2d847f2aa
9 changed files with 69 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/web/bgservice.rs
View File

@@ -10,6 +10,7 @@ use futures::StreamExt;
use log::info;
use pingora_core::server::ShutdownWatch;
use pingora_core::services::background::BackgroundService;
use std::sync::Arc;
#[async_trait]
impl BackgroundService for LB {
@@ -53,19 +54,11 @@ impl BackgroundService for LB {
Some(ss) => {
clone_dashmap_into(&ss.upstreams, &self.ump_full);
clone_dashmap_into(&ss.upstreams, &self.ump_upst);
self.proxyconf.clear();
{
let mut write_guard = self.extraparams.write().await;
write_guard.stickysessions = ss.extraparams.stickysessions;
}
match ss.globals {
Some(globals) => {
for (k,v) in globals {
self.proxyconf.insert(k, v);
}
}
None => {}
}
let current = self.extraparams.load_full();
let mut new = (*current).clone();
new.stickysessions = ss.extraparams.stickysessions;
new.authentication = ss.extraparams.authentication.clone();
self.extraparams.store(Arc::new(new));
self.headers.clear();
for entry in ss.upstreams.iter() {

65
src/web/proxyhttp.rs
View File

@@ -1,8 +1,8 @@
use crate::utils::auth::authenticate;
use crate::utils::structs::{AppConfig, Extraparams, Headers, UpstreamsDashMap, UpstreamsIdMap};
use crate::web::gethosts::GetHost;
use arc_swap::ArcSwap;
use async_trait::async_trait;
use dashmap::DashMap;
use log::{debug, warn};
use pingora::http::RequestHeader;
use pingora::prelude::*;
@@ -10,9 +10,7 @@ use pingora_core::listeners::ALPN;
use pingora_core::prelude::HttpPeer;
use pingora_http::ResponseHeader;
use pingora_proxy::{ProxyHttp, Session};
use std::ops::Deref;
use std::sync::Arc;
use tokio::sync::RwLock;
pub struct LB {
pub ump_upst: Arc<UpstreamsDashMap>,
@@ -20,9 +18,7 @@ pub struct LB {
pub ump_byid: Arc<UpstreamsIdMap>,
pub headers: Arc<Headers>,
pub config: Arc<AppConfig>,
pub local: Arc<(String, u16)>,
pub proxyconf: Arc<DashMap<String, Vec<String>>>,
pub extraparams: Arc<RwLock<Extraparams>>,
pub extraparams: Arc<ArcSwap<Extraparams>>,
}
pub struct Context {
@@ -45,17 +41,14 @@ impl ProxyHttp for LB {
// session.req_header_mut().headers.insert("X-Host-Name", host.to_string().parse().unwrap());
let mut backend_id = None;
{
let read_guard = self.extraparams.read().await;
if read_guard.stickysessions {
if let Some(cookies) = session.req_header().headers.get("cookie") {
if let Ok(cookie_str) = cookies.to_str() {
for cookie in cookie_str.split(';') {
let trimmed = cookie.trim();
if let Some(value) = trimmed.strip_prefix("backend_id=") {
backend_id = Some(value);
break;
}
if self.extraparams.load().stickysessions {
if let Some(cookies) = session.req_header().headers.get("cookie") {
if let Ok(cookie_str) = cookies.to_str() {
for cookie in cookie_str.split(';') {
let trimmed = cookie.trim();
if let Some(value) = trimmed.strip_prefix("backend_id=") {
backend_id = Some(value);
break;
}
}
}
@@ -76,20 +69,18 @@ impl ProxyHttp for LB {
}
None => {
warn!("Upstream not found. Host: {:?}, Path: {}", host, session.req_header().uri);
let peer = Box::new(HttpPeer::new(self.local.deref(), false, String::new()));
Ok(peer)
Ok(return_no_host(&self.config.local_server))
}
}
}
None => {
warn!("Upstream not found. Host: {:?}, Path: {}", host_name, session.req_header().uri);
let peer = Box::new(HttpPeer::new(self.local.deref(), false, String::new()));
Ok(peer)
Ok(return_no_host(&self.config.local_server))
}
}
}
async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<bool> {
if let Some(auth) = self.proxyconf.get("authorization") {
if let Some(auth) = self.extraparams.load().authentication.get("authorization") {
let authenticated = authenticate(&auth.value(), &session);
if !authenticated {
let _ = session.respond_error(401).await;
@@ -97,11 +88,11 @@ impl ProxyHttp for LB {
return Ok(true);
}
};
if session.req_header().uri.path().starts_with("/denied") {
let _ = session.respond_error(403).await;
warn!("Forbidden: {:?}, {}", session.client_addr(), session.req_header().uri.path().to_string());
return Ok(true);
};
// if session.req_header().uri.path().starts_with("/denied") {
// let _ = session.respond_error(403).await;
// warn!("Forbidden: {:?}, {}", session.client_addr(), session.req_header().uri.path().to_string());
// return Ok(true);
// };
Ok(false)
}
async fn upstream_request_filter(&self, _session: &mut Session, _upstream_request: &mut RequestHeader, _ctx: &mut Self::CTX) -> Result<()> {
@@ -128,14 +119,11 @@ impl ProxyHttp for LB {
async fn response_filter(&self, _session: &mut Session, _upstream_response: &mut ResponseHeader, _ctx: &mut Self::CTX) -> Result<()> {
// _upstream_response.insert_header("X-Proxied-From", "Fooooooooooooooo").unwrap();
{
let read_guard = self.extraparams.read().await;
if read_guard.stickysessions {
let backend_id = _ctx.backend_id.clone();
if let Some(bid) = self.ump_byid.get(&backend_id) {
// let _ = _upstream_response.insert_header("set-cookie", format!("backend {}", bid.0));
let _ = _upstream_response.insert_header("set-cookie", format!("backend_id={}; Path=/; Max-Age=600; HttpOnly; SameSite=Lax", bid.0));
}
if self.extraparams.load().stickysessions {
let backend_id = _ctx.backend_id.clone();
if let Some(bid) = self.ump_byid.get(&backend_id) {
// let _ = _upstream_response.insert_header("set-cookie", format!("backend {}", bid.0));
let _ = _upstream_response.insert_header("set-cookie", format!("backend_id={}; Path=/; Max-Age=600; HttpOnly; SameSite=Lax", bid.0));
}
}
@@ -191,3 +179,10 @@ fn return_header_host(session: &Session) -> Option<&str> {
}
}
}
fn return_no_host(inp: &Option<(String, u16)>) -> Box<HttpPeer> {
match inp {
Some(t) => Box::new(HttpPeer::new(t, false, String::new())),
None => Box::new(HttpPeer::new(("0.0.0.0", 0), false, String::new())),
}
}

40
src/web/start.rs
View File

@@ -1,56 +1,38 @@
use crate::utils::structs::{Extraparams, Headers, UpstreamsDashMap, UpstreamsIdMap};
use crate::utils::structs::Extraparams;
use crate::web::proxyhttp::LB;
use arc_swap::ArcSwap;
use dashmap::DashMap;
use log::info;
use pingora_core::prelude::{background_service, Opt};
use pingora_core::server::Server;
use std::env;
use std::sync::Arc;
use tokio::sync::RwLock;
pub fn run() {
let parameters = Some(Opt::parse_args()).unwrap();
let file = parameters.conf.clone().unwrap();
let maincfg = crate::utils::parceyaml::parce_main_config(file.as_str());
// println!("{:?}", maincfg);
let mut local_conf: (String, u16) = ("0.0.0.0".to_string(), 0);
if let Some((ip, port_str)) = maincfg.config_address.split_once(':') {
if let Ok(port) = port_str.parse::<u16>() {
local_conf = (ip.to_string(), port);
}
}
let mut server = Server::new(parameters).unwrap();
server.bootstrap();
let uf: UpstreamsDashMap = DashMap::new();
let ff: UpstreamsDashMap = DashMap::new();
let im: UpstreamsIdMap = DashMap::new();
let hh: Headers = DashMap::new();
let ec: Extraparams = Extraparams { stickysessions: false };
let uf_config = Arc::new(uf);
let ff_config = Arc::new(ff);
let im_config = Arc::new(im);
let hh_config = Arc::new(hh);
let ec_config = Arc::new(RwLock::new(ec));
let uf_config = Arc::new(DashMap::new());
let ff_config = Arc::new(DashMap::new());
let im_config = Arc::new(DashMap::new());
let hh_config = Arc::new(DashMap::new());
let ec_config = Arc::new(ArcSwap::from_pointee(Extraparams {
stickysessions: false,
authentication: DashMap::new(),
}));
let cfg = Arc::new(maincfg);
let local = Arc::new(local_conf);
let proxyconf: DashMap<String, Vec<String>> = Default::default();
let pconf = Arc::new(proxyconf);
let lb = LB {
ump_upst: uf_config.clone(),
ump_full: ff_config.clone(),
ump_byid: im_config.clone(),
config: cfg.clone(),
local: local.clone(),
headers: hh_config.clone(),
proxyconf: pconf.clone(),
extraparams: ec_config.clone(),
};
let bg = LB {
@@ -58,9 +40,7 @@ pub fn run() {
ump_full: ff_config.clone(),
ump_byid: im_config.clone(),
config: cfg.clone(),
local: local.clone(),
headers: hh_config.clone(),
proxyconf: pconf.clone(),
extraparams: ec_config.clone(),
};