From 3251fb5c07047218ff0be7d306cb4560dd27f7d2 Mon Sep 17 00:00:00 2001
From: Your Name <name@gmail.ccom>
Date: Mon, 18 Jan 2021 02:37:53 +0300
Subject: [PATCH] updated

---
 ql/src/experimental/CWE-369/DivideByZero.ql | 50 ++++++++++++---------
 1 file changed, 30 insertions(+), 20 deletions(-)

diff --git a/ql/src/experimental/CWE-369/DivideByZero.ql b/ql/src/experimental/CWE-369/DivideByZero.ql
index 741f0155195..a7a16b2f9c5 100644
--- a/ql/src/experimental/CWE-369/DivideByZero.ql
+++ b/ql/src/experimental/CWE-369/DivideByZero.ql
@@ -12,13 +12,23 @@
 
 import go
 import DataFlow::PathGraph
+import semmle.go.dataflow.internal.TaintTrackingUtil
 
-class DivideByZeroSanitizeGuard extends DataFlow::BarrierGuard, DataFlow::EqualityTestNode {
+class DivideByZeroSanitizeGuard extends DataFlow::BarrierGuard {
   override predicate checks(Expr e, boolean branch) {
-    exists(DataFlow::Node zero, DataFlow::Node sink |
+    exists(
+      DataFlow::Node zero, DataFlow::Node sink, DataFlow::EqualityTestNode eqNode,
+      DataFlow::RelationalComparisonNode compNode
+    |
       zero.getNumericValue() = 0 and
-      sink.getType().getUnderlyingType() instanceof SignedIntegerType and
-      this.eq(branch.booleanNot(), sink, zero) and
+      (
+        sink.getType().getUnderlyingType() instanceof SignedIntegerType or
+        sink.getType().getUnderlyingType() instanceof UnsignedIntegerType
+      ) and
+      (
+        eqNode.eq(branch.booleanNot(), sink, zero) or
+        compNode.leq(branch.booleanNot(), sink, zero, 0)
+      ) and
       globalValueNumber(DataFlow::exprNode(e)) = globalValueNumber(sink)
     )
   }
@@ -27,31 +37,31 @@ class DivideByZeroSanitizeGuard extends DataFlow::BarrierGuard, DataFlow::Equali
 class DivideByZeroCheckConfig extends TaintTracking::Configuration {
   DivideByZeroCheckConfig() { this = "DivideByZeroCheckConfig" }
 
-  override predicate isSource(DataFlow::Node source) {
-    exists(DataFlow::CallNode c, IntegerParser::Range ip |
-      c.getTarget() = ip and source = c.getResult(0)
-    )
-    or
-    exists(IntegerType integerType | source.getType().getUnderlyingType() = integerType)
-  }
+  override predicate isSource(DataFlow::Node source) { source instanceof UntrustedFlowSource }
 
-  override predicate isSink(DataFlow::Node sink) {
-    exists(IntegerType integerType, QuoExpr e |
-      sink.asExpr().getParent().(QuoExpr).getRightOperand() = e.getAnOperand() and
-      not sink.asExpr().getParent().(QuoExpr).getRightOperand().isConst() and
-      sink.getType().getUnderlyingType() = integerType
+  override predicate isAdditionalTaintStep(DataFlow::Node node1, DataFlow::Node node2) {
+    exists(Function f |
+      (
+        f.getName() = "Atoi" or
+        f.getName() = "ParseInt" or
+        f.getName() = "ParseUint"
+      ) and
+      node1 = f.getACall().getArgument(0) and
+      node2 = f.getACall().getResult(0)
     )
   }
 
   override predicate isSanitizerGuard(DataFlow::BarrierGuard guard) {
     guard instanceof DivideByZeroSanitizeGuard
   }
+
+  override predicate isSink(DataFlow::Node sink) {
+    exists(QuoExpr e | sink.asExpr().getParent().(QuoExpr).getRightOperand() = e.getAnOperand())
+  }
 }
 
-from
-  DataFlow::PathNode source, DataFlow::PathNode sink, DivideByZeroCheckConfig cfg,
-  DataFlow::CallNode call
-where cfg.hasFlowPath(source, sink) and call.getResult(0) = source.getNode()
+from DataFlow::PathNode source, DataFlow::PathNode sink, DivideByZeroCheckConfig cfg
+where cfg.hasFlowPath(source, sink)
 select sink, source, sink,
   "Variable $@, which is used at division statement might be zero and leads to division by zero exception.",
   sink, sink.getNode().toString()