From 93a633750c45942cb607ce446954b9c8314ef010 Mon Sep 17 00:00:00 2001
From: Anand Krishnamoorthi <anakrish@microsoft.com>
Date: Thu, 18 Dec 2025 10:31:00 -0600
Subject: [PATCH] feat: Guard against runtime panics in lexer

Add guardrails for operations to ensure that they
won't panic at runtime.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
---
 src/lexer.rs       | 41 +++++++++++++++++++++++++++++++++++++----
 tests/lexer/mod.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/src/lexer.rs b/src/lexer.rs
index e01eb77..6b37e78 100644
--- a/src/lexer.rs
+++ b/src/lexer.rs
@@ -98,7 +98,13 @@ impl SourceStr {
     }
 
     pub fn text(&self) -> &str {
-        &self.source.contents()[self.start as usize..self.end as usize]
+        let start = self.start as usize;
+        let end = self.end as usize;
+        // Use safe slicing to avoid panics on malformed spans
+        self.source
+            .contents()
+            .get(start..end)
+            .unwrap_or("<invalid-span>")
     }
 
     pub fn clone_empty(&self) -> SourceStr {
@@ -241,7 +247,13 @@ pub struct Span {
 
 impl Span {
     pub fn text(&self) -> &str {
-        &self.source.contents()[self.start as usize..self.end as usize]
+        let start = self.start as usize;
+        let end = self.end as usize;
+        // Use safe slicing to avoid panics on malformed spans
+        self.source
+            .contents()
+            .get(start..end)
+            .unwrap_or("<invalid-span>")
     }
 
     pub fn source_str(&self) -> SourceStr {
@@ -445,7 +457,13 @@ impl<'source> Lexer<'source> {
         }
 
         // Ensure that the number is parsable in Rust.
-        match serde_json::from_str::<Value>(&self.source.contents()[start..end]) {
+        let num_slice = self
+            .source
+            .contents()
+            .get(start..end)
+            .ok_or_else(|| self.source.error(self.line, col, "invalid number span"))?;
+
+        match serde_json::from_str::<Value>(num_slice) {
             Ok(_) => (),
             Err(e) => {
                 let serde_msg = &e.to_string();
@@ -506,6 +524,10 @@ impl<'source> Lexer<'source> {
             }
         }
         let end = self.peek().0;
+        if end <= start {
+            // Guard against invalid span that would underflow end - 1
+            return Err(self.source.error(line, col, "invalid raw string span"));
+        }
         Ok(Token(
             TokenKind::RawString,
             Span {
@@ -574,8 +596,19 @@ impl<'source> Lexer<'source> {
         let end = self.peek().0;
         self.col += (end - start) as u32;
 
+        if start == 0 || end <= start {
+            // Reject invalid spans before slicing/serde to avoid panic
+            return Err(self.source.error(line, col, "invalid string span"));
+        }
+
+        let str_slice = self
+            .source
+            .contents()
+            .get(start - 1..end)
+            .ok_or_else(|| self.source.error(line, col, "invalid string span"))?;
+
         // Ensure that the string is parsable in Rust.
-        match serde_json::from_str::<String>(&self.source.contents()[start - 1..end]) {
+        match serde_json::from_str::<String>(str_slice) {
             Ok(_) => (),
             Err(e) => {
                 let serde_msg = &e.to_string();
diff --git a/tests/lexer/mod.rs b/tests/lexer/mod.rs
index ec8f8c6..d21c946 100644
--- a/tests/lexer/mod.rs
+++ b/tests/lexer/mod.rs
@@ -224,6 +224,34 @@ fn invalid_line() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn invalid_span_text_fallbacks() -> Result<()> {
+    let rego = "abc";
+    let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
+
+    let ss = SourceStr::new(source.clone(), 100, 200);
+    assert_eq!(
+        ss.text(),
+        "<invalid-span>",
+        "SourceStr should return fallback for out-of-bounds span"
+    );
+
+    let span = Span {
+        source: source.clone(),
+        line: 1,
+        col: 1,
+        start: 5,
+        end: 2,
+    };
+    assert_eq!(
+        span.text(),
+        "<invalid-span>",
+        "Span should return fallback for malformed span"
+    );
+
+    Ok(())
+}
+
 #[test]
 #[cfg(feature = "std")]
 fn file_more_than_64_kb_size() -> Result<()> {