From 5864ca0c8674ff05b2e9d509168371465510aa71 Mon Sep 17 00:00:00 2001
From: Karan Batavia <118820668+karan-batavia@users.noreply.github.com>
Date: Tue, 6 Aug 2024 19:40:49 +0530
Subject: [PATCH] Wrap all custom nodes in an exception handler (#1238)

wrap all custom nodes in an exception handler
---
 .../java/language/package.scala               | 14 ++++++++-----
 .../scala/ai/privado/semantic/Language.scala  | 20 +++++++++++--------
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/src/main/scala/ai/privado/languageEngine/java/language/package.scala b/src/main/scala/ai/privado/languageEngine/java/language/package.scala
index 744f0732f..4425057ae 100644
--- a/src/main/scala/ai/privado/languageEngine/java/language/package.scala
+++ b/src/main/scala/ai/privado/languageEngine/java/language/package.scala
@@ -25,26 +25,30 @@ package ai.privado.languageEngine.java
 
 import io.shiftleft.codepropertygraph.generated.{Cpg, EdgeTypes, NodeTypes}
 import io.shiftleft.codepropertygraph.generated.nodes.{AstNode, CfgNode, File, JavaProperty, Literal, MethodParameterIn}
-import overflowdb.traversal._
+import overflowdb.traversal.*
 
 import scala.jdk.CollectionConverters.IteratorHasAsScala
+import scala.util.Try
 
 package object language {
 
   implicit class NodeStarters(cpg: Cpg) {
     def property: Traversal[JavaProperty] =
-      cpg.graph.nodes(NodeTypes.JAVA_PROPERTY).asScala.cast[JavaProperty]
+      Try(cpg.graph.nodes(NodeTypes.JAVA_PROPERTY).asScala.cast[JavaProperty]).toOption
+        .getOrElse(Iterator.empty[JavaProperty])
   }
 
   implicit class StepsForProperty(val trav: Traversal[JavaProperty]) extends AnyVal {
 
-    def usedAt: Traversal[CfgNode] = trav.out(EdgeTypes.IS_USED_AT).cast[CfgNode]
-    def file: Traversal[File]      = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
+    def usedAt: Traversal[CfgNode] =
+      Try(trav.out(EdgeTypes.IS_USED_AT).cast[CfgNode]).toOption.getOrElse(Iterator.empty[CfgNode])
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
 
   }
 
   implicit class NodeTravToProperty(val trav: Traversal[AstNode]) {
-    def originalProperty: Traversal[JavaProperty] = trav.out(EdgeTypes.ORIGINAL_PROPERTY).cast[JavaProperty]
+    def originalProperty: Traversal[JavaProperty] =
+      Try(trav.out(EdgeTypes.ORIGINAL_PROPERTY).cast[JavaProperty]).toOption.getOrElse(Iterator.empty[JavaProperty])
   }
 
   implicit class NodeToProperty(val node: AstNode) {
diff --git a/src/main/scala/ai/privado/semantic/Language.scala b/src/main/scala/ai/privado/semantic/Language.scala
index 52f3e4c47..c10529e6b 100644
--- a/src/main/scala/ai/privado/semantic/Language.scala
+++ b/src/main/scala/ai/privado/semantic/Language.scala
@@ -63,21 +63,23 @@ object Language {
   }
 
   implicit class StepsForPropertyForDbNode(val trav: Traversal[DbNode]) extends AnyVal {
-    def file: Traversal[File] = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
 
   }
   implicit class StepsForPropertyForSqlQueryNode(val trav: Traversal[SqlQueryNode]) extends AnyVal {
-    def file: Traversal[File]             = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
-    def sqlTable: Traversal[SqlTableNode] = trav.out(EdgeTypes.AST).cast[SqlTableNode]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
+    def sqlTable: Traversal[SqlTableNode] =
+      Try(trav.out(EdgeTypes.AST).cast[SqlTableNode]).toOption.getOrElse(Iterator.empty[SqlTableNode])
   }
 
   implicit class StepsForPropertyForSqlTableNode(val trav: Traversal[SqlTableNode]) extends AnyVal {
-    def file: Traversal[File]               = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
-    def sqlColumn: Traversal[SqlColumnNode] = trav.out(EdgeTypes.AST).cast[SqlColumnNode]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
+    def sqlColumn: Traversal[SqlColumnNode] =
+      Try(trav.out(EdgeTypes.AST).cast[SqlColumnNode]).toOption.getOrElse(Iterator.empty[SqlColumnNode])
   }
 
   implicit class StepsForPropertyForSqlColumnNode(val trav: Traversal[SqlColumnNode]) extends AnyVal {
-    def file: Traversal[File] = trav.out(EdgeTypes.SOURCE_FILE).cast[File]
+    def file: Traversal[File] = Try(trav.out(EdgeTypes.SOURCE_FILE).cast[File]).toOption.getOrElse(Iterator.empty[File])
 
   }
 
@@ -97,12 +99,14 @@ object Language {
 
   implicit class NodeStarterForAndroidXmlLayoutNode(cpg: Cpg) {
     def androidXmlLayoutNode: Traversal[AndroidXmlLayoutNode] =
-      cpg.graph.nodes(NodeTypes.ANDROID_XML_LAYOUT_NODE).asScala.cast[AndroidXmlLayoutNode]
+      Try(cpg.graph.nodes(NodeTypes.ANDROID_XML_LAYOUT_NODE).asScala.cast[AndroidXmlLayoutNode]).toOption
+        .getOrElse(Iterator.empty[AndroidXmlLayoutNode])
   }
 
   implicit class NodeStarterForAndroidXmlPermissionNode(cpg: Cpg) {
     def androidXmlPermissionNode: Traversal[AndroidXmlPermissionNode] =
-      cpg.graph.nodes(NodeTypes.ANDROID_XML_PERMISSION_NODE).asScala.cast[AndroidXmlPermissionNode]
+      Try(cpg.graph.nodes(NodeTypes.ANDROID_XML_PERMISSION_NODE).asScala.cast[AndroidXmlPermissionNode]).toOption
+        .getOrElse(Iterator.empty[AndroidXmlPermissionNode])
   }
 
 }