Added whitelisted or all services support t Consul

This commit is contained in:
Ara Sadoyan
2025-04-04 17:38:18 +02:00
parent 23dae51e22
commit b2823508b7
6 changed files with 111 additions and 90 deletions

View File

@@ -1,10 +1,7 @@
use crate::utils::tools::{Headers, UpstreamsDashMap};
use futures::channel::mpsc::Sender;
use std::collections::HashMap;
use std::time::Duration;
use crate::utils::parceyaml::load_configuration;
use crate::utils::tools::{clone_dashmap_into, compare_dashmaps, Headers, UpstreamsDashMap};
use dashmap::DashMap;
use futures::channel::mpsc::Sender;
use futures::SinkExt;
use hickory_client::client::{Client, ClientHandle};
use hickory_client::proto::rr::{DNSClass, Name, RecordType};
@@ -14,8 +11,10 @@ use log::info;
use pingora::prelude::sleep;
use rand::Rng;
use serde::Deserialize;
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::atomic::AtomicUsize;
use std::time::Duration;
#[derive(Debug, Deserialize)]
struct Service {
@@ -36,53 +35,77 @@ struct TaggedAddress {
pub async fn start(fp: String, mut toreturn: Sender<(UpstreamsDashMap, Headers)>) {
let config = load_configuration(fp.as_str(), "filepath");
let headers = DashMap::new();
// println!("{:?}", config);
match config {
Some(config) => {
let conf: Vec<&str> = config.2.split_whitespace().collect();
let y = conf.get(0).unwrap();
if y.to_string() != "consul" {
info!("Not running Consul discovery, requested type is: {}", config.2);
if config.typecfg.to_string() != "consul" {
info!("Not running Consul discovery, requested type is: {}", config.typecfg);
return;
}
info!("Consul Discovery is enabled : {}", config.2);
let end = conf.len();
loop {
let num = rand::thread_rng().gen_range(1..end);
sleep(Duration::from_secs(5)).await;
headers.clear();
for (k, v) in config.1.clone() {
headers.insert(k.to_string(), v);
}
let consul = "http://".to_string() + conf.get(num).unwrap();
let upstreams = http_request(consul, "GET");
match upstreams.await {
Some(upstreams) => {
toreturn.send((upstreams, headers.clone())).await.unwrap();
info!("Consul Discovery is enabled : {}", config.typecfg);
let consul = config.consul;
let prev_upstreams = UpstreamsDashMap::new();
match consul {
Some(consul) => {
let servers = consul.servers.unwrap();
info!("Consul Servers => {:?}", servers);
let end = servers.len();
loop {
let num = rand::thread_rng().gen_range(1..end);
headers.clear();
for (k, v) in config.headers.clone() {
headers.insert(k.to_string(), v);
}
let consul_data = servers.get(num).unwrap().to_string();
let upstreams = http_request(consul_data, consul.whitelist.clone());
match upstreams.await {
Some(upstreams) => {
if !compare_dashmaps(&upstreams, &prev_upstreams) {
clone_dashmap_into(&upstreams, &prev_upstreams);
toreturn.send((upstreams, headers.clone())).await.unwrap();
}
}
None => {}
}
sleep(Duration::from_secs(5)).await;
}
None => {}
}
None => {}
}
}
None => {}
}
}
async fn http_request(url: String, method: &str) -> Option<UpstreamsDashMap> {
async fn http_request(url: String, whitelist: Option<Vec<String>>) -> Option<UpstreamsDashMap> {
let client = reqwest::Client::new();
let to = Duration::from_secs(1);
let upstreams = UpstreamsDashMap::new();
let excludes = vec!["consul", "nomad", "nomad-client"];
match method {
"GET" => {
let ss = url.clone() + "/v1/catalog/service";
let response = client.get(ss.clone() + "s").timeout(to).send().await;
match response {
Ok(r) => {
let json = r.json::<HashMap<String, Vec<String>>>().await;
match json {
Ok(_j) => {
for (k, _v) in _j {
let ss = url.clone() + "/v1/catalog/service";
let response = client.get(ss.clone() + "s").timeout(to).send().await;
match response {
Ok(r) => {
let json = r.json::<HashMap<String, Vec<String>>>().await;
match json {
Ok(_j) => {
for (k, _v) in _j {
match whitelist.clone() {
Some(whitelist) => {
if whitelist.iter().any(|i| *i == k) {
let mut pref: String = ss.clone() + "/";
pref.push_str(&k);
let list = get_by_http(pref).await;
match list {
Some(list) => {
upstreams.insert(k.to_string(), list);
}
None => {}
}
}
}
None => {
if !excludes.iter().any(|&i| i == k) {
let mut pref: String = ss.clone() + "/";
pref.push_str(&k);
@@ -95,19 +118,18 @@ async fn http_request(url: String, method: &str) -> Option<UpstreamsDashMap> {
}
}
}
// print_upstreams(&upstreams);
Some(upstreams)
}
Err(_) => None,
}
// print_upstreams(&upstreams);
Some(upstreams)
}
Err(e) => {
println!("Error: {:?}", e);
None
}
Err(_) => None,
}
}
_ => None,
Err(e) => {
println!("Error: {:?}", e);
None
}
}
}
@@ -121,8 +143,8 @@ async fn get_by_http(url: String) -> Option<DashMap<String, (Vec<(String, u16, b
Ok(r) => {
let jason = r.json::<Vec<Service>>().await;
match jason {
Ok(services) => {
for service in services {
Ok(whitelist) => {
for service in whitelist {
let addr = service.tagged_addresses.get("lan_ipv4").unwrap().address.clone();
let prt = service.tagged_addresses.get("lan_ipv4").unwrap().port.clone();
let to_add = (addr, prt, false);

View File

@@ -20,7 +20,7 @@ pub async fn start(fp: String, mut toreturn: Sender<(UpstreamsDashMap, Headers)>
match snd {
Some(snd) => {
toreturn.send((snd.0, snd.1)).await.unwrap();
toreturn.send((snd.upstreams, snd.headers)).await.unwrap();
}
None => {}
}
@@ -53,7 +53,7 @@ pub async fn start(fp: String, mut toreturn: Sender<(UpstreamsDashMap, Headers)>
let snd = load_configuration(file_path, "filepath");
match snd {
Some(snd) => {
toreturn.send((snd.0, snd.1)).await.unwrap();
toreturn.send((snd.upstreams, snd.headers)).await.unwrap();
}
None => {}
}

View File

@@ -8,9 +8,9 @@ use std::fs;
use std::sync::atomic::AtomicUsize;
#[derive(Debug, Serialize, Deserialize)]
struct Consul {
servers: Option<Vec<String>>,
services: Option<Vec<String>>,
pub struct Consul {
pub servers: Option<Vec<String>>,
pub whitelist: Option<Vec<String>>,
}
#[derive(Debug, Serialize, Deserialize)]
struct Config {
@@ -32,18 +32,27 @@ struct PathConfig {
headers: Option<Vec<String>>,
}
// #[derive(Debug, Serialize, Deserialize)]
// pub struct Allconfig {
// pub upstreams: Option<UpstreamsDashMap>,
// pub headers: Option<Headers>,
// pub consul: Option<Consul>,
// pub typecfg: String,
// }
pub struct Configuration {
pub upstreams: UpstreamsDashMap,
pub headers: Headers,
pub consul: Option<Consul>,
pub typecfg: String,
}
// pub fn load_configuration(d: &str, kind: &str) -> Option<(UpstreamsDashMap, Headers, String)> {
pub fn load_configuration(d: &str, kind: &str) -> Option<(UpstreamsDashMap, Headers, String)> {
let upstreamsmap = UpstreamsDashMap::new();
let headersmap = DashMap::new();
pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
let mut toreturn: Configuration = Configuration {
upstreams: Default::default(),
headers: Default::default(),
consul: None,
typecfg: "".to_string(),
};
toreturn.upstreams = UpstreamsDashMap::new();
toreturn.headers = Headers::new();
// let upstreamsmap = UpstreamsDashMap::new();
// let headersmap = DashMap::new();
let mut yaml_data = d.to_string();
match kind {
"filepath" => {
@@ -79,11 +88,11 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<(UpstreamsDashMap, Head
}
}
global_headers.insert("/".to_string(), hl);
headersmap.insert("GLOBAL_HEADERS".to_string(), global_headers);
toreturn.headers.insert("GLOBAL_HEADERS".to_string(), global_headers);
}
match parsed.provider.as_str() {
"file" => {
toreturn.typecfg = "file".to_string();
if let Some(upstream) = parsed.upstreams {
for (hostname, host_config) in upstream {
let path_map = DashMap::new();
@@ -119,33 +128,22 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<(UpstreamsDashMap, Head
}
path_map.insert(path, (server_list, AtomicUsize::new(0)));
}
headersmap.insert(hostname.clone(), header_list);
upstreamsmap.insert(hostname, path_map);
toreturn.headers.insert(hostname.clone(), header_list);
toreturn.upstreams.insert(hostname, path_map);
}
}
Some((upstreamsmap, headersmap, String::from("file")))
Some(toreturn)
}
"consul" => {
toreturn.typecfg = "consul".to_string();
let consul = parsed.consul;
match consul {
Some(consul) => {
// println!("{:?}", consul.services);
if let Some(srv) = consul.servers {
let joined = srv.join(" ");
Some((upstreamsmap, headersmap, String::from("consul ") + &*joined))
} else {
None
}
toreturn.consul = Some(consul);
Some(toreturn)
}
None => None,
}
// if let Some(srv) = parsed.consul?.servers {
// let joined = srv.join(" ");
// Some((upstreamsmap, headersmap, String::from("consul ") + &*joined))
// } else {
// None
// }
// Some((upstreamsmap, headersmap, String::from("consul ")))
}
"kubernetes" => None,
_ => {

View File

@@ -95,7 +95,8 @@ impl BackgroundService for LB {
}
}
}
// print_upstreams(&self.ump_full);
info!("Upstreams list is changed, updating to:");
print_upstreams(&self.ump_full);
}
None => {}
}

View File

@@ -25,7 +25,7 @@ pub async fn run_server(bindaddress: String, mut toreturn: Sender<(UpstreamsDash
match serverlist {
Some(serverlist) => {
let _ = tr.send((serverlist.0, serverlist.1)).await.unwrap();
let _ = tr.send((serverlist.upstreams, serverlist.headers)).await.unwrap();
Response::builder().status(StatusCode::CREATED).body(Body::from("Config, conf file, updated!\n")).unwrap()
}
None => Response::builder()