mod predicates;
pub use predicates::*;
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::RwLock;
use serde::{Serialize, Deserialize};
use crate::config::Config;
use crate::core::{ProxyRequest, ProxyError, Route};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouteConfig {
pub id: String,
pub target: String,
#[serde(default)]
pub filters: Vec<String>,
#[serde(default = "default_priority")]
pub priority: i32,
#[serde(default)]
pub predicates: Vec<PredicateConfig>,
}
fn default_priority() -> i32 {
0
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredicateConfig {
pub type_: String,
pub config: serde_json::Value,
}
#[async_trait]
pub trait Predicate: Send + Sync + std::fmt::Debug {
async fn matches(&self, request: &ProxyRequest) -> bool;
fn predicate_type(&self) -> &str;
}
#[derive(Debug)]
pub struct PredicateRouter {
routes: RwLock<HashMap<String, RouteWithPredicates>>,
sorted_routes: RwLock<Vec<RouteWithPredicates>>,
config: Arc<Config>,
}
#[derive(Debug, Clone)]
struct RouteWithPredicates {
route: Route,
predicates: Vec<Arc<dyn Predicate>>,
priority: i32,
}
impl PredicateRouter {
pub async fn new(config: Arc<Config>) -> Result<Self, ProxyError> {
let router = Self {
routes: RwLock::new(HashMap::new()),
sorted_routes: RwLock::new(Vec::new()),
config,
};
router.load_routes_from_config().await?;
Ok(router)
}
async fn load_routes_from_config(&self) -> Result<(), ProxyError> {
let route_configs: Option<Vec<RouteConfig>> = self.config.get("routes")?;
if let Some(route_configs) = route_configs {
for route_config in route_configs {
let route = Route {
id: route_config.id.clone(),
target_base_url: route_config.target.clone(),
path_pattern: String::new(), filter_ids: route_config.filters.clone(),
};
let mut predicates = Vec::new();
for predicate_config in &route_config.predicates {
let predicate = PredicateFactory::create_predicate(
&predicate_config.type_,
predicate_config.config.clone(),
)?;
predicates.push(predicate);
}
self.add_route_with_predicates(
route,
predicates,
route_config.priority,
).await?;
}
}
Ok(())
}
async fn add_route_with_predicates(
&self,
route: Route,
predicates: Vec<Arc<dyn Predicate>>,
priority: i32,
) -> Result<(), ProxyError> {
let route_with_predicates = RouteWithPredicates {
route: route.clone(),
predicates,
priority,
};
{
let mut routes = self.routes.write().await;
routes.insert(route.id.clone(), route_with_predicates.clone());
}
{
let mut sorted_routes = self.sorted_routes.write().await;
sorted_routes.push(route_with_predicates);
sorted_routes.sort_by(|a, b| b.priority.cmp(&a.priority));
}
Ok(())
}
}
#[async_trait]
impl crate::core::Router for PredicateRouter {
async fn route(&self, request: &ProxyRequest) -> Result<Route, ProxyError> {
let sorted_routes = self.sorted_routes.read().await;
for route_with_predicates in sorted_routes.iter() {
let mut all_match = true;
for predicate in &route_with_predicates.predicates {
if !predicate.matches(request).await {
all_match = false;
break;
}
}
if all_match {
return Ok(route_with_predicates.route.clone());
}
}
Err(ProxyError::RoutingError(format!("No route matched the request: {} {}",
request.method, request.path)))
}
async fn get_routes(&self) -> Vec<Route> {
let routes = self.routes.read().await;
routes.values().map(|r| r.route.clone()).collect()
}
async fn add_route(&self, route: Route) -> Result<(), ProxyError> {
self.add_route_with_predicates(route, Vec::new(), 0).await
}
async fn remove_route(&self, route_id: &str) -> Result<(), ProxyError> {
{
let mut routes = self.routes.write().await;
if routes.remove(route_id).is_none() {
return Err(ProxyError::RoutingError(format!("Route not found: {}", route_id)));
}
}
{
let mut sorted_routes = self.sorted_routes.write().await;
sorted_routes.retain(|r| r.route.id != route_id);
}
Ok(())
}
}
#[derive(Debug)]
pub struct PredicateFactory;
impl PredicateFactory {
pub fn create_predicate(
predicate_type: &str,
config: serde_json::Value,
) -> Result<Arc<dyn Predicate>, ProxyError> {
match predicate_type {
"path" => {
let path_config: PathPredicateConfig = serde_json::from_value(config)
.map_err(|e| ProxyError::RoutingError(
format!("Invalid path predicate config: {}", e)
))?;
Ok(Arc::new(PathPredicate::new(path_config)))
},
"method" => {
let method_config: MethodPredicateConfig = serde_json::from_value(config)
.map_err(|e| ProxyError::RoutingError(
format!("Invalid method predicate config: {}", e)
))?;
Ok(Arc::new(MethodPredicate::new(method_config)))
},
"header" => {
let header_config: HeaderPredicateConfig = serde_json::from_value(config)
.map_err(|e| ProxyError::RoutingError(
format!("Invalid header predicate config: {}", e)
))?;
Ok(Arc::new(HeaderPredicate::new(header_config)))
},
"query" => {
let query_config: QueryPredicateConfig = serde_json::from_value(config)
.map_err(|e| ProxyError::RoutingError(
format!("Invalid query predicate config: {}", e)
))?;
Ok(Arc::new(QueryPredicate::new(query_config)))
},
_ => Err(ProxyError::RoutingError(
format!("Unknown predicate type: {}", predicate_type)
)),
}
}
}