kata-containers/regorus
1
0
Fork 0
mirror of https://github.com/microsoft/regorus.git synced 2026-08-05 02:16:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(memory): Allocator-backed global memory limits (#544)

Browse Source 
Policy evaluation at scale needs to be able to set memory limits
so that a bad policy does not hog memory or to ensure that
policy evaluation itself does not use too much memory which could
cause other components to suffer.

This PR introduces capability to set and enforce global memory limits.
It also lays the groundwork for enabling per evaluation limits in future.

Once a global memory limit is set, Regorus maintains per thread counters
to track memory activity (allocation, deallocation) of a thread.
These counters are periodically flushed to global memory counters.
Per thread counters avoid the contention that updating global counters
on each alloc/free would cause.

Policy evaluation periodically checks these counters and raises errors
if allocated memory has exceeded the configured limit.

Currently memory limit capability is exposed only to FFI and C#.

Also update mimalloc to v2.2.6

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2026-01-24 07:08:54 +05:30
committed by GitHub
parent 80686d6ed1
commit fd59bb5a91
93 changed files with 9772 additions and 3779 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/builtins/aggregates.rs
+10 -2
View File
@@ -5,7 +5,7 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::{ensure_args_count, ensure_numeric};
use crate::builtins::utils::{enforce_limit, ensure_args_count, ensure_numeric};
use crate::lexer::Span;
use crate::number::Number;
use crate::value::Value;
@@ -104,7 +104,15 @@ fn sort(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> Res
Value::from(ac)
}
// Sorting a set produces array.
Value::Set(a) => Value::from(a.iter().cloned().collect::<Vec<Value>>()),
Value::Set(a) => {
let mut items = Vec::with_capacity(a.len());
for value in a.iter() {
items.push(value.clone());
// Guard array growth while materializing the sorted set.
enforce_limit()?;
}
Value::from(items)
}
a => {
let span = params[0].span();
bail!(span.error(format!("`sort` requires array/set argument. Got `{a}`.").as_str()))
src/builtins/encoding.rs
+22 -11
View File
@@ -4,6 +4,8 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
#[cfg(feature = "urlquery")]
use crate::builtins::utils::enforce_limit;
#[allow(unused)]
use crate::builtins::utils::{
ensure_args_count, ensure_object, ensure_string, ensure_string_collection,
@@ -255,7 +257,9 @@ fn urlquery_decode_object(
let key = Value::String(k.clone().into());
let value = Value::String(v.clone().into());
if let Ok(a) = map.entry(key).or_insert(Value::new_array()).as_array_mut() {
a.push(value)
a.push(value);
// Guard decoded parameter accumulation while grouping duplicate keys.
enforce_limit()?;
}
}
Ok(Value::from_map(map))
@@ -302,16 +306,23 @@ fn urlquery_encode_object(
Err(_) => bail!(params[0].span().error("not a valid url query")),
};
for (key, value) in obj.iter() {
let key = ensure_string(name, &params[0], key)?;
match value {
Value::String(v) => {
url.query_pairs_mut().append_pair(key.as_ref(), v.as_ref());
}
_ => {
let values = ensure_string_collection(name, &params[0], value)?;
for v in values {
url.query_pairs_mut().append_pair(key.as_ref(), v);
{
let mut pairs = url.query_pairs_mut();
for (key, value) in obj.iter() {
let key = ensure_string(name, &params[0], key)?;
match value {
Value::String(v) => {
pairs.append_pair(key.as_ref(), v.as_ref());
// Guard encoded parameter growth when serializing string fields.
enforce_limit()?;
}
_ => {
let values = ensure_string_collection(name, &params[0], value)?;
for v in values {
pairs.append_pair(key.as_ref(), v);
// Guard encoded parameter growth when serializing multi-valued fields.
enforce_limit()?;
}
}
}
}
src/builtins/graph.rs
+51 -5
View File
@@ -5,7 +5,7 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::{ensure_args_count, ensure_object};
use crate::builtins::utils::{enforce_limit, ensure_args_count, ensure_object};
use crate::lexer::Span;
use crate::value::Value;
use crate::*;
@@ -28,8 +28,20 @@ fn reachable(span: &Span, params: &[Ref<Expr>], args: &[Value], strict: bool) ->
let mut worklist = vec![];
match &args[1] {
Value::Array(arr) => worklist.extend(arr.iter().cloned()),
Value::Set(set) => worklist.extend(set.iter().cloned()),
Value::Array(arr) => {
for node in arr.iter() {
worklist.push(node.clone());
// Guard worklist growth when seeding traversal from an array.
enforce_limit()?;
}
}
Value::Set(set) => {
for node in set.iter() {
worklist.push(node.clone());
// Guard worklist growth when seeding traversal from a set.
enforce_limit()?;
}
}
_ if strict => bail!(params[1].span().error("initial vertices must be array/set")),
_ => return Ok(Value::Undefined),
}
@@ -41,13 +53,27 @@ fn reachable(span: &Span, params: &[Ref<Expr>], args: &[Value], strict: bool) ->
}
match graph.get(&v) {
Some(Value::Array(arr)) => worklist.extend(arr.iter().cloned()),
Some(Value::Set(set)) => worklist.extend(set.iter().cloned()),
Some(Value::Array(arr)) => {
for neighbor in arr.iter() {
worklist.push(neighbor.clone());
// Guard worklist growth when enqueuing array neighbors.
enforce_limit()?;
}
}
Some(Value::Set(set)) => {
for neighbor in set.iter() {
worklist.push(neighbor.clone());
// Guard worklist growth when enqueuing set neighbors.
enforce_limit()?;
}
}
Some(_) => (),
_ => continue,
}
reachable.insert(v);
// Guard reachable set size as discovered vertices accumulate.
enforce_limit()?;
}
Ok(Value::from_set(reachable))
@@ -64,6 +90,8 @@ fn visit(
if s.as_ref() == "" {
if !path.is_empty() {
paths.insert(Value::from_array(path.clone()));
// Guard path result growth when terminating at empty edge.
enforce_limit()?;
}
return Ok(());
}
@@ -74,15 +102,23 @@ fn visit(
// Current node is not valid. Add path as is.
if !path.is_empty() {
paths.insert(Value::from_array(path.clone()));
// Guard path set growth when encountering missing nodes.
enforce_limit()?;
}
return Ok(());
}
if visited.contains(node) {
paths.insert(Value::from_array(path.clone()));
// Guard path set growth when detecting a cycle.
enforce_limit()?;
} else {
path.push(node.clone());
// Guard path stack growth while descending the graph.
enforce_limit()?;
visited.insert(node.clone());
// Guard visited set growth while marking nodes as seen.
enforce_limit()?;
let n = match neighbors {
Some(Value::Array(arr)) => {
for n in arr.iter().rev() {
@@ -104,6 +140,8 @@ fn visit(
// Current node has no neighbors.
if !path.is_empty() {
paths.insert(Value::from_array(path.clone()));
// Guard path set growth when recording leaf nodes.
enforce_limit()?;
}
}
@@ -150,11 +188,15 @@ fn walk_visit(path: &mut Vec<Value>, value: &Value, paths: &mut Vec<Value>) -> R
{
let path = Value::from_array(path.clone());
paths.push(Value::from_array([path, value.clone()].into()));
// Guard walk result growth when emitting a new path/value pair.
enforce_limit()?;
}
match value {
Value::Array(arr) => {
for (idx, elem) in arr.iter().enumerate() {
path.push(Value::from(idx));
// Guard path stack growth while traversing array members.
enforce_limit()?;
walk_visit(path, elem, paths)?;
path.pop();
}
@@ -162,6 +204,8 @@ fn walk_visit(path: &mut Vec<Value>, value: &Value, paths: &mut Vec<Value>) -> R
Value::Set(set) => {
for elem in set.iter() {
path.push(elem.clone());
// Guard path stack growth while traversing set members.
enforce_limit()?;
walk_visit(path, elem, paths)?;
path.pop();
}
@@ -169,6 +213,8 @@ fn walk_visit(path: &mut Vec<Value>, value: &Value, paths: &mut Vec<Value>) -> R
Value::Object(obj) => {
for (key, value) in obj.iter() {
path.push(key.clone());
// Guard path stack growth while traversing object entries.
enforce_limit()?;
walk_visit(path, value, paths)?;
path.pop();
}
src/builtins/net.rs
+13 -7
View File
@@ -14,7 +14,7 @@ use std::vec::Vec;
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::ensure_args_count;
use crate::builtins::utils::{enforce_limit, ensure_args_count};
use crate::lexer::Span;
use crate::value::Value;
@@ -134,17 +134,23 @@ fn _cidr_expand(cidr: Arc<str>) -> Result<Value> {
.parse::<IpNet>()
.map_err(|e| anyhow!("Error parsing {cidr}: {e}"))?;
let mut hosts: Vec<Value> = net
.hosts()
.map(|h| Value::String(h.to_string().into()))
.collect();
let mut hosts: Vec<Value> = Vec::new();
for host in net.hosts() {
hosts.push(Value::String(host.to_string().into()));
// Guard expanded host list growth while iterating over CIDR addresses.
enforce_limit()?;
}
// the IpNet library has some different behavior regarding CIDR expansion from the go implementation
// that OPA uses; it will exclude the IPv4 CIDR network address and broadcast address when the netmask < 31.
// Adjust accordingly for parity.
if matches!(net, IpNet::V4(_) if net.prefix_len() < 31) {
hosts.push(net.broadcast().to_string().into());
hosts.insert(0, net.network().to_string().into());
hosts.push(Value::String(net.broadcast().to_string().into()));
// Guard expanded host list growth when adding the broadcast address.
enforce_limit()?;
hosts.insert(0, Value::String(net.network().to_string().into()));
// Guard expanded host list growth when reintroducing the network address.
enforce_limit()?;
}
Ok(Value::Array(Arc::from(hosts)))
src/builtins/numbers.rs
+7 -1
View File
@@ -10,7 +10,7 @@
use crate::ast::{ArithOp, Expr, Ref};
use crate::builtins;
use crate::builtins::utils::{ensure_args_count, ensure_numeric};
use crate::builtins::utils::{enforce_limit, ensure_args_count, ensure_numeric};
use crate::lexer::Span;
use crate::number::Number;
use crate::value::Value;
@@ -112,8 +112,12 @@ fn range(span: &Span, params: &[Ref<Expr>], args: &[Value], strict: bool) -> Res
while v != v2 {
values.push(Value::from(v.clone()));
v.add_assign(&incr)?;
// Guard vector growth while we enumerate the range.
enforce_limit()?;
}
values.push(Value::from(v));
// Guard the last push before materializing the array.
enforce_limit()?;
Ok(Value::from_array(values))
}
@@ -152,6 +156,8 @@ fn range_step(span: &Span, params: &[Ref<Expr>], args: &[Value], strict: bool) -
while (v <= v2 && incr.is_positive()) || (v >= v2 && !incr.is_positive()) {
values.push(Value::from(v.clone()));
v.add_assign(&incr)?;
// Guard vector growth as the stepped range accumulates.
enforce_limit()?;
}
Ok(Value::from_array(values))
src/builtins/objects.rs
+53 -24
View File
@@ -5,7 +5,7 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::{ensure_args_count, ensure_array, ensure_object};
use crate::builtins::utils::{enforce_limit, ensure_args_count, ensure_array, ensure_object};
use crate::lexer::Span;
use crate::Rc;
use crate::Value;
@@ -34,13 +34,13 @@ pub fn register(m: &mut builtins::BuiltinsMap<&'static str, builtins::BuiltinFcn
}
}
fn json_filter_impl(v: &Value, filter: &Value) -> Value {
fn json_filter_impl(v: &Value, filter: &Value) -> Result<Value> {
let filters = match filter {
Value::Object(fields) if fields.len() == 1 && filter[&Value::Null] == Value::Null => {
return v.clone()
return Ok(v.clone())
}
Value::Object(fields) if !fields.is_empty() => fields,
_ => return v.clone(),
_ => return Ok(v.clone()),
};
match v {
@@ -51,53 +51,59 @@ fn json_filter_impl(v: &Value, filter: &Value) -> Value {
// TODO: support integer indexes?
if let Value::String(idx) = idx {
if let Ok(idx) = Value::from_json_str(idx) {
let item = json_filter_impl(&v[&idx], filter);
let item = json_filter_impl(&v[&idx], filter)?;
if item != Value::Undefined {
items.push(item);
// Guard array growth while filtering nested structures.
enforce_limit()?;
}
}
}
}
Value::from_array(items)
Ok(Value::from_array(items))
}
Value::Set(s) => {
let mut items = BTreeSet::new();
for (item, filter) in filters.iter() {
if s.contains(item) {
let item = json_filter_impl(item, filter);
let item = json_filter_impl(item, filter)?;
if item != Value::Undefined {
items.insert(item);
// Guard set growth when preserving matched entries.
enforce_limit()?;
}
}
}
Value::from_set(items)
Ok(Value::from_set(items))
}
Value::Object(_) => {
let mut items = BTreeMap::new();
for (key, filter) in filters.iter() {
let item = json_filter_impl(&v[key], filter);
let item = json_filter_impl(&v[key], filter)?;
if item != Value::Undefined {
items.insert(key.clone(), item);
// Guard map growth as filtered keys accumulate.
enforce_limit()?;
}
}
Value::from_map(items)
Ok(Value::from_map(items))
}
_ => Value::Undefined,
_ => Ok(Value::Undefined),
}
}
fn json_remove_impl(v: &Value, filter: &Value) -> Value {
fn json_remove_impl(v: &Value, filter: &Value) -> Result<Value> {
let filters = match filter {
Value::Object(fields) if !fields.is_empty() => fields,
_ => return v.clone(),
_ => return Ok(v.clone()),
};
if filter[&Value::Null] == Value::Null {
return Value::Undefined;
return Ok(Value::Undefined);
}
match v {
@@ -106,50 +112,62 @@ fn json_remove_impl(v: &Value, filter: &Value) -> Value {
for (idx, item) in a.iter().enumerate() {
let idx = Value::String(format!("{idx}").into());
if let Some(f) = filters.get(&idx) {
let v = json_remove_impl(item, f);
let v = json_remove_impl(item, f)?;
if v != Value::Undefined {
items.push(v);
// Guard array size while removing JSON paths.
enforce_limit()?;
}
} else {
// Retain the item.
items.push(item.clone());
// Guard array size while copying retained entries.
enforce_limit()?;
}
}
Value::from_array(items)
Ok(Value::from_array(items))
}
Value::Set(s) => {
let mut items = BTreeSet::new();
for item in s.iter() {
if let Some(f) = filters.get(item) {
let v = json_remove_impl(item, f);
let v = json_remove_impl(item, f)?;
if v != Value::Undefined {
items.insert(v);
// Guard set size during filtered retention.
enforce_limit()?;
}
} else {
// Retain the item.
items.insert(item.clone());
// Guard set size when keeping unmatched entries.
enforce_limit()?;
}
}
Value::from_set(items)
Ok(Value::from_set(items))
}
Value::Object(obj) => {
let mut items = BTreeMap::new();
for (key, value) in obj.iter() {
if let Some(f) = filters.get(key) {
let v = json_remove_impl(value, f);
let v = json_remove_impl(value, f)?;
if v != Value::Undefined {
items.insert(key.clone(), v);
// Guard map size as filtered properties accumulate.
enforce_limit()?;
}
} else {
items.insert(key.clone(), value.clone());
// Guard map size while copying retained properties.
enforce_limit()?;
}
}
Value::from_map(items)
Ok(Value::from_map(items))
}
_ => Value::Undefined,
_ => Ok(Value::Undefined),
}
}
@@ -170,9 +188,13 @@ fn merge_filters(
*vref = Value::new_object();
}
f = vref;
// Guard recursive filter construction as path objects materialize.
enforce_limit()?;
}
if let Ok(f) = f.as_object_mut() {
f.insert(Value::Null, Value::Null);
// Guard filter map growth when marking terminal entries.
enforce_limit()?;
};
filters = fc;
}
@@ -183,14 +205,21 @@ fn merge_filters(
let vref = match f {
Value::Object(obj) => {
let obj = Rc::make_mut(obj);
obj.entry(p.clone()).or_insert_with(Value::new_object)
let entry = obj.entry(p.clone()).or_insert_with(Value::new_object);
// Guard filter map growth when creating nested objects.
enforce_limit()?;
entry
}
_ => break,
};
f = vref;
// Guard recursive descent as additional path components attach.
enforce_limit()?;
}
if let Ok(f) = f.as_object_mut() {
f.insert(Value::Null, Value::Null);
// Guard filter map growth when sealing terminal markers.
enforce_limit()?;
};
filters = fc;
}
@@ -223,7 +252,7 @@ fn json_filter(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool)
}
}
Ok(json_filter_impl(&args[0], &filters))
json_filter_impl(&args[0], &filters)
}
fn json_remove(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> Result<Value> {
@@ -237,7 +266,7 @@ fn json_remove(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool)
_ => bail!(span.error(format!("`{name}` requires set/array argument").as_str())),
};
Ok(json_remove_impl(&args[0], &filters))
json_remove_impl(&args[0], &filters)
}
fn filter(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> Result<Value> {
src/builtins/opa.rs
+15 -13
View File
@@ -3,7 +3,7 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::ensure_args_count;
use crate::builtins::utils::{enforce_limit, ensure_args_count};
use crate::*;
use crate::lexer::Span;
@@ -83,27 +83,29 @@ fn opa_runtime(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool)
];
let features = &features[..features.len() - 1];
let mut feature_values = Vec::with_capacity(features.len());
for feature in features.iter() {
feature_values.push(Value::String((*feature).to_string().into()));
// Guard feature list growth while reporting enabled capabilities.
enforce_limit()?;
}
obj.insert(
Value::String("features".into()),
Value::from_array(
features
.iter()
.map(|f| Value::String(f.to_string().into()))
.collect(),
),
Value::from_array(feature_values),
);
let mut builtins: Vec<&&str> = builtins::BUILTINS.keys().collect();
builtins.sort();
let mut builtin_values = Vec::with_capacity(builtins.len());
for builtin in builtins.iter() {
builtin_values.push(Value::String((**builtin).to_string().into()));
// Guard builtin list growth while reporting registered functions.
enforce_limit()?;
}
obj.insert(
Value::String("builtins".into()),
Value::from_array(
builtins
.iter()
.map(|f| Value::String(f.to_string().into()))
.collect(),
),
Value::from_array(builtin_values),
);
Ok(Value::from_map(obj))
src/builtins/regex.rs
+32 -17
View File
@@ -4,7 +4,7 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::{ensure_args_count, ensure_numeric, ensure_string};
use crate::builtins::utils::{enforce_limit, ensure_args_count, ensure_numeric, ensure_string};
use crate::lexer::Span;
use crate::value::Value;
use crate::*;
@@ -56,20 +56,25 @@ fn find_all_string_submatch_n(
pattern
.captures_iter(&value)
.map(|capture| {
Value::from_array(
capture
.iter()
.map(|group| {
Value::String(match group {
Some(s) => s.as_str().into(),
_ => "".into(),
})
})
.collect(),
)
let groups = capture
.iter()
.map(|group| {
let value = Value::String(match group {
Some(s) => s.as_str().into(),
_ => "".into(),
});
// Guard match accumulation while adding each capture group.
enforce_limit()?;
Ok(value)
})
.collect::<Result<Vec<Value>>>()?;
let array = Value::from_array(groups);
// Guard outer match accumulation as nested arrays grow.
enforce_limit()?;
Ok(array)
})
.take(n)
.collect(),
.collect::<Result<Vec<Value>>>()?,
))
}
@@ -97,9 +102,14 @@ fn find_n(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> R
Ok(Value::from_array(
pattern
.find_iter(&value)
.map(|m| Value::String(m.as_str().into()))
.map(|m| {
let value = Value::String(m.as_str().into());
// Guard match accumulation while pushing each substring.
enforce_limit()?;
Ok(value)
})
.take(n)
.collect(),
.collect::<Result<Vec<Value>>>()?,
))
}
@@ -161,8 +171,13 @@ fn regex_split(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool)
Ok(Value::from_array(
pattern
.split(&value)
.map(|s| Value::String(s.into()))
.collect::<Vec<Value>>(),
.map(|s| {
let value = Value::String(s.into());
// Guard output accumulation as each split segment is emitted.
enforce_limit()?;
Ok(value)
})
.collect::<Result<Vec<Value>>>()?,
))
}
src/builtins/strings.rs
+18 -4
View File
@@ -10,7 +10,7 @@
use crate::ast::{Expr, Ref};
use crate::builtins;
use crate::builtins::utils::{
ensure_args_count, ensure_array, ensure_numeric, ensure_object, ensure_string,
enforce_limit, ensure_args_count, ensure_array, ensure_numeric, ensure_object, ensure_string,
ensure_string_collection,
};
use crate::lexer::Span;
@@ -125,6 +125,8 @@ fn indexof_n(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -
for (pos, (idx, _)) in s1.char_indices().enumerate() {
if s1[idx..].starts_with(s2.as_ref()) {
positions.push(Value::from(Number::from(pos)));
// Guard position vector growth while tracking matches.
enforce_limit()?;
}
}
Ok(Value::from_array(positions))
@@ -156,11 +158,23 @@ fn split(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> Re
let parts: Vec<Value> = if delimiter.as_ref() == "" {
// If delimiter is "", str::split returns a leading and trailing "" whereas Golang's split doesn't.
// Therefore avoid str::split and instead return each char as a Value::String.
s.chars().map(|c| Value::from(c.to_string())).collect()
s.chars()
.map(|c| {
let value = Value::from(c.to_string());
// Guard part accumulation when splitting into characters.
enforce_limit()?;
Ok(value)
})
.collect::<Result<Vec<Value>>>()?
} else {
s.split(delimiter.as_ref())
.map(|s| Value::String(s.into()))
.collect()
.map(|s| {
let value = Value::String(s.into());
// Guard part accumulation when splitting by delimiter.
enforce_limit()?;
Ok(value)
})
.collect::<Result<Vec<Value>>>()?
};
Ok(Value::from(parts))
src/builtins/utils.rs
+9
View File
@@ -13,6 +13,11 @@ use alloc::collections::{BTreeMap, BTreeSet};
use anyhow::{bail, Result};
#[inline]
pub fn enforce_limit() -> Result<()> {
crate::utils::limits::check_memory_limit_if_needed().map_err(anyhow::Error::new)
}
pub fn ensure_args_count(
span: &Span,
fcn: &'static str,
@@ -124,11 +129,15 @@ pub fn ensure_string_collection<'a>(fcn: &str, arg: &Expr, v: &'a Value) -> Resu
Value::Array(a) => {
for (idx, elem) in a.iter().enumerate() {
collection.push(ensure_string_element(fcn, arg, elem, idx)?);
// Enforce allocator limit while materializing the string collection.
enforce_limit()?;
}
}
Value::Set(s) => {
for (idx, elem) in s.iter().enumerate() {
collection.push(ensure_string_element(fcn, arg, elem, idx)?);
// Enforce allocator limit while materializing the string collection.
enforce_limit()?;
}
}
_ => {
src/engine.rs
+6 -1
View File
@@ -8,7 +8,7 @@ use crate::interpreter::*;
use crate::lexer::*;
use crate::parser::*;
use crate::scheduler::*;
use crate::utils::gather_functions;
use crate::utils::{gather_functions, limits};
use crate::value::*;
use crate::*;
use crate::{Extension, QueryResults};
@@ -131,6 +131,7 @@ impl Engine {
let source = Source::from_contents(path, rego)?;
let mut parser = self.make_parser(&source)?;
let module = Ref::new(parser.parse()?);
limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
Rc::make_mut(&mut self.modules).push(module.clone());
// if policies change, interpreter needs to be prepared again
self.prepared = false;
@@ -164,6 +165,7 @@ impl Engine {
let source = Source::from_file(path)?;
let mut parser = self.make_parser(&source)?;
let module = Ref::new(parser.parse()?);
limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
Rc::make_mut(&mut self.modules).push(module.clone());
// if policies change, interpreter needs to be prepared again
self.prepared = false;
@@ -942,6 +944,9 @@ impl Engine {
#[doc(hidden)]
fn prepare_for_eval(&mut self, enable_tracing: bool, for_target: bool) -> Result<()> {
// Fail fast if the engine already exceeds the global memory limit before evaluation work.
limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
self.interpreter.set_traces(enable_tracing);
// if the data/policies have changed or the interpreter has never been prepared
src/interpreter.rs
+27
View File
@@ -359,6 +359,18 @@ impl Interpreter {
self.builtins_cache.clear();
}
#[cfg(feature = "allocator-memory-limits")]
fn memory_check(&mut self) -> Result<()> {
let _ = self; // quiet clippy::unused_self; retained for symmetry with VM path
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| anyhow!(err))
}
#[cfg(not(feature = "allocator-memory-limits"))]
fn memory_check(&mut self) -> Result<()> {
let _ = self; // quiet clippy::unused_self; retained for symmetry with VM path
Ok(())
}
// Helper methods for working with ExprLookup
fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) -> Result<()> {
let module_idx = self.current_module_index;
@@ -1038,6 +1050,8 @@ impl Interpreter {
}
fn eval_stmt_impl(&mut self, stmt: &LiteralStmt, stmts: &[&LiteralStmt]) -> Result<bool> {
self.memory_check()?;
Ok(match &stmt.literal {
Literal::Expr { span, expr, .. } => {
let value = match expr.as_ref() {
@@ -1328,6 +1342,8 @@ impl Interpreter {
stmts: &[&LiteralStmt],
loops: &[HoistedLoop],
) -> Result<bool> {
self.memory_check()?;
if loops.is_empty() {
if let Some((first_stmt, tail_stmts)) = stmts.split_first() {
// Evaluate the current statement whose loop expressions have been hoisted.
@@ -1397,6 +1413,7 @@ impl Interpreter {
match loop_value {
Value::Array(items) => {
for item in items.iter() {
self.memory_check()?;
self.set_loop_var_value(loop_target_expr, item.clone())?;
if self.execute_destructuring_plan(&walk_plan, item)?
@@ -1476,6 +1493,7 @@ impl Interpreter {
match loop_value {
Value::Array(items) => {
for (idx, v) in items.iter().enumerate() {
self.memory_check()?;
self.set_loop_var_value(loop_target_expr, v.clone())?;
if self.execute_destructuring_plan(&index_plan, &Value::from(idx))?
@@ -1497,6 +1515,7 @@ impl Interpreter {
}
Value::Set(items) => {
for v in items.iter() {
self.memory_check()?;
self.set_loop_var_value(loop_target_expr, v.clone())?;
// For sets, index is also the value.
@@ -1516,6 +1535,7 @@ impl Interpreter {
}
Value::Object(obj) => {
for (k, v) in obj.iter() {
self.memory_check()?;
self.set_loop_var_value(loop_target_expr, v.clone())?;
// For objects, index is key.
if self.execute_destructuring_plan(&index_plan, k)? == Value::from(true) {
@@ -1975,6 +1995,8 @@ impl Interpreter {
let mut eval_success = true;
for (idx, stmt) in stmts.iter().enumerate() {
self.memory_check()?;
if !eval_success {
break;
}
@@ -2334,6 +2356,7 @@ impl Interpreter {
if name == "trace" {
if let (Some(traces), Value::String(msg)) = (&mut self.traces, &v) {
traces.push(msg.clone());
self.memory_check()?;
return Ok(Value::Bool(true));
}
}
@@ -2341,6 +2364,8 @@ impl Interpreter {
if let Some(cached_key) = cache {
self.builtins_cache.insert((cached_key, args), v.clone());
}
self.memory_check()?;
Ok(v)
}
@@ -3782,6 +3807,8 @@ impl Interpreter {
false => None,
};
self.memory_check()?;
// Store the query schedule for lookup during evaluation
self.query_schedule = Some(query_schedule);
src/lexer.rs
+192 -169
View File
@@ -20,6 +20,11 @@ use crate::Value;
use anyhow::{anyhow, bail, Result};
#[inline]
fn check_memory_limit() -> Result<()> {
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| anyhow!(err))
}
// Maximum column width to prevent overflow and catch pathological input.
// Lines exceeding this are likely minified/generated code or attack attempts.
const MAX_COL: u32 = 1024;
@@ -185,6 +190,8 @@ impl Source {
bail!("{file} exceeds maximum allowed line count {MAX_LINES}");
}
lines.push((start, end));
// Enforce the current global memory cap after recording each line span.
check_memory_limit()?;
start = i_u32.saturating_add(1);
}
prev_ch = ch;
@@ -197,14 +204,20 @@ impl Source {
bail!("{file} exceeds maximum allowed line count {MAX_LINES}");
}
lines.push((start, usize_to_u32(contents.len())?));
// Enforce the global limit after appending the final line span.
check_memory_limit()?;
} else if contents.is_empty() {
lines.push((0, 0));
// Enforce the global limit even for empty sources.
check_memory_limit()?;
} else {
let s = usize_to_u32(contents.len().saturating_sub(1))?;
if lines.len() >= MAX_LINES {
bail!("{file} exceeds maximum allowed line count {MAX_LINES}");
}
lines.push((s, s));
// Enforce the global limit after storing the trailing span.
check_memory_limit()?;
}
Ok(Self {
src: Rc::new(SourceInternal {
@@ -536,8 +549,8 @@ impl<'source> Lexer<'source> {
.get(start..end)
.ok_or_else(|| self.source.error(self.line, col, "invalid number span"))?;
match serde_json::from_str::<Value>(num_slice) {
Ok(_) => (),
let parsed_number = match serde_json::from_str::<Value>(num_slice) {
Ok(value) => value,
Err(e) => {
let serde_msg = &e.to_string();
let msg = match &serde_msg {
@@ -558,7 +571,11 @@ impl<'source> Lexer<'source> {
msg
)
}
}
};
// Enforce the global memory limit after serde allocates the temporary Value.
check_memory_limit()?;
drop(parsed_number);
Ok(Token(
TokenKind::Number,
@@ -601,6 +618,7 @@ impl<'source> Lexer<'source> {
// Guard against invalid span that would underflow end - 1
return Err(self.source.error(line, col, "invalid raw string span"));
}
check_memory_limit()?;
Ok(Token(
TokenKind::RawString,
Span {
@@ -710,6 +728,8 @@ impl<'source> Lexer<'source> {
}
}
check_memory_limit()?;
Ok(Token(
TokenKind::String,
Span {
@@ -778,6 +798,8 @@ impl<'source> Lexer<'source> {
let end = self.peek().0;
self.advance_col(usize_to_u32(end.saturating_sub(start))?)?;
check_memory_limit()?;
Ok(Token(
TokenKind::String,
Span {
@@ -846,174 +868,175 @@ impl<'source> Lexer<'source> {
let start_u32 = usize_to_u32(start)?;
let col = self.col;
match chr {
// Special case for - followed by digit which is a
// negative json number.
// . followed by digit is invalid number.
'-' | '.' if self.peekahead(1).1.is_ascii_digit() => {
self.read_number()
}
// grouping characters
'{' | '}' | '[' | ']' | '(' | ')' |
// arith operator
'+' | '-' | '*' | '/' | '%' |
// separators
',' | ';' | '.' => {
self.advance_col(1)?;
self.iter.next();
Ok(Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
}))
}
#[cfg(feature = "azure-rbac")]
// RBAC logical AND operator (&&)
'&' if self.enable_rbac_tokens && self.peekahead(1).1 == '&' => {
self.advance_col(2)?;
self.iter.next();
self.iter.next();
Ok(Token(TokenKind::AzureRbac(AzureRbacTokenKind::LogicalAnd), Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(2),
}))
}
#[cfg(feature = "azure-rbac")]
// RBAC logical OR operator (||)
'|' if self.enable_rbac_tokens && self.peekahead(1).1 == '|' => {
self.advance_col(2)?;
self.iter.next();
self.iter.next();
Ok(Token(TokenKind::AzureRbac(AzureRbacTokenKind::LogicalOr), Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(2),
}))
}
// Generic bin operators (when RBAC tokens not enabled or single & |)
'&' | '|' => {
self.advance_col(1)?;
self.iter.next();
Ok(Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
}))
}
':' => {
self.advance_col(1)?;
self.iter.next();
let mut end = start_u32.saturating_add(1);
if self.peek().1 == '=' || (self.peek().1 == ':' && self.double_colon_token) {
self.advance_col(1)?;
self.iter.next();
end = end.saturating_add(1);
}
Ok(Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end
}))
}
// < <= > >= = ==
'<' | '>' | '=' => {
self.advance_col(1)?;
self.iter.next();
if self.peek().1 == '=' {
self.advance_col(1)?;
self.iter.next();
};
Ok(Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: usize_to_u32(self.peek().0)?,
}))
}
'!' if self.peekahead(1).1 == '=' => {
self.advance_col(2)?;
self.iter.next();
self.iter.next();
Ok(Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: usize_to_u32(self.peek().0)?,
}))
}
#[cfg(feature = "azure-rbac")]
// RBAC @ token for attribute references
'@' if self.enable_rbac_tokens => {
self.advance_col(1)?;
self.iter.next();
Ok(Token(TokenKind::AzureRbac(AzureRbacTokenKind::At), Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
}))
}
'"' => self.read_string(),
#[cfg(feature = "azure-rbac")]
'\'' if self.allow_single_quoted_strings => self.read_single_quoted_string(),
'`' => self.read_raw_string(),
'\x00' => Ok(Token(TokenKind::Eof, Span {
source: self.source.clone(),
line:self.line,
col,
start: start_u32,
end: start_u32
})),
_ if chr.is_ascii_digit() => self.read_number(),
_ if chr.is_ascii_alphabetic() || chr == '_' => {
let mut ident = self.read_ident()?;
if ident.1.text() == "set" && self.peek().1 == '(' {
// set immediately followed by ( is treated as set( if
// the next token is ).
let state = (self.iter.clone(), self.line, self.col);
self.iter.next();
let token = match chr {
// Special case for - followed by digit which is a
// negative json number.
// . followed by digit is invalid number.
'-' | '.' if self.peekahead(1).1.is_ascii_digit() => self.read_number()?,
// grouping characters
'{' | '}' | '[' | ']' | '(' | ')' |
// arith operator
'+' | '-' | '*' | '/' | '%' |
// separators
',' | ';' | '.' => {
self.advance_col(1)?;
self.iter.next();
Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
})
}
#[cfg(feature = "azure-rbac")]
// RBAC logical AND operator (&&)
'&' if self.enable_rbac_tokens && self.peekahead(1).1 == '&' => {
self.advance_col(2)?;
self.iter.next();
self.iter.next();
Token(TokenKind::AzureRbac(AzureRbacTokenKind::LogicalAnd), Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(2),
})
}
#[cfg(feature = "azure-rbac")]
// RBAC logical OR operator (||)
'|' if self.enable_rbac_tokens && self.peekahead(1).1 == '|' => {
self.advance_col(2)?;
self.iter.next();
self.iter.next();
Token(TokenKind::AzureRbac(AzureRbacTokenKind::LogicalOr), Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(2),
})
}
// Generic bin operators (when RBAC tokens not enabled or single & |)
'&' | '|' => {
self.advance_col(1)?;
self.iter.next();
Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
})
}
':' => {
self.advance_col(1)?;
self.iter.next();
let mut end = start_u32.saturating_add(1);
if self.peek().1 == '=' || (self.peek().1 == ':' && self.double_colon_token) {
self.advance_col(1)?;
self.iter.next();
end = end.saturating_add(1);
}
Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end,
})
}
// < <= > >= = ==
'<' | '>' | '=' => {
self.advance_col(1)?;
self.iter.next();
if self.peek().1 == '=' {
self.advance_col(1)?;
self.iter.next();
};
Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: usize_to_u32(self.peek().0)?,
})
}
'!' if self.peekahead(1).1 == '=' => {
self.advance_col(2)?;
self.iter.next();
self.iter.next();
Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: usize_to_u32(self.peek().0)?,
})
}
#[cfg(feature = "azure-rbac")]
// RBAC @ token for attribute references
'@' if self.enable_rbac_tokens => {
self.advance_col(1)?;
self.iter.next();
Token(TokenKind::AzureRbac(AzureRbacTokenKind::At), Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
})
}
'"' => self.read_string()?,
#[cfg(feature = "azure-rbac")]
'\'' if self.allow_single_quoted_strings => self.read_single_quoted_string()?,
'`' => self.read_raw_string()?,
'\x00' => Token(TokenKind::Eof, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32,
}),
_ if chr.is_ascii_digit() => self.read_number()?,
_ if chr.is_ascii_alphabetic() || chr == '_' => {
let mut ident = self.read_ident()?;
if ident.1.text() == "set" && self.peek().1 == '(' {
// set immediately followed by ( is treated as set( if
// the next token is ).
let state = (self.iter.clone(), self.line, self.col);
self.iter.next();
// Check it next token is ).
let next_tok = self.next_token()?;
let is_setp = next_tok.1.text() == ")";
// Check it next token is ).
let next_tok = self.next_token()?;
let is_setp = next_tok.1.text() == ")";
// Restore state
(self.iter, self.line, self.col) = state;
// Restore state
(self.iter, self.line, self.col) = state;
if is_setp {
self.iter.next();
self.advance_col(1)?;
ident.1.end = ident.1.end.saturating_add(1);
}
}
Ok(ident)
}
_ if self.unknown_char_is_symbol => {
self.advance_col(1)?;
self.iter.next();
Ok(Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
}))
}
_ => Err(self.source.error(self.line, self.col, "invalid character"))
}
if is_setp {
self.iter.next();
self.advance_col(1)?;
ident.1.end = ident.1.end.saturating_add(1);
}
}
ident
}
_ if self.unknown_char_is_symbol => {
self.advance_col(1)?;
self.iter.next();
Token(TokenKind::Symbol, Span {
source: self.source.clone(),
line: self.line,
col,
start: start_u32,
end: start_u32.saturating_add(1),
})
}
_ => return Err(self.source.error(self.line, self.col, "invalid character")),
};
check_memory_limit()?;
Ok(token)
}
}
src/lib.rs
+7
View File
@@ -166,6 +166,13 @@ pub use compiled_policy::CompiledPolicy;
pub use engine::Engine;
pub use lexer::Source;
pub use policy_info::PolicyInfo;
pub use utils::limits::LimitError;
#[cfg(feature = "allocator-memory-limits")]
pub use utils::limits::{
check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters,
global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override,
thread_memory_flush_threshold,
};
pub use value::Value;
#[cfg(feature = "arc")]
src/number.rs
+1 -1
View File
@@ -20,7 +20,7 @@ use core::str::FromStr;
use anyhow::{anyhow, bail, Result};
use num_bigint::BigInt as NumBigInt;
#[cfg(not(feature = "std"))]
#[allow(unused)]
use num_traits::float::FloatCore;
use num_traits::{One, Signed, ToPrimitive, Zero};
src/parser.rs
+71 -6
View File
@@ -25,6 +25,11 @@ use core::str::FromStr;
use anyhow::{anyhow, bail, Result};
#[inline]
fn check_memory_limit() -> Result<()> {
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| anyhow!(err))
}
#[derive(Clone)]
pub struct Parser<'source> {
source: Source,
@@ -481,12 +486,18 @@ impl<'source> Parser<'source> {
let mut items = vec![];
if self.token_text() != "]" {
items.push(Ref::new(self.parse_in_expr()?));
// Guard array literal growth against allocator limits.
check_memory_limit()?;
while self.token_text() == "," {
self.next_token()?;
match self.token_text() {
"]" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => items.push(Ref::new(self.parse_in_expr()?)),
_ => {
items.push(Ref::new(self.parse_in_expr()?));
// Guard array literal growth against allocator limits.
check_memory_limit()?;
}
}
}
}
@@ -552,12 +563,18 @@ impl<'source> Parser<'source> {
if self.token_text() != ":" {
// Parse as set.
let mut items = vec![Ref::new(first)];
// Guard set literal growth against allocator limits.
check_memory_limit()?;
while self.token_text() == "," {
self.next_token()?;
match self.token_text() {
"}" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => items.push(Ref::new(self.parse_in_expr()?)),
_ => {
items.push(Ref::new(self.parse_in_expr()?));
// Guard set literal growth against allocator limits.
check_memory_limit()?;
}
}
}
self.expect("}", "while parsing set")?;
@@ -607,6 +624,8 @@ impl<'source> Parser<'source> {
let value = self.parse_in_expr()?;
item_span.end = self.end;
items.push((item_span, Ref::new(first), Ref::new(value)));
// Guard object literal growth against allocator limits.
check_memory_limit()?;
while self.token_text() == "," {
self.next_token()?;
@@ -624,6 +643,8 @@ impl<'source> Parser<'source> {
item_span.end = self.end;
items.push((item_span, Ref::new(key), Ref::new(value)));
// Guard object literal growth against allocator limits.
check_memory_limit()?;
}
self.expect("}", "while parsing object")?;
@@ -763,12 +784,18 @@ impl<'source> Parser<'source> {
let mut args = vec![];
if self.token_text() != ")" {
args.push(Ref::new(self.parse_in_expr()?));
// Guard call argument list against allocator limits.
check_memory_limit()?;
while self.token_text() == "," {
self.next_token()?;
match self.token_text() {
")" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => args.push(Ref::new(self.parse_in_expr()?)),
_ => {
args.push(Ref::new(self.parse_in_expr()?));
// Guard call argument list against allocator limits.
check_memory_limit()?;
}
}
}
}
@@ -1075,6 +1102,8 @@ impl<'source> Parser<'source> {
refr: Ref::new(refr),
r#as: Ref::new(r#as),
});
// Guard with-modifier list against allocator limits.
check_memory_limit()?;
}
Ok(modifiers)
}
@@ -1123,8 +1152,12 @@ impl<'source> Parser<'source> {
self.expect("some", "while parsing some-decl")?;
// parse any vars.
let mut vars = vec![self.tok.1.clone()];
let mut refs = vec![Ref::new(self.parse_ref()?)];
let first_var = self.tok.1.clone();
let first_ref = Ref::new(self.parse_ref()?);
let mut vars = vec![first_var];
let mut refs = vec![first_ref];
// Guard some-statement bindings against allocator limits.
check_memory_limit()?;
while self.token_text() == "," {
self.next_token()?;
@@ -1132,6 +1165,8 @@ impl<'source> Parser<'source> {
refs.push(Ref::new(self.parse_ref()?));
span.end = self.end;
vars.push(span);
// Guard some-statement bindings against allocator limits.
check_memory_limit()?;
}
if self.token_text() != "in" || !self.is_imported_future_keyword("in") {
@@ -1291,6 +1326,8 @@ impl<'source> Parser<'source> {
}
literals.push(stmt);
// Guard query literal accumulation against allocator limits.
check_memory_limit()?;
loop {
match self.token_text() {
@@ -1306,6 +1343,8 @@ impl<'source> Parser<'source> {
}
let stmt = self.parse_literal_stmt()?;
literals.push(stmt);
// Guard query literal accumulation against allocator limits.
check_memory_limit()?;
}
if !end_delim.is_empty() {
@@ -1529,12 +1568,18 @@ impl<'source> Parser<'source> {
let mut args = vec![];
if self.token_text() != ")" {
args.push(Ref::new(self.parse_term()?));
// Guard rule head arguments against allocator limits.
check_memory_limit()?;
while self.token_text() == "," {
self.next_token()?;
match self.token_text() {
")" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => args.push(Ref::new(self.parse_term()?)),
_ => {
args.push(Ref::new(self.parse_term()?));
// Guard rule head arguments against allocator limits.
check_memory_limit()?;
}
}
}
}
@@ -1647,6 +1692,8 @@ impl<'source> Parser<'source> {
assign,
query,
});
// Guard rule body accumulation against allocator limits.
check_memory_limit()?;
true
}
"if" => {
@@ -1665,6 +1712,8 @@ impl<'source> Parser<'source> {
assign,
query,
});
// Guard rule body accumulation against allocator limits.
check_memory_limit()?;
true
}
_ => false,
@@ -1690,6 +1739,8 @@ impl<'source> Parser<'source> {
assign: None,
query,
});
// Guard rule body accumulation against allocator limits.
check_memory_limit()?;
}
Ok(())
}
@@ -1722,6 +1773,8 @@ impl<'source> Parser<'source> {
assign,
query,
});
// Guard rule body accumulation against allocator limits.
check_memory_limit()?;
}
"{" => {
if self.rego_v1 {
@@ -1735,6 +1788,8 @@ impl<'source> Parser<'source> {
assign,
query,
});
// Guard rule body accumulation against allocator limits.
check_memory_limit()?;
}
_ if assign.is_none() => {
if self.token_text() == "if" {
@@ -1760,6 +1815,8 @@ impl<'source> Parser<'source> {
assign,
query,
});
// Guard rule body accumulation against allocator limits.
check_memory_limit()?;
break;
}
}
@@ -1782,6 +1839,8 @@ impl<'source> Parser<'source> {
bail!(arg.error("repeating parameter name"));
}
args.push(arg);
// Guard default rule parameters against allocator limits.
check_memory_limit()?;
if self.token_text() == ")" || self.tok.0 == TokenKind::Eof {
break;
}
@@ -1882,6 +1941,8 @@ impl<'source> Parser<'source> {
if comps.len() >= 2 && comps[0] == "future" && comps[1] == "keywords" {
imports.push(import);
// Guard import accumulation against allocator limits.
check_memory_limit()?;
return Ok(());
}
@@ -1914,6 +1975,8 @@ impl<'source> Parser<'source> {
}
imports.push(import);
// Guard import accumulation against allocator limits.
check_memory_limit()?;
Ok(())
}
@@ -2037,6 +2100,8 @@ impl<'source> Parser<'source> {
let mut policy = vec![];
while self.tok.0 != TokenKind::Eof {
policy.push(Ref::new(self.parse_rule()?));
// Guard policy rule accumulation against allocator limits.
check_memory_limit()?;
if self.token_text() == "__target__" {
bail!(self
.tok
src/rvm/vm/dispatch.rs
+1
View File
@@ -26,6 +26,7 @@ impl RegoVM {
program: &Program,
instruction: Instruction,
) -> Result<InstructionOutcome> {
self.memory_check()?;
self.execute_load_and_move(program, instruction)
}
src/rvm/vm/errors.rs
+3
View File
@@ -17,6 +17,9 @@ pub enum VmError {
pc: usize,
},
#[error("Execution stopped: exceeded maximum memory limit of {limit} bytes with usage {usage} bytes (pc={pc})")]
MemoryLimitExceeded { usage: u64, limit: u64, pc: usize },
#[error("Literal index {index} out of bounds (pc={pc})")]
LiteralIndexOutOfBounds { index: u16, pc: usize },
src/rvm/vm/execution.rs
+2
View File
@@ -127,6 +127,7 @@ impl RegoVM {
let target = self.convert_pc(target, "jump target")?;
self.pc = target;
while self.pc < program.instructions.len() {
self.memory_check()?;
if self.executed_instructions >= self.max_instructions {
return Err(VmError::InstructionLimitExceeded {
limit: self.max_instructions,
@@ -293,6 +294,7 @@ impl RegoVM {
fn run_stackless_loop(&mut self, program: &Program, last_result: &mut Value) -> Result<()> {
while !self.execution_stack.is_empty() {
self.memory_check()?;
self.frame_pc_overridden = false;
let should_finalize_rule = self.execution_stack.last().is_some_and(|frame| {
matches!(
src/rvm/vm/functions.rs
+10 -2
View File
@@ -21,7 +21,7 @@ impl RegoVM {
pc: self.pc,
available: self.program.instruction_data.function_call_params.len(),
})?;
match self.execution_mode {
let call_result = match self.execution_mode {
ExecutionMode::RunToCompletion => {
self.execute_call_rule_common(params.dest, params.func_rule_index, Some(&params))
}
@@ -30,7 +30,12 @@ impl RegoVM {
params.func_rule_index,
Some(&params),
),
}
};
call_result?;
self.memory_check()?;
Ok(())
}
pub(super) fn execute_builtin_call(&mut self, params_index: u16) -> Result<()> {
@@ -69,6 +74,7 @@ impl RegoVM {
if args.iter().any(|a| a == &Value::Undefined) {
self.set_register(params.dest, Value::Undefined)?;
self.memory_check()?;
return Ok(());
}
@@ -117,6 +123,8 @@ impl RegoVM {
if let Some(name) = cache_name {
self.builtins_cache.insert((name, args), result);
}
self.memory_check()?;
} else {
return Err(VmError::BuiltinNotResolved {
name: builtin_info.name.clone(),
src/rvm/vm/machine.rs
+23
View File
@@ -2,9 +2,12 @@
// Licensed under the MIT License.
use crate::rvm::program::Program;
#[cfg(feature = "allocator-memory-limits")]
use crate::utils::limits::{self, LimitError};
use crate::value::Value;
use crate::CompiledPolicy;
use alloc::collections::{btree_map::Entry, BTreeMap, VecDeque};
use alloc::format;
use alloc::string::String;
use alloc::sync::Arc;
use alloc::vec;
@@ -356,4 +359,24 @@ impl RegoVM {
*slot = value;
Ok(())
}
#[cfg(feature = "allocator-memory-limits")]
pub(super) fn memory_check(&mut self) -> Result<()> {
limits::check_memory_limit_if_needed().map_err(|err| match err {
LimitError::MemoryLimitExceeded { usage, limit } => VmError::MemoryLimitExceeded {
usage,
limit,
pc: self.pc,
},
other => VmError::Internal {
message: format!("unexpected limit error: {other}"),
pc: self.pc,
},
})
}
#[cfg(not(feature = "allocator-memory-limits"))]
pub(super) fn memory_check(&mut self) -> Result<()> {
Ok(())
}
}
src/utils.rs
+2
View File
@@ -8,6 +8,8 @@
clippy::as_conversions
)] // small arithmetic checks are intentional
pub mod limits;
use crate::ast::*;
use crate::builtins::*;
use crate::lexer::*;
src/utils/limits/error.rs
+68
View File
@@ -0,0 +1,68 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
#![allow(dead_code)]
use core::fmt;
use core::time::Duration;
/// Errors reported when execution time or memory ceilings are enforced.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum LimitError {
/// Reported when the execution timer observes elapsed time beyond the configured limit.
TimeLimitExceeded {
/// Elapsed work duration when the threshold was exceeded.
elapsed: Duration,
/// Configured time limit.
limit: Duration,
},
/// Reported when the memory tracker estimates usage beyond the configured limit.
MemoryLimitExceeded {
/// Estimated bytes in use when the limit was detected.
usage: u64,
/// Configured memory ceiling in bytes.
limit: u64,
},
}
impl fmt::Debug for LimitError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::TimeLimitExceeded { elapsed, limit } => f
.debug_struct("TimeLimitExceeded")
.field("elapsed", elapsed)
.field("limit", limit)
.finish(),
Self::MemoryLimitExceeded { usage, limit } => f
.debug_struct("MemoryLimitExceeded")
.field("usage", usage)
.field("limit", limit)
.finish(),
}
}
}
impl fmt::Display for LimitError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::TimeLimitExceeded { elapsed, limit } => {
let elapsed_ns = elapsed.as_nanos();
let limit_ns = limit.as_nanos();
write!(
f,
"execution exceeded time limit (elapsed={}ns, limit={}ns)",
elapsed_ns, limit_ns
)
}
Self::MemoryLimitExceeded { usage, limit } => {
write!(
f,
"execution exceeded memory limit (usage={} bytes, limit={} bytes)",
usage, limit
)
}
}
}
}
impl core::error::Error for LimitError {}
src/utils/limits/memory.rs
+213
View File
@@ -0,0 +1,213 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
use super::error::LimitError;
use core::cell::Cell;
use core::sync::atomic::{AtomicU64, Ordering};
use std::thread_local;
static GLOBAL_MEMORY_LIMIT: AtomicU64 = AtomicU64::new(u64::MAX);
// Maximum iteration count before forcing a global memory check.
const MEMORY_CHECK_STRIDE: u32 = 16;
// Pending per-thread allocator delta (bytes) that triggers a memory check. Chosen at 32 KiB to
// catch short bursts before they exceed typical entry budgets while still amortizing the atomic.
const MEMORY_CHECK_DELTA_BYTES: u64 = 32 * 1024;
thread_local! {
// Per-thread stride counter used to amortize global memory checks.
static MEMORY_CHECK_TICKS: Cell<u32> = const { Cell::new(0) };
}
/// Sets the global memory limit in bytes; `None` disables enforcement.
///
/// # Examples
///
/// ```rust
/// #![cfg(feature = "allocator-memory-limits")]
/// use regorus::{set_global_memory_limit, Engine, LimitError, Value};
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// const LIMIT: u64 = 32 * 1024;
/// set_global_memory_limit(Some(LIMIT));
///
/// let mut engine = Engine::new();
/// engine.add_policy(
/// "limit.rego".to_string(),
/// "package limit\nallow if input.blob != \"\"".to_string(),
/// )?;
///
/// let payload = format!("{{\"blob\":\"{}\"}}", "x".repeat(128 * 1024));
///
/// // Prepare input while the limit is relaxed so parsing succeeds on constrained builds.
/// set_global_memory_limit(Some(LIMIT * 32));
/// engine.set_input(Value::from_json_str(&payload)?);
///
/// // Tighten the budget and observe evaluation fail fast on the smaller limit.
/// set_global_memory_limit(Some(LIMIT));
/// let err = engine
/// .eval_query("data.limit.allow".to_string(), false)
/// .unwrap_err();
/// let limit = err.downcast_ref::<LimitError>().copied();
/// assert!(matches!(
/// limit,
/// Some(LimitError::MemoryLimitExceeded { limit: LIMIT, .. })
/// ));
///
/// // Raise the ceiling and the same evaluation succeeds.
/// set_global_memory_limit(Some(LIMIT * 32));
/// let result = engine.eval_query("data.limit.allow".to_string(), false)?;
/// assert_eq!(result.result.len(), 1);
/// # Ok(())
/// # }
/// ```
pub fn set_global_memory_limit(limit: Option<u64>) {
let value = limit.unwrap_or(u64::MAX);
GLOBAL_MEMORY_LIMIT.store(value, Ordering::Relaxed);
}
/// Flushes this thread's Regorus allocation counters into the global aggregates.
///
/// Regorus batches per-thread deltas to keep hot paths uncontended. Automatic publication occurs when
/// [`check_global_memory_limit`] runs, when the configured threshold (see
/// [`thread_memory_flush_threshold`]) is exceeded, or when allocation statistics are queried.
/// Workloads that burst and then idle—or threads that are about to terminate—can call this helper to
/// push their pending deltas immediately and avoid dropping buffered usage.
///
/// ```rust
/// #![cfg(feature = "allocator-memory-limits")]
/// use regorus::flush_thread_memory_counters;
///
/// flush_thread_memory_counters();
/// ```
///
/// Flush before a worker thread terminates:
/// ```no_run
/// # use regorus::flush_thread_memory_counters;
/// std::thread::spawn(|| {
/// // ... do work ...
/// flush_thread_memory_counters(); // publish before exiting
/// });
/// ```
pub fn flush_thread_memory_counters() {
mimalloc::limits::flush_thread_counters();
}
/// Configures the per-thread flush threshold in bytes.
///
/// Each thread buffers allocation deltas locally and publishes them automatically once the absolute
/// difference since the last flush exceeds this threshold. Passing `None` restores the default of
/// 1&nbsp;MiB. Setting the threshold to zero disables automatic flushing, requiring manual calls to
/// [`flush_thread_memory_counters()`]. Values larger than [`i64::MAX`] are saturated.
pub fn set_thread_flush_threshold_override(bytes: Option<u64>) {
mimalloc::limits::set_thread_flush_threshold(bytes);
}
/// Returns the current per-thread flush threshold in bytes, if automatic flushing is enabled.
///
/// When the threshold is disabled (zero or negative), `None` is returned. Otherwise the value
/// represents the absolute delta that will trigger an automatic flush.
///
/// ```rust
/// #![cfg(feature = "allocator-memory-limits")]
/// use regorus::{set_thread_flush_threshold_override, thread_memory_flush_threshold};
///
/// set_thread_flush_threshold_override(Some(256 * 1024));
/// assert_eq!(thread_memory_flush_threshold(), Some(256 * 1024));
///
/// set_thread_flush_threshold_override(Some(0));
/// assert_eq!(thread_memory_flush_threshold(), None);
/// ```
pub fn thread_memory_flush_threshold() -> Option<u64> {
mimalloc::limits::thread_flush_threshold()
}
/// Validates the current allocation usage against the configured global limit.
///
/// This helper consults the global memory limit and returns [`LimitError::MemoryLimitExceeded`]
/// if tracked usage exceeds the configured ceiling.
///
/// ```rust
/// #![cfg(feature = "allocator-memory-limits")]
/// use regorus::{check_global_memory_limit, set_global_memory_limit, LimitError};
///
/// const LIMIT: u64 = 32 * 1024;
/// set_global_memory_limit(Some(LIMIT));
/// let _buffer = vec![0u8; 128 * 1024];
///
/// let outcome = check_global_memory_limit();
/// assert!(matches!(
/// outcome,
/// Err(LimitError::MemoryLimitExceeded { limit: LIMIT, .. })
/// ));
///
/// set_global_memory_limit(Some(LIMIT * 32));
/// check_global_memory_limit().unwrap();
/// ```
pub fn check_global_memory_limit() -> Result<(), LimitError> {
if let Some(limit) = global_memory_limit() {
let (stats, _) = mimalloc::allocation_stats_snapshot();
let usage = stats.allocated;
if usage > limit {
return Err(LimitError::MemoryLimitExceeded { usage, limit });
}
}
Ok(())
}
/// Enforces the currently configured memory ceiling, if any.
#[inline]
pub fn enforce_memory_limit() -> Result<(), LimitError> {
check_global_memory_limit()
}
/// Performs a throttled global memory check, combining a lightweight stride counter
/// with the allocator's pending per-thread delta to avoid excessive atomics on hot paths.
/// The global limit is only evaluated when the stride expires or the pending delta crosses
/// the configured watermark.
pub(super) fn check_memory_limit_if_needed() -> Result<(), LimitError> {
if global_memory_limit().is_none() {
// Reset state when enforcement is disabled to avoid stale counters.
MEMORY_CHECK_TICKS.with(|ticks| ticks.set(0));
let _ = mimalloc::limits::take_thread_flushed_since_check_flag();
return Ok(());
}
// Inspect unflushed per-thread allocator usage to catch large bursts early.
let pending_delta = mimalloc::limits::thread_allocation_pending_delta().unsigned_abs() as u64;
let flush_hint = mimalloc::limits::take_thread_flushed_since_check_flag();
MEMORY_CHECK_TICKS.with(|ticks| {
let next = ticks.get().wrapping_add(1);
if flush_hint || next >= MEMORY_CHECK_STRIDE || pending_delta >= MEMORY_CHECK_DELTA_BYTES {
// Publish usage when either the stride or delta threshold is hit.
ticks.set(0);
check_global_memory_limit()
} else {
ticks.set(next);
Ok(())
}
})
}
/// Returns the currently-configured global memory limit in bytes, if any.
///
/// When [`set_global_memory_limit`] is called with `Some(value)`, that value is reported here.
/// Passing `None` to [`set_global_memory_limit`] removes the limit, causing this function to return
/// `None`.
///
/// # Examples
///
/// ```rust
/// #![cfg(feature = "allocator-memory-limits")]
/// use regorus::{global_memory_limit, set_global_memory_limit};
/// set_global_memory_limit(Some(123));
/// assert_eq!(global_memory_limit(), Some(123));
///
/// set_global_memory_limit(None);
/// assert_eq!(global_memory_limit(), None);
/// ```
pub fn global_memory_limit() -> Option<u64> {
let limit = GLOBAL_MEMORY_LIMIT.load(Ordering::Relaxed);
(limit != u64::MAX).then_some(limit)
}
src/utils/limits/mod.rs
+37
View File
@@ -0,0 +1,37 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
//! Helpers for cooperative execution time and memory limits.
#![allow(dead_code)]
mod error;
#[cfg(feature = "allocator-memory-limits")]
mod memory;
#[allow(unused_imports)]
pub use error::LimitError;
#[allow(unused_imports)]
#[cfg(feature = "allocator-memory-limits")]
pub use memory::{
check_global_memory_limit, enforce_memory_limit, flush_thread_memory_counters,
global_memory_limit, set_global_memory_limit, set_thread_flush_threshold_override,
thread_memory_flush_threshold,
};
#[cfg(feature = "allocator-memory-limits")]
#[inline]
pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
memory::check_memory_limit_if_needed()
}
#[cfg(not(feature = "allocator-memory-limits"))]
#[inline]
pub fn enforce_memory_limit() -> core::result::Result<(), LimitError> {
Ok(())
}
#[cfg(not(feature = "allocator-memory-limits"))]
#[inline]
pub fn check_memory_limit_if_needed() -> core::result::Result<(), LimitError> {
Ok(())
}
src/value.rs
+45 -4
View File
@@ -22,7 +22,7 @@ use core::convert::AsRef;
use core::str::FromStr;
use anyhow::{anyhow, bail, Result};
use serde::de::{self, Deserializer, MapAccess, SeqAccess, Visitor};
use serde::de::{self, Deserializer, Error as DeError, MapAccess, SeqAccess, Visitor};
use serde::ser::{SerializeMap, Serializer};
use serde::{Deserialize, Serialize};
@@ -70,6 +70,16 @@ pub enum Value {
Undefined,
}
#[inline]
fn enforce_limit_anyhow() -> Result<()> {
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| anyhow!(err))
}
#[inline]
fn enforce_limit_for<E: DeError>() -> core::result::Result<(), E> {
crate::utils::limits::check_memory_limit_if_needed().map_err(|err| E::custom(err.to_string()))
}
#[doc(hidden)]
impl Serialize for Value {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
@@ -199,6 +209,8 @@ impl<'de> Visitor<'de> for ValueVisitor {
let mut arr = vec![];
while let Some(v) = visitor.next_element()? {
arr.push(v);
// Enforce allocator limit while expanding a deserialized array.
enforce_limit_for::<V::Error>()?;
}
Ok(Value::from(arr))
}
@@ -218,8 +230,12 @@ impl<'de> Visitor<'de> for ValueVisitor {
}
let mut map = BTreeMap::new();
map.insert(key, value);
// Enforce allocator limit while expanding a deserialized object.
enforce_limit_for::<V::Error>()?;
while let Some((key, value)) = visitor.next_entry()? {
map.insert(key, value);
// Enforce allocator limit while expanding a deserialized object.
enforce_limit_for::<V::Error>()?;
}
Ok(Value::from(map))
} else {
@@ -334,7 +350,24 @@ impl Value {
/// # }
/// ```
pub fn from_json_str(json: &str) -> Result<Value> {
serde_json::from_str(json).map_err(anyhow::Error::msg)
match serde_json::from_str::<Value>(json) {
Ok(value) => Ok(value),
Err(err) => {
#[cfg(feature = "allocator-memory-limits")]
{
// Re-validate allocator limits when serde parsing fails to surface LimitError.
match crate::utils::limits::check_global_memory_limit() {
Err(limit_err) => Err(anyhow!(limit_err)),
Ok(_) => Err(anyhow!(err)),
}
}
#[cfg(not(feature = "allocator-memory-limits"))]
{
Err(anyhow!(err))
}
}
}
}
/// Deserialize a [`Value`] from a file containing JSON.
@@ -1305,6 +1338,8 @@ impl Value {
if let Value::Object(map) = self {
if map.get(&key).is_none() {
Rc::make_mut(map).insert(key.clone(), Value::Undefined);
// Enforce allocator limit while creating nested object entries.
enforce_limit_anyhow()?;
}
}
@@ -1330,7 +1365,9 @@ impl Value {
match (self, &mut new) {
(v @ Value::Undefined, _) => *v = new,
(Value::Set(ref mut set), Value::Set(new)) => {
Rc::make_mut(set).append(Rc::make_mut(new))
Rc::make_mut(set).append(Rc::make_mut(new));
// Enforce allocator limit after merging set entries.
enforce_limit_anyhow()?;
}
(Value::Object(map), Value::Object(new)) => {
for (k, v) in new.iter() {
@@ -1343,7 +1380,11 @@ impl Value {
serde_json::to_string_pretty(&v).map_err(anyhow::Error::msg)?,
)
}
_ => Rc::make_mut(map).insert(k.clone(), v.clone()),
_ => {
Rc::make_mut(map).insert(k.clone(), v.clone());
// Enforce allocator limit after merging object entries.
enforce_limit_anyhow()?;
}
};
}
}