From 22d0600da8d8a58fba75ca358b7e439111af5441 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alvaro=20Mu=C3=B1oz?= <pwntester@github.com>
Date: Thu, 14 Mar 2024 13:28:39 +0100
Subject: [PATCH] Support more PR head checkouts

---
 ql/lib/codeql/actions/Ast.qll                            | 4 +++-
 ql/lib/codeql/actions/ast/internal/Ast.qll               | 2 ++
 ql/src/Security/CWE-829/UntrustedCheckout.ql             | 9 +++++++++
 .../Security/CWE-829/UntrustedCheckout.expected          | 8 ++++++++
 4 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/ql/lib/codeql/actions/Ast.qll b/ql/lib/codeql/actions/Ast.qll
index 70424a46f95..4a7ff12b4f9 100644
--- a/ql/lib/codeql/actions/Ast.qll
+++ b/ql/lib/codeql/actions/Ast.qll
@@ -232,7 +232,9 @@ abstract class SimpleReferenceExpression extends AstNode instanceof SimpleRefere
   AstNode getTarget() { result = super.getTarget() }
 }
 
-class StepsExpression extends SimpleReferenceExpression instanceof StepsExpressionImpl { }
+class StepsExpression extends SimpleReferenceExpression instanceof StepsExpressionImpl {
+  string getStepId() { result = super.getStepId() }
+}
 
 class NeedsExpression extends SimpleReferenceExpression instanceof NeedsExpressionImpl { }
 
diff --git a/ql/lib/codeql/actions/ast/internal/Ast.qll b/ql/lib/codeql/actions/ast/internal/Ast.qll
index 084474b4020..5c6ce37fa92 100644
--- a/ql/lib/codeql/actions/ast/internal/Ast.qll
+++ b/ql/lib/codeql/actions/ast/internal/Ast.qll
@@ -858,6 +858,8 @@ class StepsExpressionImpl extends SimpleReferenceExpressionImpl {
     this.getEnclosingJob() = result.getEnclosingJob() and
     result.(StepImpl).getId() = stepId
   }
+
+  string getStepId() { result = stepId }
 }
 
 /**
diff --git a/ql/src/Security/CWE-829/UntrustedCheckout.ql b/ql/src/Security/CWE-829/UntrustedCheckout.ql
index a24c80a2f60..f12f1102087 100644
--- a/ql/src/Security/CWE-829/UntrustedCheckout.ql
+++ b/ql/src/Security/CWE-829/UntrustedCheckout.ql
@@ -66,6 +66,7 @@ predicate containsHeadRef(string s) {
             "\\bgithub\\.event\\.check_run\\.pull_requests\\[\\d+\\]\\.head\\.sha\\b",
             "\\bgithub\\.event\\.check_run\\.pull_requests\\[\\d+\\]\\.id\\b",
             "\\bgithub\\.event\\.check_run\\.pull_requests\\[\\d+\\]\\.number\\b",
+            "\\bhead\\.sha\\b", "\\bhead\\.ref\\b"
           ], _, _)
   )
 }
@@ -80,6 +81,14 @@ class ActionsCheckout extends PRHeadCheckoutStep instanceof UsesStep {
     (
       containsHeadRef(this.getArgumentExpr("ref").getExpression())
       or
+      exists(StepsExpression e |
+        this.getArgumentExpr("ref") = e and
+        (
+          e.getStepId().matches(["%sha%", "%head%", "branch"]) or
+          e.getFieldName().matches(["%sha%", "%head%", "branch"])
+        )
+      )
+      or
       exists(UsesStep head |
         head.getCallee() = ["eficode/resolve-pr-refs", "xt0rted/pull-request-comment-branch"] and
         DataFlow::hasLocalFlowExpr(head, this.getArgumentExpr("ref"))
diff --git a/ql/test/query-tests/Security/CWE-829/UntrustedCheckout.expected b/ql/test/query-tests/Security/CWE-829/UntrustedCheckout.expected
index c7f4e4ad1c2..a6f02e7752a 100644
--- a/ql/test/query-tests/Security/CWE-829/UntrustedCheckout.expected
+++ b/ql/test/query-tests/Security/CWE-829/UntrustedCheckout.expected
@@ -10,4 +10,12 @@
 | .github/workflows/issue_comment_direct.yml:28:9:32:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
 | .github/workflows/issue_comment_direct.yml:35:9:40:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
 | .github/workflows/issue_comment_direct.yml:43:9:46:126 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_heuristic.yml:28:9:33:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_heuristic.yml:48:7:50:46 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_octokit.yml:26:9:30:6 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_octokit.yml:30:9:35:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_octokit.yml:57:9:62:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_octokit.yml:79:9:83:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_octokit.yml:95:9:100:2 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
+| .github/workflows/issue_comment_octokit.yml:109:9:114:66 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |
 | .github/workflows/untrusted_checkout.yml:9:7:13:4 | Uses Step | Potential unsafe checkout of untrusted pull request on privileged workflow. |