Support more PR head checkouts

Author: Alvaro Muñoz

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 { }
 

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 }
 }
 
 /**

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"))

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. |