relay_core_lib/rule/engine/

executor.rs

use crate::rule::engine::actions;
use crate::rule::engine::compiled::CompiledRule;
use crate::rule::engine::compiler;
use crate::rule::engine::matcher;
use crate::rule::engine::state::{InMemoryRuleStateStore, RuleStateStore};
use crate::rule::engine::validator;
use crate::rule::model::{
    Rule, RuleExecutionEvent, RuleGroup, RuleOutcome, RuleStage, RuleTermination, RuleTraceSummary,
};
use relay_core_api::flow::Flow;
use relay_core_api::policy::ProxyPolicy;
use std::collections::HashMap;
use std::sync::Arc;

pub struct ExecutionContext {
    pub trace: Vec<RuleExecutionEvent>,
    pub variables: HashMap<String, String>,
    pub policy: Option<Arc<ProxyPolicy>>,
    pub summary: RuleTraceSummary,
    pub state_store: Arc<dyn RuleStateStore>,
}

#[derive(Debug)]
pub struct RuleEngine {
    compiled_rules: Vec<CompiledRule>,
    policy: Option<Arc<ProxyPolicy>>,
    state_store: Arc<dyn RuleStateStore>,
}

impl RuleEngine {
    pub fn new(
        rules: Vec<Rule>,
        rule_groups: Vec<RuleGroup>,
        policy: Option<Arc<ProxyPolicy>>,
        state_store: Option<Arc<dyn RuleStateStore>>,
    ) -> Self {
        let mut all_rules = Vec::new();

        // Flatten rules
        for rule in rules {
            all_rules.push(rule);
        }

        // Flatten rule groups
        for group in rule_groups {
            if group.active {
                for rule in group.rules {
                    all_rules.push(rule);
                }
            }
        }

        // Sort by priority (descending)
        all_rules.sort_by_key(|r| std::cmp::Reverse(r.priority));

        // Compile all rules
        let compiled_rules = all_rules.into_iter().map(compiler::compile_rule).collect();

        Self {
            compiled_rules,
            policy,
            state_store: state_store.unwrap_or_else(|| Arc::new(InMemoryRuleStateStore::new())),
        }
    }

    pub fn has_rules_for_stage(&self, stage: RuleStage) -> bool {
        self.compiled_rules
            .iter()
            .any(|r| r.original.active && r.original.stage == stage)
    }

