From cdee44bbd11c718d76f3cd078877235cca6f964c Mon Sep 17 00:00:00 2001
From: Owen Mansel-Chan <owen-mc@github.com>
Date: Wed, 27 Oct 2021 10:31:03 +0100
Subject: [PATCH] Add barrier guard for comparison

---
 .../Security/CWE-326/InsufficientKeySize.ql   | 12 ++++++++
 .../CWE-326/InsufficientKeySize.expected      |  8 +++++
 .../Security/CWE-326/InsufficientKeySize.go   | 30 +++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/ql/src/Security/CWE-326/InsufficientKeySize.ql b/ql/src/Security/CWE-326/InsufficientKeySize.ql
index 6e1096f7708..6166d62c6a1 100644
--- a/ql/src/Security/CWE-326/InsufficientKeySize.ql
+++ b/ql/src/Security/CWE-326/InsufficientKeySize.ql
@@ -27,6 +27,18 @@ class RsaKeyTrackingConfiguration extends DataFlow::Configuration {
       c.getTarget().hasQualifiedName("crypto/rsa", "GenerateKey")
     )
   }
+
+  override predicate isBarrierGuard(DataFlow::BarrierGuard guard) {
+    guard instanceof ComparisonBarrierGuard
+  }
+}
+
+class ComparisonBarrierGuard extends DataFlow::BarrierGuard, DataFlow::RelationalComparisonNode {
+  override predicate checks(Expr e, boolean branch) {
+    exists(DataFlow::Node lesser , DataFlow::Node greater, int bias | this.leq(branch, lesser, greater, bias) |
+      globalValueNumber(DataFlow::exprNode(e)) = globalValueNumber(greater) and lesser.getIntValue() - bias >= 2048
+    )
+  }
 }
 
 from RsaKeyTrackingConfiguration cfg, DataFlow::PathNode source, DataFlow::PathNode sink
diff --git a/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.expected b/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.expected
index c7840f6adb5..f8c4e0998c8 100644
--- a/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.expected
+++ b/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.expected
@@ -2,6 +2,8 @@ edges
 | InsufficientKeySize.go:13:10:13:13 | 1024 : int | InsufficientKeySize.go:14:31:14:34 | size |
 | InsufficientKeySize.go:18:7:18:10 | 1024 : int | InsufficientKeySize.go:25:11:25:14 | definition of size : int |
 | InsufficientKeySize.go:25:11:25:14 | definition of size : int | InsufficientKeySize.go:26:31:26:34 | size |
+| InsufficientKeySize.go:30:13:30:16 | 1024 : int | InsufficientKeySize.go:32:32:32:38 | keyBits |
+| InsufficientKeySize.go:44:13:44:16 | 1024 : int | InsufficientKeySize.go:47:32:47:38 | keyBits |
 nodes
 | InsufficientKeySize.go:9:31:9:34 | 1024 | semmle.label | 1024 |
 | InsufficientKeySize.go:13:10:13:13 | 1024 : int | semmle.label | 1024 : int |
@@ -9,7 +11,13 @@ nodes
 | InsufficientKeySize.go:18:7:18:10 | 1024 : int | semmle.label | 1024 : int |
 | InsufficientKeySize.go:25:11:25:14 | definition of size : int | semmle.label | definition of size : int |
 | InsufficientKeySize.go:26:31:26:34 | size | semmle.label | size |
+| InsufficientKeySize.go:30:13:30:16 | 1024 : int | semmle.label | 1024 : int |
+| InsufficientKeySize.go:32:32:32:38 | keyBits | semmle.label | keyBits |
+| InsufficientKeySize.go:44:13:44:16 | 1024 : int | semmle.label | 1024 : int |
+| InsufficientKeySize.go:47:32:47:38 | keyBits | semmle.label | keyBits |
 #select
 | InsufficientKeySize.go:9:31:9:34 | 1024 | InsufficientKeySize.go:9:31:9:34 | 1024 | InsufficientKeySize.go:9:31:9:34 | 1024 | The size of this RSA key should be at least 2048 bits. |
 | InsufficientKeySize.go:14:31:14:34 | size | InsufficientKeySize.go:13:10:13:13 | 1024 : int | InsufficientKeySize.go:14:31:14:34 | size | The size of this RSA key should be at least 2048 bits. |
 | InsufficientKeySize.go:26:31:26:34 | size | InsufficientKeySize.go:18:7:18:10 | 1024 : int | InsufficientKeySize.go:26:31:26:34 | size | The size of this RSA key should be at least 2048 bits. |
+| InsufficientKeySize.go:32:32:32:38 | keyBits | InsufficientKeySize.go:30:13:30:16 | 1024 : int | InsufficientKeySize.go:32:32:32:38 | keyBits | The size of this RSA key should be at least 2048 bits. |
+| InsufficientKeySize.go:47:32:47:38 | keyBits | InsufficientKeySize.go:44:13:44:16 | 1024 : int | InsufficientKeySize.go:47:32:47:38 | keyBits | The size of this RSA key should be at least 2048 bits. |
diff --git a/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.go b/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.go
index ad1f1c8de75..d91d2cfa7cd 100644
--- a/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.go
+++ b/ql/test/query-tests/Security/CWE-326/InsufficientKeySize.go
@@ -25,3 +25,33 @@ func foo4() {
 func foo5(size int) {
 	rsa.GenerateKey(rand.Reader, size)
 }
+
+func foo6() {
+	keyBits := 1024
+	if keyBits >= 2047 {
+		rsa.GenerateKey(rand.Reader, keyBits) // BAD
+	}
+}
+
+func foo7() {
+	keyBits := 1024
+	if keyBits >= 2048 {
+		rsa.GenerateKey(rand.Reader, keyBits) // GOOD
+	}
+}
+
+func foo8() {
+	keyBits := 1024
+	switch {
+	case keyBits >= 2047:
+		rsa.GenerateKey(rand.Reader, keyBits) // BAD
+	}
+}
+
+func foo9() {
+	keyBits := 1024
+	switch {
+	case keyBits >= 2048:
+		rsa.GenerateKey(rand.Reader, keyBits) // GOOD
+	}
+}