use async_trait::async_trait;
use pingora::prelude::*;
use pingora::protocols::l4::socket::SocketAddr as PingoraSocketAddr;
use pingora_core::server::Server;
use pingora_core::upstreams::peer::HttpPeer;
use pingora_http::{RequestHeader, ResponseHeader};
use pingora_load_balancing::{
Backend, Backends, LoadBalancer, health_check, selection::RoundRobin,
};
use pingora_proxy::{ProxyHttp, Session, http_proxy_service};
use revoke_core::ServiceRegistry;
use std::collections::{BTreeSet, HashMap};
use std::sync::{Arc, RwLock};
use structopt::StructOpt;
use tracing::{error, info};
pub struct RevokeGateway {
server: Server,
service_registry: Arc<dyn ServiceRegistry>,
}
pub struct GatewayProxyService {
service_registry: Arc<dyn ServiceRegistry>,
load_balancers: Arc<RwLock<HashMap<String, Arc<LoadBalancer<RoundRobin>>>>>,
config: Arc<GatewayConfig>,
}
#[async_trait]
impl ProxyHttp for GatewayProxyService {
type CTX = ();
fn new_ctx(&self) -> Self::CTX {}
async fn upstream_peer(
&self,
session: &mut Session,
_ctx: &mut Self::CTX,
) -> Result<Box<HttpPeer>> {
let path = session.req_header().uri.path();
let service_name = extract_service_name(path);
let load_balancer = self.get_or_create_load_balancer(&service_name).await?;
let backend = load_balancer.select(b"", 256).ok_or_else(|| {
error!("No healthy backend available for service: {}", service_name);
Error::new(ErrorType::HTTPStatus(503))
})?;
let peer = Box::new(HttpPeer::new(
backend.addr.to_string(),
false,
service_name.clone(),
));
Ok(peer)
}
async fn upstream_request_filter(
&self,
session: &mut Session,
upstream_request: &mut RequestHeader,
_ctx: &mut Self::CTX,
) -> Result<()> {
upstream_request.insert_header("X-Gateway", "revoke-gateway")?;
if let Some(client_addr) = session.client_addr() {
upstream_request.insert_header("X-Forwarded-For", client_addr.to_string())?;
}
Ok(())
}
async fn response_filter(
&self,
_session: &mut Session,
upstream_response: &mut ResponseHeader,
_ctx: &mut Self::CTX,
) -> Result<()> {
upstream_response.insert_header("X-Gateway-Version", "revoke-gateway/0.1.0")?;
Ok(())
}
}
impl GatewayProxyService {
async fn get_or_create_load_balancer(
&self,
service_name: &str,
) -> Result<Arc<LoadBalancer<RoundRobin>>> {
{
let lb_map = self.load_balancers.read().unwrap();
if let Some(lb) = lb_map.get(service_name) {
return Ok(Arc::clone(lb));
}
}
let services = self
.service_registry
.get_service(service_name)
.await
.map_err(|e| {
error!("Failed to get service {}: {:?}", service_name, e);
Error::new(ErrorType::HTTPStatus(503))
})?;
if services.is_empty() {
return Err(Error::new(ErrorType::HTTPStatus(503)));
}
let backends: BTreeSet<Backend> = services
.into_iter()
.filter_map(|service| {
let addr_str = format!("{}:{}", service.address, service.port);
use std::net::SocketAddr as StdSocketAddr;
addr_str.parse::<StdSocketAddr>().ok().map(|std_addr| Backend {
addr: PingoraSocketAddr::Inet(std_addr),
weight: 1,
ext: pingora_load_balancing::Extensions::new(),
})
})
.collect();
if backends.is_empty() {
return Err(Error::new(ErrorType::HTTPStatus(503)));
}
let discovery = pingora_load_balancing::discovery::Static::new(backends);
let backends_obj = Backends::new(discovery);
let mut lb = LoadBalancer::from_backends(backends_obj);
if self.config.enable_health_check {
let hc = health_check::TcpHealthCheck::new();
lb.set_health_check(hc);
lb.health_check_frequency = Some(std::time::Duration::from_secs(
self.config.health_check_interval,
));
}
let lb = Arc::new(lb);
{
let mut lb_map = self.load_balancers.write().unwrap();
lb_map.insert(service_name.to_string(), Arc::clone(&lb));
}
info!(
"Created load balancer for service '{}'",
service_name
);
Ok(lb)
}
}
fn extract_service_name(path: &str) -> String {
path.split('/').nth(1).unwrap_or("default").to_string()
}
#[derive(StructOpt, Debug, Clone)]
#[structopt(name = "revoke-gateway")]
pub struct GatewayConfig {
#[structopt(short, long, default_value = "0.0.0.0:8080")]
pub bind: String,
#[structopt(short = "w", long, default_value = "4")]
pub workers: usize,
#[structopt(long, default_value = "warn")]
pub log_level: String,
#[structopt(long)]
pub enable_health_check: bool,
#[structopt(long, default_value = "10")]
pub health_check_interval: u64,
#[structopt(long, default_value = "30")]
pub backend_refresh_interval: u64,
}
impl RevokeGateway {
pub fn new(service_registry: Arc<dyn ServiceRegistry>) -> anyhow::Result<Self> {
let mut server = Server::new(None)?;
server.bootstrap();
Ok(Self {
server,
service_registry,
})
}
pub fn run(mut self, config: GatewayConfig) -> anyhow::Result<()> {
let config = Arc::new(config);
let load_balancers = Arc::new(RwLock::new(HashMap::new()));
let proxy_service = GatewayProxyService {
service_registry: self.service_registry.clone(),
load_balancers,
config: config.clone(),
};
let mut lb = http_proxy_service(&self.server.configuration, proxy_service);
lb.add_tcp(&config.bind);
self.server.add_service(lb);
info!("Revoke Gateway starting on {}", config.bind);
info!(
"Health check: {}",
if config.enable_health_check {
"enabled"
} else {
"disabled"
}
);
self.server.run_forever();
}
}
pub mod router {
#[derive(Debug, Clone)]
pub struct Route {
pub path_prefix: String,
pub service_name: String,
pub strip_prefix: bool,
}
pub struct Router {
routes: Vec<Route>,
}
impl Default for Router {
fn default() -> Self {
Self::new()
}
}
impl Router {
pub fn new() -> Self {
Self { routes: Vec::new() }
}
pub fn add_route(&mut self, route: Route) {
self.routes.push(route);
}
pub fn match_route(&self, path: &str) -> Option<&Route> {
self.routes
.iter()
.find(|r| path.starts_with(&r.path_prefix))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_extract_service_name() {
assert_eq!(extract_service_name("/users/123"), "users");
assert_eq!(extract_service_name("/products/list"), "products");
assert_eq!(extract_service_name("/"), "default");
}
}