Merge pull request #758 from DataDog/jf/K9VULN-8057

[K9VULN-8057] Fix "imports" parser to gracefully handle edge cases

Author: jasonforal

crates/static-analysis-kernel/src/analysis/ddsa_lib/js/flow/graph.rs

@@ -207,8 +207,9 @@ impl VertexId {
 
     /// Returns the kind of this vertex.
     pub fn kind(&self) -> VertexKind {
-        VertexKind::try_from_id(self.0 & Self::KIND_BIT_MASK)
-            .expect("js should serialize VertexId correctly")
+        let kind = VertexKind::try_from_id(self.0 & Self::KIND_BIT_MASK);
+        debug_assert!(kind.is_ok(), "js should serialize VertexId correctly");
+        kind.unwrap_or(VertexKind::Invalid)
     }
 
     /// Returns the internal node id for this vertex. This will return a [`ddsa_lib::common::NodeId`]
@@ -244,6 +245,7 @@ impl std::fmt::Display for VertexId {
 pub enum VertexKind {
     Cst = 0,
     Phi,
+    Invalid,
 }
 
 impl std::fmt::Display for VertexKind {
@@ -282,6 +284,7 @@ impl VertexKind {
         match self {
             VertexKind::Cst => Self::CST_STR,
             VertexKind::Phi => Self::PHI_STR,
+            VertexKind::Invalid => "<invalid>",
         }
     }
 
@@ -312,7 +315,11 @@ pub(crate) fn id_str(id: &dot_structures::Id) -> Cow<str> {
 
             while let Some(ch) = chars.next() {
                 unescaped.push(match ch {
-                    '\\' => chars.next().expect("string should never end on `\\`"),
+                    '\\' => {
+                        let next_char = chars.next();
+                        debug_assert!(next_char.is_some(), "string should never end on `\\`");
+                        next_char.unwrap_or_default()
+                    }
                     _ => ch,
                 });
             }
@@ -854,6 +861,7 @@ graph.adjacencyList;
             let js_const = match rust_kind {
                 VertexKind::Cst => "VERTEX_CST",
                 VertexKind::Phi => "VERTEX_PHI",
+                VertexKind::Invalid => unreachable!(),
             };
             let js_value = try_execute(scope, &format!("{};", js_const)).unwrap();
             assert!(js_value.is_number());

crates/static-analysis-kernel/src/analysis/ddsa_lib/js/flow/graph_test_utils.rs

@@ -249,6 +249,7 @@ impl NodeSearchAttrs {
                         panic!("phi node: unexpected attribute `{key}`");
                     }
                 }
+                VertexKind::Invalid => unreachable!(),
             }
         }
 
@@ -266,6 +267,7 @@ impl NodeSearchAttrs {
                 }
             }
             VertexKind::Phi => Self::Phi,
+            VertexKind::Invalid => unreachable!(),
         }
     }
 }

crates/static-analysis-kernel/src/analysis/ddsa_lib/ops.rs

@@ -12,9 +12,9 @@ use std::rc::Rc;
 #[op2(fast)]
 pub fn op_console_push(state: &mut OpState, #[string] line: &str) {
     let console = state.borrow::<Rc<RefCell<runtime::JsConsole>>>();
-    let mut console = console
-        .try_borrow_mut()
-        .expect("console should only be accessed via sequential executions");
+    let Ok(mut console) = console.try_borrow_mut() else {
+        unreachable!("parallel access of console is impossible");
+    };
     console.push(line);
 }
 
@@ -268,11 +268,12 @@ pub fn op_digraph_adjacency_list_to_dot(
                 let node_text = ts_node
                     .utf8_text(text.as_bytes())
                     .expect("bytes should be utf8 sequence");
-                LocatedNode::new_cst(ts_node, node_text)
+                Some(LocatedNode::new_cst(ts_node, node_text))
             }
-            VertexKind::Phi => LocatedNode::new_phi(vid.internal_id()),
+            VertexKind::Phi => Some(LocatedNode::new_phi(vid.internal_id())),
+            VertexKind::Invalid => None,
         };
-        Some(located.into())
+        located.map(Into::into)
     };
 
     V8DotGraph::try_new(scope, adjacency_list)

crates/static-analysis-kernel/src/analysis/languages/java/imports.rs

@@ -113,15 +113,23 @@ pub fn parse_imports_with_tree<'text>(
             }
         }
 
-        let package_node = package_node.expect("query invariant: value should be Some");
-        let target = target.expect("query invariant: value should be Some");
+        let Some(package_node) = package_node else {
+            debug_assert!(false, "query invariant: value should be Some");
+            continue;
+        };
+        let Some(target) = target else {
+            debug_assert!(false, "query invariant: value should be Some");
+            continue;
+        };
         // Due to the way the tree-sitter-java grammar is defined, when there is an asterisk, what
         // we've captured as the `import_child` will represent the full text of the package.
         // Otherwise, what we've captured as `package` will.
+        let Some(import_child_node) = import_child_node else {
+            debug_assert!(false, "query invariant: value should be Some");
+            continue;
+        };
         let package = match target {
-            ImportTarget::Wildcard => {
-                import_child_node.expect("query invariant: value should be Some")
-            }
+            ImportTarget::Wildcard => import_child_node,
             ImportTarget::Class(_) => package_node,
         };
         imports.push(Import { package, target });

crates/static-analysis-kernel/src/analysis/languages/javascript/imports.rs

@@ -216,8 +216,10 @@ pub(crate) fn parse_imports_with_tree_inner<'text>(
                 name = imported_from.take();
             }
 
+            debug_assert!(name.is_some(), "name should always be set");
+            let name = name?;
             Some(PackageImport {
-                name: name.expect("name should always be set"),
+                name,
                 imported_from,
             })
         })

crates/static-analysis-kernel/src/analysis/languages/python/imports.rs

@@ -104,9 +104,10 @@ impl<'a> Import<'a> {
     /// from ..common_utils.local_utils import print_array # Some(Import("common_utils"))
     /// ```
     pub fn parent_import(&self) -> Option<Import<'a>> {
-        self.full_text.rsplit_once('.').map(|(parent, _)| {
-            Import::try_new(parent, None, None, self.is_relative)
-                .expect("full_text should not have whitespace")
+        self.full_text.rsplit_once('.').and_then(|(parent, _)| {
+            let import = Import::try_new(parent, None, None, self.is_relative);
+            debug_assert!(import.is_ok());
+            import.ok()
         })
     }
 }