Merge branch 'feat/audit-fix-0513' of https://github.com/zkemail/zk-r…

…egex into feat/audit-fix-0513
zkemail · May 17, 2024 · 25cc837 · 25cc837
2 parents 5421798 + d295713
commit 25cc837
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/packages/compiler/src/lib.rs b/packages/compiler/src/lib.rs
@@ -26,6 +26,8 @@ pub enum CompilerError {
     IoError(#[from] std::io::Error),
     #[error(transparent)]
     RegexError(#[from] fancy_regex::Error),
+    #[error("Generic error: {0}")]
+    GenericError(String),
 }
 
 /// A configuration of decomposed regexes.
@@ -91,7 +93,7 @@ pub struct SubstrsDefsJson {
 
 impl DecomposedRegexConfig {
     pub fn to_regex_and_dfa(&mut self) -> Result<RegexAndDFA, CompilerError> {
-        Ok(regex_and_dfa(self))
+        regex_and_dfa(self)
     }
 }
 

diff --git a/packages/compiler/src/regex.rs b/packages/compiler/src/regex.rs
@@ -1,3 +1,4 @@
+use super::CompilerError;
 use crate::{DFAGraph, DFAState, DecomposedRegexConfig, RegexAndDFA, RegexPartConfig, SubstrsDefs};
 use regex::Regex;
 use regex_automata::dfa::{dense::DFA, StartKind};
@@ -252,7 +253,9 @@ fn add_dfa(net_dfa: &DFAGraph, graph: &DFAGraph) -> DFAGraph {
     net_dfa
 }
 
-pub fn regex_and_dfa(decomposed_regex: &mut DecomposedRegexConfig) -> RegexAndDFA {
+pub fn regex_and_dfa(
+    decomposed_regex: &mut DecomposedRegexConfig,
+) -> Result<RegexAndDFA, CompilerError> {
     let mut config = DFA::config().minimize(true);
     config = config.start_kind(StartKind::Anchored);
     config = config.byte_classes(false);
@@ -316,15 +319,19 @@ pub fn regex_and_dfa(decomposed_regex: &mut DecomposedRegexConfig) -> RegexAndDF
             1 => regex.regex_def.ends_with("$"),
             2 => {
                 if idx == 0 && regex.regex_def.ends_with("$") {
-                    panic!("Invalid regex");
+                    return Err(CompilerError::GenericError(
+                        "Invalid regex, $ can only be at the end of the regex".to_string(),
+                    ));
                 }
                 idx == 1 && regex.regex_def.ends_with("$")
             }
             _ => match idx {
                 0 | _ if idx == decomposed_regex.parts.len() - 1 => {
                     if regex.regex_def.ends_with("$") {
                         if idx == 0 {
-                            panic!("Invalid regex");
+                            return Err(CompilerError::GenericError(
+                                "Invalid regex, $ can only be at the end of the regex".to_string(),
+                            ));
                         }
                         true
                     } else {
@@ -336,9 +343,14 @@ pub fn regex_and_dfa(decomposed_regex: &mut DecomposedRegexConfig) -> RegexAndDF
         };
         let re = DFA::builder()
             .configure(config.clone())
-            .build(&format!(r"^({})$", regex.regex_def.as_str()))
-            .unwrap();
-        let re_str = format!("{:?}", re);
+            .build(&format!(r"^({})$", regex.regex_def.as_str()));
+        if re.is_err() {
+            return Err(CompilerError::GenericError(format!(
+                "Failed to build DFA for regex: "{}", please check your regex",
+                regex.regex_def
+            )));
+        }
+        let re_str = format!("{:?}", re.unwrap());
         // println!("{:?}", re_str);
         let mut graph = dfa_to_graph(&parse_dfa_output(&re_str));
         if idx == 0 && caret_regex_index.is_some() {
@@ -435,15 +447,15 @@ pub fn regex_and_dfa(decomposed_regex: &mut DecomposedRegexConfig) -> RegexAndDF
         regex_str += &regex.regex_def;
     }
 
-    RegexAndDFA {
+    Ok(RegexAndDFA {
         regex_str: regex_str,
         dfa_val: net_dfa,
         end_anchor,
         substrs_defs: SubstrsDefs {
             substr_defs_array: substr_defs_array,
             substr_endpoints_array: None,
         },
-    }
+    })
 }
 
 pub fn dfa_from_regex_str(regex: &str) -> DFAGraph {

diff --git a/packages/compiler/src/wasm.rs b/packages/compiler/src/wasm.rs
@@ -31,7 +31,7 @@ pub fn genFromRaw(rawRegex: &str, substrsJson: &str, circomTemplateName: &str) -
 #[allow(non_snake_case)]
 pub fn genRegexAndDfa(decomposedRegex: JsValue) -> JsValue {
     let mut decomposed_regex_config: DecomposedRegexConfig = from_value(decomposedRegex).unwrap();
-    let regex_and_dfa = regex_and_dfa(&mut decomposed_regex_config);
+    let regex_and_dfa = regex_and_dfa(&mut decomposed_regex_config).unwrap();
     let dfa_val_str = serde_json::to_string(&regex_and_dfa).unwrap();
     JsValue::from_str(&dfa_val_str)
 }
@@ -42,6 +42,7 @@ pub fn genCircom(decomposedRegex: JsValue, circomTemplateName: &str) -> String {
     let mut decomposed_regex_config: DecomposedRegexConfig = from_value(decomposedRegex).unwrap();
     let regex_and_dfa = regex_and_dfa(&mut decomposed_regex_config);
     regex_and_dfa
+        .expect("failed to convert the decomposed regex to dfa")
         .gen_circom_str(&circomTemplateName)
         .expect("failed to generate circom")
 }