From ef16e645b7fb293390e335a7cd2a227966b1839e Mon Sep 17 00:00:00 2001
From: sadeli413 <53066432+sadeli413@users.noreply.github.com>
Date: Mon, 19 Feb 2024 15:46:28 +0000
Subject: [PATCH] Support raw byte grammar in NautilusContext (#1868)

* Support raw byte grammar in NautilusContext

* Add `use` statement in doc comment
---
 libafl/src/generators/nautilus.rs | 41 +++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/libafl/src/generators/nautilus.rs b/libafl/src/generators/nautilus.rs
index fb842000d5..0514d0493a 100644
--- a/libafl/src/generators/nautilus.rs
+++ b/libafl/src/generators/nautilus.rs
@@ -39,6 +39,47 @@ impl NautilusContext {
         Self { ctx }
     }
 
+    /// Returns a new [`NautilusContext`] with support for non UTF-8 rules.
+    ///
+    /// - This has the same behaviour as [`NautilusContext::new`] but returns `None` if the `rules` are empty.
+    /// - The starting rule is `rules[0]`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use libafl::generators::nautilus::NautilusContext;
+    ///
+    /// // Create a simple grammar for a series of null-terminated data.
+    /// let null = vec![0];
+    /// let mut rules = vec![
+    ///     // rules[0] is considered the starting rule
+    ///     ("SERIES", "{DATA}{NULL}".as_bytes()),
+    ///     ("SERIES", "{SERIES}{SERIES}".as_bytes()),
+    ///     ("DATA", "".as_bytes()),
+    ///     ("DATA", "{BYTE}{DATA}".as_bytes()),
+    ///     ("NULL", &null),
+    /// ];
+    ///
+    /// let bytes: Vec<u8> = (1..=u8::MAX).collect();
+    /// for i in 0..bytes.len() {
+    ///     rules.push(("BYTE", &bytes[i..=i]));
+    /// }
+    ///
+    /// let context = NautilusContext::with_rules(100, &rules).unwrap();
+    /// ```
+    #[must_use]
+    pub fn with_rules(tree_depth: usize, rules: &[(&str, &[u8])]) -> Option<Self> {
+        let mut ctx = Context::new();
+        for (symbol, rule) in rules {
+            ctx.add_rule(symbol, rule);
+        }
+
+        let root = format!("{{{}}}", rules.first()?.0);
+        ctx.add_rule("START", root.as_bytes());
+        ctx.initialize(tree_depth);
+        Some(Self { ctx })
+    }
+
     /// Create a new [`NautilusContext`] from a file
     #[must_use]
     pub fn from_file<P: AsRef<Path>>(tree_depth: usize, grammar_file: P) -> Self {