Enable policy files greater than 64KB in size (#217)

fixes #214

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2024-04-26 01:04:21 -07:00
committed by GitHub
parent 744dad6126
commit 3743f32edc
8 changed files with 2711 additions and 69 deletions

View File

@@ -65,7 +65,7 @@ impl Engine {
/// ```
///
pub fn add_policy(&mut self, path: String, rego: String) -> Result<()> {
let source = Source::new(path, rego);
let source = Source::from_contents(path, rego)?;
let mut parser = Parser::new(&source)?;
self.modules.push(Ref::new(parser.parse()?));
// if policies change, interpreter needs to be prepared again
@@ -294,15 +294,15 @@ impl Engine {
self.interpreter.create_rule_prefixes()?;
let query_module = {
let source = Source::new(
let source = Source::from_contents(
"<query_module.rego>".to_owned(),
"package __internal_query_module".to_owned(),
);
)?;
Ref::new(Parser::new(&source)?.parse()?)
};
// Parse the query.
let query_source = Source::new("<query.rego>".to_string(), query);
let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
let mut parser = Parser::new(&query_source)?;
let query_node = parser.parse_user_query()?;
if query_node.span.text() == "data" {
@@ -411,15 +411,15 @@ impl Engine {
self.eval_modules(enable_tracing)?;
let query_module = {
let source = Source::new(
let source = Source::from_contents(
"<query_module.rego>".to_owned(),
"package __internal_query_module".to_owned(),
);
)?;
Ref::new(Parser::new(&source)?.parse()?)
};
// Parse the query.
let query_source = Source::new("<query.rego>".to_string(), query);
let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
let mut parser = Parser::new(&query_source)?;
let query_node = parser.parse_user_query()?;
let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;

View File

@@ -18,7 +18,7 @@ use anyhow::{anyhow, bail, Result};
struct SourceInternal {
pub file: String,
pub contents: String,
pub lines: Vec<(u16, u16)>,
pub lines: Vec<(u32, u32)>,
}
#[derive(Clone)]
@@ -61,8 +61,8 @@ impl Debug for Source {
#[derive(Clone)]
pub struct SourceStr {
source: Source,
start: u16,
end: u16,
start: u32,
end: u32,
}
impl Debug for SourceStr {
@@ -78,7 +78,7 @@ impl std::fmt::Display for SourceStr {
}
impl SourceStr {
pub fn new(source: Source, start: u16, end: u16) -> Self {
pub fn new(source: Source, start: u32, end: u32) -> Self {
Self { source, start, end }
}
@@ -116,39 +116,43 @@ impl std::cmp::Ord for SourceStr {
}
impl Source {
pub fn new(file: String, contents: String) -> Source {
pub fn from_contents(file: String, contents: String) -> Result<Source> {
let max_size = u32::MAX as usize - 2; // Account for rows, cols possibly starting at 1, EOF etc.
if contents.len() > max_size {
bail!("{file} exceeds maximum allowed policy file size {max_size}");
}
let mut lines = vec![];
let mut prev_ch = ' ';
let mut prev_pos = 0u16;
let mut start = 0u16;
let mut prev_pos = 0u32;
let mut start = 0u32;
for (i, ch) in contents.char_indices() {
if ch == '\n' {
let end = match prev_ch {
'\r' => prev_pos,
_ => i as u16,
_ => i as u32,
};
lines.push((start, end));
start = i as u16 + 1;
start = i as u32 + 1;
}
prev_ch = ch;
prev_pos = i as u16;
prev_pos = i as u32;
}
if (start as usize) < contents.len() {
lines.push((start, contents.len() as u16));
lines.push((start, contents.len() as u32));
} else if contents.is_empty() {
lines.push((0, 0));
} else {
let s = (contents.len() - 1) as u16;
let s = (contents.len() - 1) as u32;
lines.push((s, s));
}
Self {
Ok(Self {
src: Rc::new(SourceInternal {
file,
contents,
lines,
}),
}
})
}
pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Source> {
@@ -157,10 +161,7 @@ impl Source {
Err(e) => bail!("Failed to read {}. {e}", path.as_ref().display()),
};
// TODO: retain path instead of converting to string
Ok(Self::new(
path.as_ref().to_string_lossy().to_string(),
contents,
))
Self::from_contents(path.as_ref().to_string_lossy().to_string(), contents)
}
pub fn file(&self) -> &String {
@@ -169,7 +170,7 @@ impl Source {
pub fn contents(&self) -> &String {
&self.src.contents
}
pub fn line(&self, idx: u16) -> &str {
pub fn line(&self, idx: u32) -> &str {
let idx = idx as usize;
if idx < self.src.lines.len() {
let (start, end) = self.src.lines[idx];
@@ -179,7 +180,7 @@ impl Source {
}
}
pub fn message(&self, line: u16, col: u16, kind: &str, msg: &str) -> String {
pub fn message(&self, line: u32, col: u32, kind: &str, msg: &str) -> String {
if line as usize > self.src.lines.len() {
return format!("{}: invalid line {} specified", self.src.file, line);
}
@@ -206,7 +207,7 @@ impl Source {
)
}
pub fn error(&self, line: u16, col: u16, msg: &str) -> anyhow::Error {
pub fn error(&self, line: u32, col: u32, msg: &str) -> anyhow::Error {
anyhow!(self.message(line, col, "error", msg))
}
}
@@ -214,10 +215,10 @@ impl Source {
#[derive(Clone)]
pub struct Span {
pub source: Source,
pub line: u16,
pub col: u16,
pub start: u16,
pub end: u16,
pub line: u32,
pub col: u32,
pub start: u32,
pub end: u32,
}
impl Span {
@@ -272,8 +273,8 @@ pub struct Token(pub TokenKind, pub Span);
pub struct Lexer<'source> {
source: Source,
iter: Peekable<CharIndices<'source>>,
line: u16,
col: u16,
line: u32,
col: u32,
}
impl<'source> Lexer<'source> {
@@ -312,15 +313,15 @@ impl<'source> Lexer<'source> {
}
}
let end = self.peek().0;
self.col += (end - start) as u16;
self.col += (end - start) as u32;
Ok(Token(
TokenKind::Ident,
Span {
source: self.source.clone(),
line: self.line,
col,
start: start as u16,
end: end as u16,
start: start as u32,
end: end as u32,
},
))
}
@@ -363,7 +364,7 @@ impl<'source> Lexer<'source> {
}
let end = self.peek().0;
self.col += (end - start) as u16;
self.col += (end - start) as u32;
// Check for invalid number.Valid number cannot be followed by
// these characters:
@@ -403,8 +404,8 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line: self.line,
col,
start: start as u16,
end: end as u16,
start: start as u32,
end: end as u32,
},
))
}
@@ -440,8 +441,8 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line,
col,
start: start as u16,
end: end as u16 - 1,
start: start as u32,
end: end as u32 - 1,
},
))
}
@@ -453,7 +454,7 @@ impl<'source> Lexer<'source> {
let (start, _) = self.peek();
loop {
let (offset, ch) = self.peek();
let col = self.col + (offset - start) as u16;
let col = self.col + (offset - start) as u32;
match ch {
'"' | '\x00' => {
break;
@@ -468,7 +469,7 @@ impl<'source> Lexer<'source> {
'u' => {
for _i in 0..4 {
let (offset, ch) = self.peek();
let col = self.col + (offset - start) as u16;
let col = self.col + (offset - start) as u32;
if !ch.is_ascii_hexdigit() {
return Err(self.source.error(
line,
@@ -484,7 +485,7 @@ impl<'source> Lexer<'source> {
}
_ => {
// check for valid json chars
let col = self.col + (offset - start) as u16;
let col = self.col + (offset - start) as u32;
if !('\u{0020}'..='\u{10FFFF}').contains(&ch) {
return Err(self.source.error(line, col, "invalid character in string"));
}
@@ -499,7 +500,7 @@ impl<'source> Lexer<'source> {
self.iter.next();
let end = self.peek().0;
self.col += (end - start) as u16;
self.col += (end - start) as u32;
// Ensure that the string is parsable in Rust.
match serde_json::from_str::<String>(&self.source.contents()[start - 1..end]) {
@@ -522,8 +523,8 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line,
col: col + 1,
start: start as u16,
end: end as u16 - 1,
start: start as u32,
end: end as u32 - 1,
},
))
}
@@ -593,14 +594,14 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line: self.line,
col,
start: start as u16,
end: start as u16 + 1,
start: start as u32,
end: start as u32 + 1,
}))
}
':' => {
self.col += 1;
self.iter.next();
let mut end = start as u16 + 1;
let mut end = start as u32 + 1;
if self.peek().1 == '=' {
self.col += 1;
self.iter.next();
@@ -610,7 +611,7 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line: self.line,
col,
start: start as u16,
start: start as u32,
end
}))
}
@@ -626,8 +627,8 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line: self.line,
col,
start: start as u16,
end: self.peek().0 as u16,
start: start as u32,
end: self.peek().0 as u32,
}))
}
'!' if self.peekahead(1).1 == '=' => {
@@ -638,8 +639,8 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line: self.line,
col,
start: start as u16,
end: self.peek().0 as u16,
start: start as u32,
end: self.peek().0 as u32,
}))
}
'"' => self.read_string(),
@@ -648,8 +649,8 @@ impl<'source> Lexer<'source> {
source: self.source.clone(),
line:self.line,
col,
start: start as u16,
end: start as u16
start: start as u32,
end: start as u32
})),
_ if chr.is_ascii_digit() => self.read_number(),
_ if chr.is_ascii_alphabetic() || chr == '_' => {

View File

@@ -47,9 +47,9 @@ use std::rc::Rc;
#[derive(Debug, Clone, Serialize, Eq, PartialEq)]
pub struct Location {
/// Line number. Starts at 1.
pub row: u16,
pub row: u32,
/// Column number. Starts at 1.
pub col: u16,
pub col: u32,
}
/// An expression in a Rego query.

View File

@@ -16,8 +16,8 @@ pub struct Parser<'source> {
source: Source,
lexer: Lexer<'source>,
tok: Token,
line: u16,
end: u16,
line: u32,
end: u32,
future_keywords: BTreeMap<String, Span>,
rego_v1: bool,
}
@@ -753,7 +753,7 @@ impl<'source> Parser<'source> {
fn parse_membership_tail(
&mut self,
start: u16,
start: u32,
mut expr1: Expr,
mut expr2: Option<Expr>,
) -> Result<Expr> {

View File

@@ -38,7 +38,10 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> {
let mut sources = vec![];
let mut modules = vec![];
for (idx, _) in regos.iter().enumerate() {
sources.push(Source::new(format!("rego_{idx}"), regos[idx].clone()));
sources.push(Source::from_contents(
format!("rego_{idx}"),
regos[idx].clone(),
)?);
}
for source in &sources {

2620
tests/coco/data/large.rego Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -77,7 +77,7 @@ fn yaml_test_impl(file: &str) -> Result<()> {
let test: Test = serde_yaml::from_str(&yaml)?;
for case in &test.cases {
let source = Source::new("case.rego".to_string(), case.rego.clone());
let source = Source::from_contents("case.rego".to_string(), case.rego.clone())?;
print!("case {} ", &case.note);
match get_tokens(&source) {
@@ -162,7 +162,7 @@ fn run(path: &str) {
#[test]
fn debug() -> Result<()> {
let rego = "\"This string is 35 characters long.\"\"short string\"";
let source = Source::new("case.rego".to_string(), rego.to_string());
let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
let mut lexer = Lexer::new(&source);
let tok = lexer.next_token()?;
@@ -184,7 +184,7 @@ fn debug() -> Result<()> {
#[test]
fn tab() -> Result<()> {
let rego = r#" "This string is 35 characters long."`raw string`p"#;
let source = Source::new("case.rego".to_string(), rego.to_string());
let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
let mut lexer = Lexer::new(&source);
@@ -214,7 +214,7 @@ fn tab() -> Result<()> {
#[test]
fn invalid_line() -> Result<()> {
let rego = "";
let source = Source::new("case.rego".to_string(), rego.to_string());
let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
assert_eq!(
source.message(2, 0, "", ""),
@@ -223,3 +223,21 @@ fn invalid_line() -> Result<()> {
Ok(())
}
#[test]
fn file_more_than_64_kb_size() -> Result<()> {
let source = Source::from_file("tests/coco/data/large.rego")?;
let mut lexer = Lexer::new(&source);
let mut count = 0;
// Read tokens until EOF.
loop {
let token = lexer.next_token()?;
count += 1;
if token.0 == TokenKind::Eof {
break;
}
}
assert_eq!(count, 8789);
Ok(())
}

View File

@@ -630,7 +630,7 @@ fn yaml_test_impl(file: &str) -> Result<()> {
for case in &test.cases {
print!("\ncase {} ", case.note);
let source = Source::new("case.rego".to_string(), case.rego.clone());
let source = Source::from_contents("case.rego".to_string(), case.rego.clone())?;
let mut parser = Parser::new(&source)?;
match parser.parse() {
Ok(module) => {