mirror of
https://github.com/sadoyan/aralez.git
synced 2026-10-23 06:01:35 +08:00
367 lines
12 KiB
Rust
367 lines
12 KiB
Rust
use crate::tls::load;
|
|
use crate::tls::load::CertificateConfig;
|
|
use crate::utils::types::{Extraparams, InnerMapForJson, UpstreamSnapshotForJson, UpstreamsDashMap, UpstreamsIdMap};
|
|
use dashmap::DashMap;
|
|
use log::{error, info};
|
|
use privdrop::PrivDrop;
|
|
use serde_json::{Value, json};
|
|
use sha2::{Digest, Sha256};
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::fmt::Write;
|
|
use std::fs::OpenOptions;
|
|
use std::io::Write as IoWrite;
|
|
use std::net::SocketAddr;
|
|
use std::net::TcpListener;
|
|
use std::os::unix::fs::MetadataExt;
|
|
use std::os::unix::fs::OpenOptionsExt;
|
|
use std::process::Command;
|
|
use std::str::FromStr;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::{env, fs, process, thread, time};
|
|
|
|
pub fn print_upstreams(upstreams: &UpstreamsDashMap, extraparams: &Extraparams) {
|
|
let mut out = String::new();
|
|
for host_entry in upstreams.iter() {
|
|
writeln!(out, "Hostname: {}", host_entry.key()).unwrap();
|
|
for path_entry in host_entry.value().iter() {
|
|
writeln!(out, " Path: {}", path_entry.key()).unwrap();
|
|
for f in path_entry.value().0.clone() {
|
|
writeln!(
|
|
out,
|
|
" IP: {}, Port: {}, SSL: {}, H2: {}, To HTTPS: {}, Rate Limit: {}, 4xx Limit: {}",
|
|
f.address,
|
|
f.port,
|
|
f.is_ssl,
|
|
f.is_http2,
|
|
f.to_https,
|
|
if f.rate_limit > 0 { f.rate_limit } else { extraparams.rate_limit },
|
|
if f.x4xx_limit > 0 { f.x4xx_limit } else { extraparams.x4xx_limit }
|
|
)
|
|
.unwrap();
|
|
}
|
|
}
|
|
}
|
|
info!("\n{}", out.trim_end());
|
|
}
|
|
pub fn clone_dashmap(original: &UpstreamsDashMap) -> UpstreamsDashMap {
|
|
let new_map: UpstreamsDashMap = DashMap::new();
|
|
|
|
for outer_entry in original.iter() {
|
|
let hostname = outer_entry.key();
|
|
let inner_map = outer_entry.value();
|
|
|
|
let new_inner_map = DashMap::new();
|
|
|
|
for inner_entry in inner_map.iter() {
|
|
let path = inner_entry.key();
|
|
let (vec, _) = inner_entry.value();
|
|
let new_vec = vec.clone();
|
|
let new_counter = AtomicUsize::new(0);
|
|
new_inner_map.insert(path.clone(), (new_vec, new_counter));
|
|
}
|
|
new_map.insert(hostname.clone(), new_inner_map);
|
|
}
|
|
new_map
|
|
}
|
|
|
|
pub fn clone_dashmap_into(original: &UpstreamsDashMap, cloned: &UpstreamsDashMap) {
|
|
cloned.clear();
|
|
for outer_entry in original.iter() {
|
|
let hostname = outer_entry.key();
|
|
let inner_map = outer_entry.value();
|
|
let new_inner_map = DashMap::new();
|
|
for inner_entry in inner_map.iter() {
|
|
let path = inner_entry.key();
|
|
let (vec, _) = inner_entry.value();
|
|
let new_vec = vec.clone();
|
|
let new_counter = AtomicUsize::new(0);
|
|
new_inner_map.insert(path.clone(), (new_vec, new_counter));
|
|
}
|
|
cloned.insert(hostname.clone(), new_inner_map);
|
|
}
|
|
}
|
|
|
|
pub fn compare_dashmaps(map1: &UpstreamsDashMap, map2: &UpstreamsDashMap) -> bool {
|
|
if map1.len() != map2.len() {
|
|
return false;
|
|
}
|
|
for entry1 in map1.iter() {
|
|
let Some(inner_map2) = map2.get(entry1.key()) else {
|
|
return false;
|
|
};
|
|
let inner_map1 = entry1.value();
|
|
if inner_map1.len() != inner_map2.len() {
|
|
return false;
|
|
}
|
|
for path_entry in inner_map1.iter() {
|
|
let Some(entry2) = inner_map2.get(path_entry.key()) else {
|
|
return false;
|
|
};
|
|
let (vec1, _) = path_entry.value();
|
|
let (vec2, _) = entry2.value();
|
|
if vec1.len() != vec2.len() {
|
|
return false;
|
|
}
|
|
let set1: HashSet<_> = vec1.iter().collect();
|
|
let set2: HashSet<_> = vec2.iter().collect();
|
|
if set1 != set2 {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
true
|
|
}
|
|
|
|
pub fn merge_headers(target: &DashMap<Arc<str>, Vec<(String, Arc<str>)>>, source: &DashMap<Arc<str>, Vec<(String, Arc<str>)>>) {
|
|
for entry in source.iter() {
|
|
let global_key = entry.key().clone();
|
|
let global_values = entry.value().clone();
|
|
let mut target_entry = target.entry(global_key).or_default();
|
|
target_entry.extend(global_values);
|
|
}
|
|
}
|
|
|
|
pub fn clone_idmap_into(original: &UpstreamsDashMap, cloned: &UpstreamsIdMap) {
|
|
cloned.clear();
|
|
for outer_entry in original.iter() {
|
|
let inner_map = outer_entry.value();
|
|
for inner_entry in inner_map.iter() {
|
|
let (vec, _) = inner_entry.value();
|
|
for x in vec.iter() {
|
|
let mut id = String::new();
|
|
write!(
|
|
&mut id,
|
|
"{}:{}:{}:{}:{}:{}:{}:{}:{:?}",
|
|
outer_entry.key(),
|
|
x.address,
|
|
x.port,
|
|
x.is_http2,
|
|
x.to_https,
|
|
x.x4xx_limit,
|
|
x.rate_limit,
|
|
x.healthcheck,
|
|
x.authorization
|
|
)
|
|
.unwrap_or(());
|
|
let mut hasher = Sha256::new();
|
|
// address: "127.0.0.3", port: 8000, is_ssl: false, is_http2: false, to_https: false, rate_limit: Some(200), healthcheck: None, authorization: None } }
|
|
|
|
hasher.update(id.clone().into_bytes());
|
|
let hash = hasher.finalize();
|
|
let hex_hash = base16ct::lower::encode_string(&hash);
|
|
let hh = hex_hash[0..50].to_string();
|
|
cloned.insert(hh, x.to_owned());
|
|
}
|
|
}
|
|
}
|
|
info!("Upstreams are fully populated. Ready to server requests");
|
|
}
|
|
|
|
pub fn listdir(dir: String) -> Vec<load::CertificateConfig> {
|
|
let mut f = HashMap::new();
|
|
let mut certificate_configs: Vec<load::CertificateConfig> = vec![];
|
|
let paths = fs::read_dir(dir).unwrap();
|
|
for path in paths {
|
|
let path_str = path.unwrap().path().to_str().unwrap().to_owned();
|
|
if path_str.ends_with(".crt") {
|
|
let name = path_str.replace(".crt", "");
|
|
let mut inner = vec![];
|
|
let domain = name.split("/").collect::<Vec<&str>>();
|
|
inner.push(name.clone() + ".crt");
|
|
inner.push(name.clone() + ".key");
|
|
f.insert(domain[domain.len() - 1].to_owned(), inner);
|
|
let y = CertificateConfig {
|
|
cert_path: name.clone() + ".crt",
|
|
key_path: name.clone() + ".key",
|
|
};
|
|
certificate_configs.push(y);
|
|
}
|
|
}
|
|
// for (_, v) in f.iter() {
|
|
// let y = CertificateConfig {
|
|
// cert_path: v[0].clone(),
|
|
// key_path: v[1].clone(),
|
|
// };
|
|
// certificate_configs.push(y);
|
|
// }
|
|
certificate_configs
|
|
}
|
|
|
|
pub fn drop_priv(user: String, group: String, http_addr: String, tls_addr: Option<String>) {
|
|
thread::sleep(time::Duration::from_millis(10));
|
|
loop {
|
|
thread::sleep(time::Duration::from_millis(10));
|
|
if TcpListener::bind(&http_addr).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
if let Some(tls_addr) = tls_addr {
|
|
loop {
|
|
thread::sleep(time::Duration::from_millis(10));
|
|
if TcpListener::bind(&tls_addr).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
info!("Dropping ROOT privileges to: {}:{}", user, group);
|
|
if let Err(e) = PrivDrop::default().user(user).group(group).apply() {
|
|
error!("Failed to drop privileges: {}", e);
|
|
process::exit(1)
|
|
}
|
|
}
|
|
|
|
pub fn check_priv(addr: &str) {
|
|
let port = SocketAddr::from_str(addr).map(|sa| sa.port()).expect("Failed to parse address port ");
|
|
if port < 1024 {
|
|
let meta = std::fs::metadata("/proc/self").map(|m| m.uid()).unwrap();
|
|
if meta != 0 {
|
|
error!("Running on privileged port requires to start as ROOT");
|
|
process::exit(1)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub fn upstreams_to_json(upstreams: &UpstreamsDashMap) -> serde_json::Result<String> {
|
|
let mut outer = HashMap::new();
|
|
|
|
for outer_entry in upstreams.iter() {
|
|
let mut inner_map = HashMap::new();
|
|
|
|
for inner_entry in outer_entry.value().iter() {
|
|
let (backends, counter) = inner_entry.value();
|
|
|
|
inner_map.insert(
|
|
inner_entry.key().to_string(),
|
|
UpstreamSnapshotForJson {
|
|
backends: backends
|
|
.iter()
|
|
.map(|a| InnerMapForJson {
|
|
address: a.address.to_string(),
|
|
port: a.port,
|
|
is_ssl: a.is_ssl,
|
|
is_http2: a.is_http2,
|
|
to_https: a.to_https,
|
|
rate_limit: a.rate_limit,
|
|
x4xx_limit: a.x4xx_limit,
|
|
healthcheck: a.healthcheck,
|
|
})
|
|
.collect(),
|
|
requests: counter.load(Ordering::Relaxed),
|
|
},
|
|
);
|
|
}
|
|
|
|
outer.insert(outer_entry.key().to_string(), inner_map);
|
|
}
|
|
|
|
// serde_json::to_string_pretty(&outer)
|
|
serde_json::to_string(&outer)
|
|
}
|
|
|
|
pub fn upstreams_liveness_json(configured: &UpstreamsDashMap, current: &UpstreamsDashMap) -> Value {
|
|
let mut result = serde_json::Map::new();
|
|
|
|
for host_entry in configured.iter() {
|
|
let hostname = host_entry.key().to_string();
|
|
let configured_paths = host_entry.value();
|
|
|
|
let mut paths_json = serde_json::Map::new();
|
|
|
|
for path_entry in configured_paths.iter() {
|
|
let path = path_entry.key().clone();
|
|
let (configured_backends, _) = path_entry.value();
|
|
let backends_json: Vec<Value> = configured_backends
|
|
.iter()
|
|
.map(|backend| {
|
|
let alive = if let Some(host_map) = current.get(&*hostname) {
|
|
if let Some(path_entry) = host_map.get(&*path) {
|
|
let list = &path_entry.value().0; // Vec<Arc<InnerMap>>
|
|
list.iter().any(|b| b.address == backend.address && b.port == backend.port)
|
|
} else {
|
|
false
|
|
}
|
|
} else {
|
|
false
|
|
};
|
|
json!({
|
|
"address": &*backend.address,
|
|
"port": backend.port,
|
|
"alive": alive
|
|
})
|
|
})
|
|
.collect();
|
|
|
|
paths_json.insert(
|
|
path.to_string(),
|
|
json!({
|
|
"backends": backends_json
|
|
}),
|
|
);
|
|
}
|
|
|
|
result.insert(hostname, Value::Object(paths_json));
|
|
}
|
|
Value::Object(result)
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub fn prepend(prefix: &str, val: &Option<Arc<str>>, uri: &str, port: &str) -> Option<String> {
|
|
val.as_ref().map(|s| {
|
|
let mut buf = String::with_capacity(32);
|
|
buf.push_str(prefix);
|
|
buf.push_str(s);
|
|
buf.push(':');
|
|
buf.push_str(port);
|
|
buf.push_str(uri);
|
|
buf
|
|
})
|
|
}
|
|
|
|
pub fn write_pid_file(path: &str) -> std::io::Result<()> {
|
|
let mut file = OpenOptions::new()
|
|
.write(true)
|
|
.create(true)
|
|
.truncate(true)
|
|
.custom_flags(libc::O_NOFOLLOW) // refuse to follow symlinks
|
|
.open(path)?;
|
|
file.write_all(process::id().to_string().as_bytes())?;
|
|
Ok(())
|
|
}
|
|
|
|
#[allow(clippy::needless_return)]
|
|
pub fn split_host_port(addr: &str, tls: bool) -> Option<(&str, u16, bool, &str)> {
|
|
match addr.split_once(':') {
|
|
Some((h, p)) => match p.parse::<u16>() {
|
|
Ok(port) => return Some((h, port, tls, h)),
|
|
Err(_) => {
|
|
log::warn!("ForwardAuth: invalid port in {}", addr);
|
|
return None;
|
|
}
|
|
},
|
|
None => {
|
|
if tls {
|
|
return Some((addr, 443u16, tls, addr));
|
|
} else {
|
|
return Some((addr, 80u16, tls, addr));
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
pub fn get_hostname() -> String {
|
|
if let Ok(host) = env::var("HOSTNAME").or_else(|_| env::var("COMPUTERNAME")) {
|
|
if !host.trim().is_empty() {
|
|
return host;
|
|
}
|
|
}
|
|
Command::new("hostname")
|
|
.output()
|
|
.ok()
|
|
.and_then(|output| String::from_utf8(output.stdout).ok())
|
|
.map(|s| s.trim().to_string())
|
|
.filter(|s| !s.is_empty())
|
|
.unwrap_or_else(|| "unknown".to_string())
|
|
}
|