online_dsl_forge/sema/
schema.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use crate::parser::{AstExpression, BinaryOp, Diagnostic, SourceSpan, UnaryOp};
4use serde::{Deserialize, Serialize};
5
6use crate::sema::profile::{BodyAccess, BodyTarget, Phase};
7
8mod oxirule;
9
10#[derive(Debug, Clone, Copy, Deserialize, Eq, PartialEq, Serialize)]
11#[serde(rename_all = "snake_case")]
12pub enum TypeClass {
13  Null,
14  Bool,
15  Int,
16  Float,
17  String,
18  Bytes,
19  Array,
20  Object,
21  Dyn,
22  RegexLiteral,
23}
24
25#[derive(Debug, Clone, Copy, Deserialize, Eq, PartialEq, Serialize)]
26#[serde(rename_all = "snake_case")]
27pub enum CapabilityKind {
28  Function,
29  Method,
30  UnaryOp,
31  BinaryOp,
32}
33
34#[derive(Debug, Clone, Copy, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
35#[serde(rename_all = "snake_case")]
36pub enum RegexFlavor {
37  Default,
38  HeaderName,
39}
40
41#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
42pub struct RegexArgMeta {
43  pub index: usize,
44  pub flavor: RegexFlavor,
45}
46
47impl RegexArgMeta {
48  pub fn new(index: usize, flavor: RegexFlavor) -> Self {
49    Self { index, flavor }
50  }
51}
52
53#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
54#[serde(rename_all = "snake_case")]
55pub enum CostModel {
56  Constant(u32),
57  LinearInput { factor: u32 },
58  LinearCollection { factor: u32 },
59  RegexMatch { factor: u32, precompiled: bool },
60}
61
62impl CostModel {
63  pub fn static_cost(&self) -> u64 {
64    match self {
65      Self::Constant(value) => u64::from(*value),
66      Self::LinearInput { factor }
67      | Self::LinearCollection { factor }
68      | Self::RegexMatch { factor, .. } => u64::from(*factor),
69    }
70  }
71}
72
73impl Default for CostModel {
74  fn default() -> Self {
75    Self::Constant(1)
76  }
77}
78
79#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
80pub struct VariableMeta {
81  pub name: String,
82  pub type_class: TypeClass,
83  pub phases: BTreeSet<Phase>,
84}
85
86impl VariableMeta {
87  pub fn new(name: impl Into<String>) -> Self {
88    Self {
89      name: name.into(),
90      type_class: TypeClass::Dyn,
91      phases: BTreeSet::new(),
92    }
93  }
94
95  pub fn with_type(mut self, type_class: TypeClass) -> Self {
96    self.type_class = type_class;
97    self
98  }
99
100  pub fn with_phases(mut self, phases: impl IntoIterator<Item = Phase>) -> Self {
101    self.phases = phases.into_iter().collect();
102    self
103  }
104
105  pub fn is_available_in(&self, phase: Phase) -> bool {
106    self.phases.is_empty() || self.phases.contains(&phase)
107  }
108}
109
110#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
111pub struct CapabilityMeta {
112  pub name: String,
113  pub kind: CapabilityKind,
114  pub arity: usize,
115  pub receiver: Option<TypeClass>,
116  pub args: Vec<TypeClass>,
117  pub result: TypeClass,
118  pub phases: BTreeSet<Phase>,
119  pub body_access: BodyAccess,
120  pub regex_args: Vec<RegexArgMeta>,
121  pub deterministic: bool,
122  pub side_effect_free: bool,
123  pub cost: CostModel,
124}
125
126impl CapabilityMeta {
127  pub fn function(name: impl Into<String>, arity: usize) -> Self {
128    Self::new(name, CapabilityKind::Function, arity)
129  }
130
131  pub fn method(name: impl Into<String>, arity: usize) -> Self {
132    Self::new(name, CapabilityKind::Method, arity)
133  }
134
135  pub fn unary_operator(op: UnaryOp) -> Self {
136    Self::new(op.as_str(), CapabilityKind::UnaryOp, 1)
137  }
138
139  pub fn binary_operator(op: BinaryOp) -> Self {
140    Self::new(op.as_str(), CapabilityKind::BinaryOp, 2)
141  }
142
143  pub fn with_phases(mut self, phases: impl IntoIterator<Item = Phase>) -> Self {
144    self.phases = phases.into_iter().collect();
145    self
146  }
147
148  pub fn with_receiver(mut self, receiver: TypeClass) -> Self {
149    self.receiver = Some(receiver);
150    self
151  }
152
153  pub fn with_args(mut self, args: impl IntoIterator<Item = TypeClass>) -> Self {
154    self.args = args.into_iter().collect();
155    self.arity = self.args.len();
156    self
157  }
158
159  pub fn with_result(mut self, result: TypeClass) -> Self {
160    self.result = result;
161    self
162  }
163
164  pub fn with_body_access(mut self, access: BodyAccess) -> Self {
165    self.body_access = access;
166    self
167  }
168
169  pub fn with_regex_arg(mut self, index: usize, flavor: RegexFlavor) -> Self {
170    self.regex_args.push(RegexArgMeta::new(index, flavor));
171    self
172  }
173
174  pub fn with_cost(mut self, cost: CostModel) -> Self {
175    self.cost = cost;
176    self
177  }
178
179  pub fn with_deterministic(mut self, deterministic: bool) -> Self {
180    self.deterministic = deterministic;
181    self
182  }
183
184  pub fn with_side_effect_free(mut self, side_effect_free: bool) -> Self {
185    self.side_effect_free = side_effect_free;
186    self
187  }
188
189  pub fn is_available_in(&self, phase: Phase) -> bool {
190    self.phases.is_empty() || self.phases.contains(&phase)
191  }
192
193  pub fn ticket(&self) -> CapabilityTicket {
194    CapabilityTicket::new(self.kind, self.name.clone(), self.arity)
195  }
196
197  fn new(name: impl Into<String>, kind: CapabilityKind, arity: usize) -> Self {
198    Self {
199      name: name.into(),
200      kind,
201      arity,
202      receiver: None,
203      args: vec![TypeClass::Dyn; arity],
204      result: TypeClass::Dyn,
205      phases: BTreeSet::new(),
206      body_access: BodyAccess::None,
207      regex_args: Vec::new(),
208      deterministic: true,
209      side_effect_free: true,
210      cost: CostModel::default(),
211    }
212  }
213}
214
215#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
216pub struct CapabilityTicket {
217  pub kind: CapabilityKind,
218  pub name: String,
219  pub arity: usize,
220}
221
222impl CapabilityTicket {
223  pub fn new(kind: CapabilityKind, name: impl Into<String>, arity: usize) -> Self {
224    Self {
225      kind,
226      name: name.into(),
227      arity,
228    }
229  }
230}
231
232impl Ord for CapabilityTicket {
233  fn cmp(&self, other: &Self) -> std::cmp::Ordering {
234    (self.kind_order(), &self.name, self.arity).cmp(&(other.kind_order(), &other.name, other.arity))
235  }
236}
237
238impl PartialOrd for CapabilityTicket {
239  fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
240    Some(self.cmp(other))
241  }
242}
243
244impl CapabilityTicket {
245  fn kind_order(&self) -> u8 {
246    match self.kind {
247      CapabilityKind::Function => 0,
248      CapabilityKind::Method => 1,
249      CapabilityKind::UnaryOp => 2,
250      CapabilityKind::BinaryOp => 3,
251    }
252  }
253}
254
255#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
256pub struct BodyPathRule {
257  pub path: Vec<String>,
258  pub target: BodyTarget,
259  pub access: BodyAccess,
260}
261
262impl BodyPathRule {
263  pub fn new(
264    path: impl IntoIterator<Item = impl Into<String>>,
265    target: BodyTarget,
266    access: BodyAccess,
267  ) -> Self {
268    Self {
269      path: path.into_iter().map(Into::into).collect(),
270      target,
271      access,
272    }
273  }
274}
275
276#[derive(Debug, Clone, Deserialize, PartialEq, Serialize)]
277pub struct ExpressionFunction {
278  pub name: String,
279  pub params: Vec<String>,
280  pub expression: AstExpression,
281  #[serde(default)]
282  pub scope: ExpressionFunctionScope,
283}
284
285#[derive(
286  Debug, Clone, Copy, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize,
287)]
288#[serde(rename_all = "snake_case")]
289pub enum ExpressionFunctionScope {
290  #[default]
291  Global,
292  Local,
293}
294
295impl ExpressionFunctionScope {
296  pub fn call_scope(self) -> Self {
297    match self {
298      Self::Global => Self::Global,
299      Self::Local => Self::Local,
300    }
301  }
302}
303
304#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize)]
305pub struct ExpressionFunctionDiagnostic {
306  pub message: String,
307  pub span: SourceSpan,
308}
309
310impl ExpressionFunctionDiagnostic {
311  fn new(message: impl Into<String>, span: SourceSpan) -> Self {
312    Self {
313      message: message.into(),
314      span,
315    }
316  }
317
318  pub fn diagnostic(&self) -> Diagnostic {
319    Diagnostic::new(self.message.clone(), self.span)
320  }
321}
322
323#[derive(Debug, Clone, Default, Deserialize, PartialEq, Serialize)]
324pub struct RuntimeSchema {
325  variables: BTreeMap<String, VariableMeta>,
326  functions: BTreeMap<String, BTreeMap<usize, CapabilityMeta>>,
327  methods: BTreeMap<String, BTreeMap<usize, CapabilityMeta>>,
328  unary_ops: BTreeMap<String, CapabilityMeta>,
329  binary_ops: BTreeMap<String, CapabilityMeta>,
330  body_paths: Vec<BodyPathRule>,
331  expression_functions: BTreeMap<String, ExpressionFunction>,
332  #[serde(default)]
333  local_expression_functions: BTreeMap<String, ExpressionFunction>,
334  #[serde(default)]
335  expression_function_diagnostics: Vec<ExpressionFunctionDiagnostic>,
336}
337
338impl RuntimeSchema {
339  pub fn new() -> Self {
340    Self::default()
341  }
342
343  pub fn waf() -> Self {
344    let mut schema = Self::new();
345    schema
346      .add_variable("Request")
347      .add_variable("Response")
348      .add_variable("Stream")
349      .add_waf_body_paths()
350      .add_waf_body_methods()
351      .add_waf_regex_methods();
352    schema
353  }
354
355  pub fn add_variable(&mut self, name: impl Into<String>) -> &mut Self {
356    self.add_variable_meta(VariableMeta::new(name))
357  }
358
359  pub fn add_variable_meta(&mut self, variable: VariableMeta) -> &mut Self {
360    self.variables.insert(variable.name.clone(), variable);
361    self
362  }
363
364  pub fn add_function(&mut self, name: impl Into<String>, arity: usize) -> &mut Self {
365    self.add_function_capability(CapabilityMeta::function(name, arity))
366  }
367
368  pub fn add_function_capability(&mut self, capability: CapabilityMeta) -> &mut Self {
369    self
370      .functions
371      .entry(capability.name.clone())
372      .or_default()
373      .insert(capability.arity, capability);
374    self
375  }
376
377  pub fn add_method(&mut self, name: impl Into<String>, arity: usize) -> &mut Self {
378    self.add_method_capability(CapabilityMeta::method(name, arity))
379  }
380
381  pub fn add_method_capability(&mut self, capability: CapabilityMeta) -> &mut Self {
382    self
383      .methods
384      .entry(capability.name.clone())
385      .or_default()
386      .insert(capability.arity, capability);
387    self
388  }
389
390  pub fn add_unary_operator_capability(&mut self, capability: CapabilityMeta) -> &mut Self {
391    self.unary_ops.insert(capability.name.clone(), capability);
392    self
393  }
394
395  pub fn add_binary_operator_capability(&mut self, capability: CapabilityMeta) -> &mut Self {
396    self.binary_ops.insert(capability.name.clone(), capability);
397    self
398  }
399
400  pub fn add_body_path(
401    &mut self,
402    path: impl IntoIterator<Item = impl Into<String>>,
403    target: BodyTarget,
404    access: BodyAccess,
405  ) -> &mut Self {
406    self
407      .body_paths
408      .push(BodyPathRule::new(path, target, access));
409    self
410  }
411
412  pub fn add_waf_body_paths(&mut self) -> &mut Self {
413    for root in ["Request", "Response"] {
414      let target = if root == "Request" {
415        BodyTarget::Request
416      } else {
417        BodyTarget::Response
418      };
419      self.add_body_path([root, "Body", "Size"], target, BodyAccess::SizeOnly);
420      self.add_body_path([root, "Body", "Bytes"], target, BodyAccess::PrefixBytes);
421      self.add_body_path([root, "Body", "Text"], target, BodyAccess::PrefixBytes);
422      self.add_body_path(
423        [root, "Body", "IsTruncated"],
424        target,
425        BodyAccess::PrefixBytes,
426      );
427    }
428    self.add_body_path(
429      ["Stream", "Payload"],
430      BodyTarget::Stream,
431      BodyAccess::PrefixBytes,
432    );
433    self
434  }
435
436  pub fn add_waf_body_methods(&mut self) -> &mut Self {
437    for (name, arity) in [
438      ("isFormat", 1),
439      ("isBinaryFormat", 1),
440      ("matchesFormat", 1),
441      ("contains", 1),
442      ("containsBytes", 1),
443      ("containsAny", 1),
444      ("matchesAny", 1),
445      ("scan", 1),
446      ("anomalyScore", 1),
447      ("malformedScore", 1),
448      ("promptInjectionScore", 0),
449    ] {
450      self.add_method_capability(
451        CapabilityMeta::method(name, arity).with_body_access(BodyAccess::PrefixBytes),
452      );
453    }
454    self.add_method_capability(
455      CapabilityMeta::method("matches", 1)
456        .with_body_access(BodyAccess::PrefixBytes)
457        .with_regex_arg(0, RegexFlavor::Default),
458    );
459    self
460  }
461
462  pub fn add_waf_regex_methods(&mut self) -> &mut Self {
463    self
464      .add_method_capability(
465        CapabilityMeta::method("anyValueMatches", 1).with_regex_arg(0, RegexFlavor::Default),
466      )
467      .add_method_capability(
468        CapabilityMeta::method("anyKeyMatches", 1).with_regex_arg(0, RegexFlavor::Default),
469      )
470      .add_method_capability(
471        CapabilityMeta::method("anyMatches", 1).with_regex_arg(0, RegexFlavor::Default),
472      )
473      .add_method_capability(
474        CapabilityMeta::method("anyNameMatches", 1)
475          .with_regex_arg(0, RegexFlavor::Default)
476          .with_regex_arg(0, RegexFlavor::HeaderName),
477      )
478      .add_method_capability(
479        CapabilityMeta::method("anyEntryMatches", 2)
480          .with_regex_arg(0, RegexFlavor::Default)
481          .with_regex_arg(0, RegexFlavor::HeaderName)
482          .with_regex_arg(1, RegexFlavor::Default),
483      )
484      .add_method_capability(
485        CapabilityMeta::method("allEntriesMatch", 2)
486          .with_regex_arg(0, RegexFlavor::HeaderName)
487          .with_regex_arg(1, RegexFlavor::Default),
488      );
489    self
490  }
491
492  pub fn add_expression_function(
493    &mut self,
494    name: impl Into<String>,
495    params: impl IntoIterator<Item = impl Into<String>>,
496    expression: AstExpression,
497  ) -> &mut Self {
498    self.add_scoped_expression_function(ExpressionFunctionScope::Global, name, params, expression)
499  }
500
501  pub fn add_local_expression_function(
502    &mut self,
503    name: impl Into<String>,
504    params: impl IntoIterator<Item = impl Into<String>>,
505    expression: AstExpression,
506  ) -> &mut Self {
507    self.add_scoped_expression_function(ExpressionFunctionScope::Local, name, params, expression)
508  }
509
510  pub fn add_scoped_expression_function(
511    &mut self,
512    scope: ExpressionFunctionScope,
513    name: impl Into<String>,
514    params: impl IntoIterator<Item = impl Into<String>>,
515    expression: AstExpression,
516  ) -> &mut Self {
517    let name = name.into();
518    let params = params.into_iter().map(Into::into).collect::<Vec<_>>();
519    self.add_function(name.clone(), params.len());
520    let target = match scope {
521      ExpressionFunctionScope::Global => &mut self.expression_functions,
522      ExpressionFunctionScope::Local => &mut self.local_expression_functions,
523    };
524    if target.contains_key(&name) {
525      self
526        .expression_function_diagnostics
527        .push(ExpressionFunctionDiagnostic::new(
528          format!("duplicate expression function {name} in {scope:?} scope"),
529          expression.span,
530        ));
531    }
532    target.insert(
533      name.clone(),
534      ExpressionFunction {
535        name,
536        params,
537        expression,
538        scope,
539      },
540    );
541    self
542  }
543
544  pub fn has_variable(&self, name: &str) -> bool {
545    self.variables.contains_key(name)
546  }
547
548  pub fn variables(&self) -> impl Iterator<Item = &VariableMeta> {
549    self.variables.values()
550  }
551
552  pub fn variable(&self, name: &str) -> Option<&VariableMeta> {
553    self.variables.get(name)
554  }
555
556  pub fn function_accepts(&self, name: &str, arity: usize) -> SignatureMatch {
557    signature_accepts(&self.functions, name, arity)
558  }
559
560  pub fn method_accepts(&self, name: &str, arity: usize) -> SignatureMatch {
561    signature_accepts(&self.methods, name, arity)
562  }
563
564  pub fn function_capability(&self, name: &str, arity: usize) -> Option<&CapabilityMeta> {
565    self
566      .functions
567      .get(name)
568      .and_then(|entries| entries.get(&arity))
569  }
570
571  pub fn method_capability(&self, name: &str, arity: usize) -> Option<&CapabilityMeta> {
572    self
573      .methods
574      .get(name)
575      .and_then(|entries| entries.get(&arity))
576  }
577
578  pub fn unary_operator_capability(&self, op: UnaryOp) -> Option<&CapabilityMeta> {
579    self.unary_ops.get(op.as_str())
580  }
581
582  pub fn binary_operator_capability(&self, op: BinaryOp) -> Option<&CapabilityMeta> {
583    self.binary_ops.get(op.as_str())
584  }
585
586  pub fn function_capabilities(&self) -> impl Iterator<Item = &CapabilityMeta> {
587    self.functions.values().flat_map(BTreeMap::values)
588  }
589
590  pub fn method_capabilities(&self) -> impl Iterator<Item = &CapabilityMeta> {
591    self.methods.values().flat_map(BTreeMap::values)
592  }
593
594  pub fn unary_operator_capabilities(&self) -> impl Iterator<Item = &CapabilityMeta> {
595    self.unary_ops.values()
596  }
597
598  pub fn binary_operator_capabilities(&self) -> impl Iterator<Item = &CapabilityMeta> {
599    self.binary_ops.values()
600  }
601
602  pub fn body_access_for_path(&self, path: &[String]) -> Option<(BodyTarget, BodyAccess)> {
603    self
604      .body_paths
605      .iter()
606      .find(|rule| rule.path == path)
607      .map(|rule| (rule.target, rule.access))
608  }
609
610  pub fn expression_function(&self, name: &str) -> Option<&ExpressionFunction> {
611    self.expression_function_for_scope(name, ExpressionFunctionScope::Local)
612  }
613
614  pub fn expression_function_for_scope(
615    &self,
616    name: &str,
617    scope: ExpressionFunctionScope,
618  ) -> Option<&ExpressionFunction> {
619    match scope {
620      ExpressionFunctionScope::Global => self.expression_functions.get(name),
621      ExpressionFunctionScope::Local => self
622        .local_expression_functions
623        .get(name)
624        .or_else(|| self.expression_functions.get(name)),
625    }
626  }
627
628  pub fn expression_functions(&self) -> impl Iterator<Item = &ExpressionFunction> {
629    self
630      .expression_functions
631      .values()
632      .chain(self.local_expression_functions.values())
633  }
634
635  pub fn expression_function_diagnostics(&self) -> &[ExpressionFunctionDiagnostic] {
636    &self.expression_function_diagnostics
637  }
638}
639
640#[derive(Debug, Clone, Copy, Eq, PartialEq)]
641pub enum SignatureMatch {
642  Unknown,
643  ArityMismatch,
644  Matches,
645}
646
647fn signature_accepts(
648  signatures: &BTreeMap<String, BTreeMap<usize, CapabilityMeta>>,
649  name: &str,
650  arity: usize,
651) -> SignatureMatch {
652  match signatures.get(name) {
653    Some(accepted) if accepted.contains_key(&arity) => SignatureMatch::Matches,
654    Some(_) => SignatureMatch::ArityMismatch,
655    None => SignatureMatch::Unknown,
656  }
657}
online_dsl_forge/sema/schema.rs

online_dsl_forge/sema/
schema.rs