    pub async fn execute(&self, stage: RuleStage, flow: &mut Flow) -> ExecutionContext {
        let mut ctx = ExecutionContext {
            trace: vec![],
            variables: HashMap::new(),
            policy: self.policy.clone(),
            summary: RuleTraceSummary::NoMatch,
            state_store: self.state_store.clone(),
        };

        let mut terminated = false;
        let mut modified_rules = Vec::new();

        for compiled_rule in &self.compiled_rules {
            let rule = &compiled_rule.original;

            if !rule.active {
                continue;
            }

            if rule.stage != stage {
                continue;
            }

            // Validate filter stage (using pre-compiled filter)
            if !validator::validate_filter_stage(&compiled_rule.filter, &stage) {
                ctx.trace.push(RuleExecutionEvent {
                    rule_id: rule.id.clone(),
                    stage: stage.clone(),
                    matched: false,
                    duration_us: 0,
                    outcome: RuleOutcome::Failed(format!("Filter invalid for stage {:?}", stage)),
                });
                continue;
            }

            // Validate action stage
            let mut actions_valid = true;
            for action in &rule.actions {
                if !validator::validate_action_stage(action, &stage) {
                    ctx.trace.push(RuleExecutionEvent {
                        rule_id: rule.id.clone(),
                        stage: stage.clone(),
                        matched: false,
                        duration_us: 0,
                        outcome: RuleOutcome::Failed(format!(
                            "Action {:?} not allowed in stage {:?}",
                            action, stage
                        )),
                    });
                    actions_valid = false;
                    break;
                }
            }
            if !actions_valid {
                continue;
            }

            // Match (using pre-compiled filter)
            let start = std::time::Instant::now();
            let matched = matcher::matches(&compiled_rule.filter, flow);

            if matched {
                let timeout_ms = rule.constraints.as_ref().and_then(|c| c.timeout_ms);

                let action_execution = async {
                    let mut rule_outcome = RuleOutcome::MatchedAndExecuted;
                    let mut rule_terminated = false;

                    for action in &rule.actions {
                        match actions::execute_action(action, flow, &mut ctx).await {
                            actions::ActionOutcome::Continue => {}
                            actions::ActionOutcome::Terminated(reason) => {
                                rule_outcome = RuleOutcome::MatchedAndTerminated;
                                ctx.summary = RuleTraceSummary::Terminated {
                                    rule_id: rule.id.clone(),
                                    reason,
                                };
                                rule_terminated = true;
                                break;
                            }
                            actions::ActionOutcome::Failed(err) => {
                                rule_outcome = RuleOutcome::Failed(err);
                                break;
                            }
                        }
                    }
                    (rule_outcome, rule_terminated)
                };

                let (rule_outcome, rule_terminated) = if let Some(ms) = timeout_ms {
                    match tokio::time::timeout(
                        std::time::Duration::from_millis(ms),
                        action_execution,
                    )
                    .await
                    {
                        Ok(res) => res,
                        Err(_) => (
                            RuleOutcome::Failed(format!("Rule execution timed out after {}ms", ms)),
                            false,
                        ),
                    }
                } else {
                    action_execution.await
                };

                if rule_terminated {
                    terminated = true;
                }

                ctx.trace.push(RuleExecutionEvent {
                    rule_id: rule.id.clone(),
                    stage: stage.clone(),
                    matched: true,
                    duration_us: start.elapsed().as_micros() as u64,
                    outcome: rule_outcome.clone(),
                });

                if matches!(rule_outcome, RuleOutcome::MatchedAndExecuted) {
                    modified_rules.push(rule.id.clone());
                }

                if terminated {
                    break;
                }

                if let RuleTermination::Stop = rule.termination {
                    break;
                }
            }
        }

        if !terminated && !modified_rules.is_empty() {
            ctx.summary = RuleTraceSummary::Modified {
                rule_ids: modified_rules,
            };
        }

        ctx
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::rule::model::{
        Action, BodySource, Filter, Rule, RuleConstraints, RuleOutcome, RuleStage, RuleTermination,
        StringMatcher,
    };
    use chrono::Utc;
    use relay_core_api::flow::{Flow, HttpRequest, Layer, NetworkInfo, TransportProtocol};
    use url::Url;
    use uuid::Uuid;

    fn create_test_flow() -> Flow {
        Flow {
            id: Uuid::new_v4(),
            start_time: Utc::now(),
            end_time: None,
            network: NetworkInfo {
                client_ip: "127.0.0.1".to_string(),
                client_port: 12345,
                server_ip: "1.1.1.1".to_string(),
                server_port: 80,
                protocol: TransportProtocol::TCP,
                tls: false,
                tls_version: None,
                sni: None,
            },
            layer: Layer::Http(relay_core_api::flow::HttpLayer {
                request: HttpRequest {
                    method: "GET".to_string(),
                    url: Url::parse("http://example.com").unwrap(),
                    version: "HTTP/1.1".to_string(),
                    headers: vec![],
                    cookies: vec![],
                    query: vec![],
                    body: None,
                },
                response: None,
                error: None,
            }),
            tags: vec![],
            meta: std::collections::HashMap::new(),
        }
    }

    #[tokio::test]
    async fn test_rule_timeout() {
        let rule = Rule {
            id: "test-rule-timeout".to_string(),
            name: "Test Rule Timeout".to_string(),
            active: true,
            stage: RuleStage::RequestHeaders,
            priority: 0,
            termination: RuleTermination::Continue,
            filter: Filter::All,
            actions: vec![
                Action::Delay { ms: 200 }, // Delay 200ms
            ],
            constraints: Some(RuleConstraints {
                timeout_ms: Some(50), // Timeout 50ms
            }),
        };

        let engine = RuleEngine::new(vec![rule], vec![], None, None);
        let mut flow = create_test_flow();

        // Execute in RequestHeaders stage
        let ctx = engine.execute(RuleStage::RequestHeaders, &mut flow).await;

        // Expect 1 trace entry with Failed outcome due to timeout
        assert_eq!(ctx.trace.len(), 1);
        if let RuleOutcome::Failed(msg) = &ctx.trace[0].outcome {
            assert!(msg.contains("timed out"));
        } else {
            panic!("Expected Failed outcome, got {:?}", ctx.trace[0].outcome);
        }
    }

    #[tokio::test]
    async fn test_filter_stage_validation_failure() {
        let rule = Rule {
            id: "test-rule-1".to_string(),
            name: "Test Rule 1".to_string(),
            active: true,
            stage: RuleStage::Connect, // Stage is Connect
            priority: 0,
            termination: RuleTermination::Continue,
            filter: Filter::Path(StringMatcher::Exact("/foo".to_string())), // Path filter invalid in Connect
            actions: vec![],
            constraints: None,
        };

        let engine = RuleEngine::new(vec![rule], vec![], None, None);
        let mut flow = create_test_flow();

        // Execute in Connect stage
        let ctx = engine.execute(RuleStage::Connect, &mut flow).await;

        // Expect 1 trace entry with Failed outcome
        assert_eq!(ctx.trace.len(), 1);
        if let RuleOutcome::Failed(msg) = &ctx.trace[0].outcome {
            assert!(msg.contains("Filter invalid"));
        } else {
            assert!(
                false,
                "Expected Failed outcome, got {:?}",
                ctx.trace[0].outcome
            );
        }
    }

    #[tokio::test]
    async fn test_action_stage_validation_failure() {
        let rule = Rule {
            id: "test-rule-2".to_string(),
            name: "Test Rule 2".to_string(),
            active: true,
            stage: RuleStage::Connect, // Stage is Connect
            priority: 0,
            termination: RuleTermination::Continue,
            filter: Filter::All, // Valid filter
            actions: vec![
                Action::SetResponseBody {
                    body: BodySource::Text("foo".to_string()),
                }, // Invalid action in Connect
            ],
            constraints: None,
        };

        let engine = RuleEngine::new(vec![rule], vec![], None, None);
        let mut flow = create_test_flow();

        // Execute in Connect stage
        let ctx = engine.execute(RuleStage::Connect, &mut flow).await;

        // Expect 1 trace entry with Failed outcome
        assert_eq!(ctx.trace.len(), 1);
        if let RuleOutcome::Failed(msg) = &ctx.trace[0].outcome {
            assert!(msg.contains("Action"));
            assert!(msg.contains("not allowed"));
        } else {
            assert!(
                false,
                "Expected Failed outcome, got {:?}",
                ctx.trace[0].outcome
            );
        }
    }

    #[test]
    fn test_has_rules_for_stage_respects_active_flag() {
        let inactive_rule = Rule {
            id: "inactive-rh".to_string(),
            name: "Inactive".to_string(),
            active: false,
            stage: RuleStage::RequestHeaders,
            priority: 10,
            termination: RuleTermination::Continue,
            filter: Filter::All,
            actions: vec![Action::Tag {
                key: "k".to_string(),
                value: "v".to_string(),
            }],
            constraints: None,
        };
        let active_connect_rule = Rule {
            id: "active-connect".to_string(),
            name: "Active Connect".to_string(),
            active: true,
            stage: RuleStage::Connect,
            priority: 10,
            termination: RuleTermination::Continue,
            filter: Filter::All,
            actions: vec![Action::Drop],
            constraints: None,
        };

        let engine = RuleEngine::new(vec![inactive_rule, active_connect_rule], vec![], None, None);
        assert!(
            !engine.has_rules_for_stage(RuleStage::RequestHeaders),
            "inactive rules should not count for stage presence"
        );
        assert!(engine.has_rules_for_stage(RuleStage::Connect));
    }

    #[tokio::test]
    async fn test_execute_orders_rules_by_priority_descending() {
        let low = Rule {
            id: "low-pri".to_string(),
            name: "low".to_string(),
            active: true,
            stage: RuleStage::RequestHeaders,
            priority: 1,
            termination: RuleTermination::Continue,
            filter: Filter::All,
            actions: vec![Action::Tag {
                key: "order".to_string(),
                value: "low".to_string(),
            }],
            constraints: None,
        };
        let high = Rule {
            id: "high-pri".to_string(),
            name: "high".to_string(),
            active: true,
            stage: RuleStage::RequestHeaders,
            priority: 100,
            termination: RuleTermination::Continue,
            filter: Filter::All,
            actions: vec![Action::Tag {
                key: "order".to_string(),
                value: "high".to_string(),
            }],
            constraints: None,
        };
        let engine = RuleEngine::new(vec![low, high], vec![], None, None);
        let mut flow = create_test_flow();

        let ctx = engine.execute(RuleStage::RequestHeaders, &mut flow).await;
        assert_eq!(ctx.trace.len(), 2);
        assert_eq!(ctx.trace[0].rule_id, "high-pri");
        assert_eq!(ctx.trace[1].rule_id, "low-pri");
        assert_eq!(
            flow.tags,
            vec!["order:high".to_string(), "order:low".to_string()]
        );
    }
}