Repository: woodruffw/zizmor (rust)

Issue:
[BUG]: couldn't parse expression: fromJson(...)
### Pre-submission checks

- [X] I am **not** filing a feature request. These should be filed via the feature request form instead.
- [X] I have looked through the [open issues](https://github.com/woodruffw/zizmor/issues?q=is%3Aissue+is%3Aopen+) for a duplicate report.

### Expected behavior

I ran zizmor on https://github.com/nedbat/coveragepy/blob/master/.github/workflows/publish.yml.  Previous versions worked fine, but 0.10.0 produced a warning.

### Actual behavior

```
2024-12-19T19:16:32.907128Z  WARN audit:audit{workflow=file:///Users/ned/coverage/trunk/.github/workflows/publish.yml}: zizmor::audit::template_injection: couldn't parse expression: fromJson(steps.runs.outputs.data).workflow_runs[0].id
```

### Reproduction steps

- Clone https://github.com/nedbat/coveragepy
- Run zizmor on it.

### Logs

```plain text
% ~/.cargo/bin/zizmor --verbose coverage/trunk/.github/workflows/publish.yml
2024-12-20T12:10:22.028500Z DEBUG zizmor::state: using cache directory: "/Users/ned/.cache/zizmor"
2024-12-20T12:10:22.029274Z DEBUG zizmor::config: no config discovered; loading default
2024-12-20T12:10:22.029285Z DEBUG zizmor::config: loaded config: Config { rules: {} }
2024-12-20T12:10:22.036755Z  WARN audit:audit{workflow=file:///Users/ned/coverage/trunk/.github/workflows/publish.yml}: zizmor::audit::template_injection: couldn't parse expression: fromJson(steps.runs.outputs.data).workflow_runs[0].id
2024-12-20T12:10:22.046189Z  INFO audit: zizmor: 🌈 completed /Users/ned/coverage/trunk/.github/workflows/publish.yml
No findings to report. Good job!
```
```


### Additional context

_No response_

Pull request description:
fix: handle context accesses following calls
~~WIP;~~ These grammar changes mean the syntax tree itself is different, so I need to update the walk as well.

Fixes #339.

Patch:
diff --git a/src/audit/template_injection.rs b/src/audit/template_injection.rs
index be4262b..35e57d7 100644
--- a/src/audit/template_injection.rs
+++ b/src/audit/template_injection.rs
@@ -90,19 +90,15 @@ impl TemplateInjection {
             Expr::String(_) => true,
             Expr::Boolean(_) => true,
             Expr::Null => true,
-            // NOTE: Currently unreachable, since we churlishly consider
-            // indexing expressions unsafe and `Expr::Star` only occurs
-            // within indices at the moment.
-            Expr::Star => unreachable!(),
-            // NOTE: Some index operations may be safe, but for now
-            // we consider them all unsafe.
-            Expr::Index { .. } => false,
+            // NOTE: Currently unreachable, since these only occur
+            // within Expr::Context and we handle that at the top-level.
+            Expr::Star | Expr::Identifier(_) | Expr::Index(_) => unreachable!(),
             // NOTE: Some function calls may be safe, but for now
             // we consider them all unsafe.
             Expr::Call { .. } => false,
             // We consider all context accesses unsafe. This isn't true,
             // but our audit filters the safe ones later on.
-            Expr::Context(_) => false,
+            Expr::Context { .. } => false,
             Expr::BinOp { lhs, op, rhs } => {
                 match op {
                     // `==` and `!=` are always safe, since they evaluate to
diff --git a/src/expr/expr.pest b/src/expr/expr.pest
index 8846b68..2e3f9c9 100644
--- a/src/expr/expr.pest
+++ b/src/expr/expr.pest
@@ -46,8 +46,6 @@ primary_expr = {
   | string
   | boolean
   | null
-  | index
-  | function_call
   | context
   | "(" ~ primary_expr ~ ")"
 }
@@ -72,12 +70,10 @@ boolean = { "true" | "false" }
 null = { "null" }
 
 /// Context references (e.g., github.event.issue.number)
-context    = ${ identifier ~ ("." ~ (identifier | star))* }
+context    = ${ (function_call | identifier) ~ (("." ~ (identifier | star)) | index)* }
 star       =  { "*" }
 identifier = @{ (ASCII_ALPHA | "_" | "-") ~ (ASCII_ALPHANUMERIC | "_" | "-")* }
+index      =  { ("[" ~ (or_expr | star) ~ "]") }
 
 /// Function calls
-function_call = { identifier ~ "(" ~ (or_expr ~ ("," ~ or_expr)*)? ~ ")" }
-
-/// Index operations
-index = { (context | function_call | ("(" ~ primary_expr ~ ")")) ~ ("[" ~ (or_expr | star) ~ "]")+ }
+function_call = !{ identifier ~ "(" ~ (or_expr ~ ("," ~ or_expr)*)? ~ ")" }
diff --git a/src/expr/mod.rs b/src/expr/mod.rs
index 16c3039..92bf60d 100644
--- a/src/expr/mod.rs
+++ b/src/expr/mod.rs
@@ -38,30 +38,16 @@ pub(crate) enum Expr {
     Boolean(bool),
     /// The `null` literal.
     Null,
-    /// The `*` literal within an index.
+    /// The `*` literal within an index or context.
     Star,
-    /// An index operation.
-    ///
-    /// Three different kinds of expressions can be indexed:
-    ///
-    /// ```
-    /// functionCall[expr]
-    /// context.reference[expr]
-    /// (<arbitrary expression>)[expr]
-    /// ```
-    ///
-    /// Arbitrarily many nestings of indices are allowed,
-    /// e.g. `functionCall()[1][2][3]`.
-    Index {
-        parent: Box<Expr>,
-        indices: Vec<Expr>,
-    },
     /// A function call.
     Call { func: String, args: Vec<Expr> },
-    /// A context reference.
-    // TODO: This should probably be a vec of parts internally,
-    // to expose the individual component/star parts.
-    Context(String),
+    /// A context identifier component, e.g. `github` in `github.actor`.
+    Identifier(String),
+    /// A context index component, e.g. `[0]` in `foo[0]`.
+    Index(Box<Expr>),
+    /// A full context reference.
+    Context { raw: String, components: Vec<Expr> },
     /// A binary operation, either logical or arithmetic.
     BinOp {
         lhs: Box<Expr>,
@@ -78,31 +64,52 @@ impl Expr {
         Self::String(s.into()).into()
     }
 
+    pub(crate) fn ident(i: impl Into<String>) -> Self {
+        Self::Identifier(i.into())
+    }
+
+    pub(crate) fn context(r: impl Into<String>, components: impl Into<Vec<Expr>>) -> Self {
+        Self::Context {
+            raw: r.into(),
+            components: components.into(),
+        }
+    }
+
     /// Returns all of the contexts used in this expression, regardless
     /// of dataflow.
+    ///
+    /// **IMPORTANT**: This only returns "well-known" contexts, i.e. ones
+    /// that are defined at the top level. In other words, if an expression
+    /// contains both `foo.bar` and `foo().bar`, only the former will be
+    /// returned, as `foo()` can be anything and thus the `bar` access against
+    /// it has no particular meaning.
     pub(crate) fn contexts(&self) -> Vec<&str> {
         let mut contexts = vec![];
 
         match self {
-            Expr::Index { parent, indices } => {
-                contexts.extend(parent.contexts());
-
-                for index in indices {
-                    contexts.extend(index.contexts());
-                }
-            }
             Expr::Call { func: _, args } => {
                 for arg in args {
                     contexts.extend(arg.contexts());
                 }
             }
-            Expr::Context(ctx) => contexts.push(ctx.as_str()),
+            Expr::Context { raw, components } => {
+                if matches!(components[0], Expr::Call { .. }) {
+                    // If the context looks something like `foo(args).a.b.c`, then
+                    // we need to check each of `args` for contexts but skip
+                    // the trailing identifiers, since those aren't well-known.
+                    contexts.extend(components[0].contexts());
+                } else {
+                    // Otherwise, if the context looks like a normal context,
+                    // we include it in its entirety.
+                    contexts.push(raw)
+                }
+            }
             Expr::BinOp { lhs, op: _, rhs } => {
                 contexts.extend(lhs.contexts());
                 contexts.extend(rhs.contexts());
             }
             Expr::UnOp { op: _, expr } => contexts.extend(expr.contexts()),
-            Expr::Number(_) | Expr::String(_) | Expr::Boolean(_) | Expr::Null | Expr::Star => (),
+            _ => (),
         }
 
         contexts
@@ -236,18 +243,6 @@ impl Expr {
                 Rule::boolean => Ok(Expr::Boolean(pair.as_str().parse().unwrap()).into()),
                 Rule::null => Ok(Expr::Null.into()),
                 Rule::star => Ok(Expr::Star.into()),
-                Rule::index => {
-                    // (context | function (expr))[expr]+
-                    let mut pairs = pair.into_inner();
-
-                    Ok(Expr::Index {
-                        parent: parse_pair(pairs.next().unwrap())?,
-                        indices: pairs
-                            .map(|pair| parse_pair(pair).map(|e| *e))
-                            .collect::<Result<_, _>>()?,
-                    }
-                    .into())
-                }
                 Rule::function_call => {
                     let mut pairs = pair.into_inner();
 
@@ -262,8 +257,28 @@ impl Expr {
                     }
                     .into())
                 }
-                Rule::context => Ok(Expr::Context(pair.as_str().into()).into()),
-                r => panic!("fuck: {r:?}"),
+                Rule::identifier => Ok(Expr::ident(pair.as_str()).into()),
+                Rule::index => {
+                    Ok(Expr::Index(parse_pair(pair.into_inner().next().unwrap())?).into())
+                }
+                Rule::context => {
+                    let raw = pair.as_str().to_string();
+                    let pairs = pair.into_inner();
+
+                    let mut inner: Vec<Expr> = pairs
+                        .map(|pair| parse_pair(pair).map(|e| *e))
+                        .collect::<Result<_, _>>()?;
+
+                    // NOTE(ww): Annoying specialization: the `context` rule
+                    // wholly encloses the `function_call` rule, so we clean up
+                    // the AST slightly to turn `Context { Call }` into just `Call`.
+                    if inner.len() == 1 && matches!(inner[0], Expr::Call { .. }) {
+                        Ok(inner.remove(0).into())
+                    } else {
+                        Ok(Expr::context(raw, inner).into())
+                    }
+                }
+                r => panic!("unrecognized rule: {r:?}"),
             }
         }
 
@@ -365,6 +380,8 @@ mod tests {
             "(true == false) == true",
             "(true == (false || true && (true || false))) == true",
             "(github.actor != 'github-actions[bot]' && github.actor) == 'BrewTestBot'",
+            "foo()[0]",
+            "fromJson(steps.runs.outputs.data).workflow_runs[0].id",
         ];
 
         for case in cases {
@@ -410,24 +427,74 @@ mod tests {
                     args: vec![Expr::Number(1.0), Expr::Number(2.0), Expr::Number(3.0)],
                 },
             ),
-            ("foo.bar.baz", Expr::Context("foo.bar.baz".into())),
+            (
+                "foo.bar.baz",
+                Expr::context(
+                    "foo.bar.baz",
+                    [Expr::ident("foo"), Expr::ident("bar"), Expr::ident("baz")]
+                )
+            ),
             (
                 "foo.bar.baz[1][2]",
-                Expr::Index {
-                    parent: Expr::Context("foo.bar.baz".into()).into(),
-                    indices: vec![Expr::Number(1.0), Expr::Number(2.0)],
-                },
+                Expr::context(
+                    "foo.bar.baz[1][2]",
+                    [
+                        Expr::ident(
+                            "foo",
+                        ),
+                        Expr::ident(
+                            "bar",
+                        ),
+                        Expr::ident(
+                            "baz",
+                        ),
+                        Expr::Index(
+                            Expr::Number(
+                                1.0,
+                            ).into(),
+                        ),
+                        Expr::Index(
+                            Expr::Number(
+                                2.0,
+                            ).into(),
+                        ),
+                    ],
+                ),
             ),
             (
                 "foo.bar.baz[*]",
-                Expr::Index {
-                    parent: Expr::Context("foo.bar.baz".into()).into(),
-                    indices: vec![Expr::Star],
-                },
+                Expr::context(
+                    "foo.bar.baz[*]",
+                    [
+                        Expr::ident(
+                            "foo",
+                        ),
+                        Expr::ident(
+                            "bar",
+                        ),
+                        Expr::ident(
+                            "baz",
+                        ),
+                        Expr::Index(
+                            Expr::Star.into(),
+                        ),
+                    ],
+                ),
             ),
             (
                 "vegetables.*.ediblePortions",
-                Expr::Context("vegetables.*.ediblePortions".into()),
+                Expr::context(
+                    "vegetables.*.ediblePortions",
+                    vec![
+                        Expr::ident(
+                            "vegetables",
+                        ),
+                        Expr::Star,
+                        Expr::ident(
+                            "ediblePortions",
+                        ),
+                    ],
+                ),
             ),
             (
                 // Sanity check for our associativity: the top level Expr here
@@ -436,8 +503,12 @@ mod tests {
                 Expr::BinOp {
                     lhs: Expr::BinOp {
                         lhs: Expr::BinOp {
-                            lhs: Expr::Context(
-                                "github.ref".into(),
+                            lhs: Expr::context(
+                                "github.ref",
+                                [
+                                    Expr::ident("github"),
+                                    Expr::ident("ref"),
+                                ],
                             ).into(),
                             op: BinOp::Eq,
                             rhs: Expr::string("refs/heads/main"),
@@ -473,7 +544,7 @@ mod tests {
                     op: BinOp::Or, rhs: Expr::Boolean(false).into()
                     }.into()
                 }
-            )
+            ),
         ];
 
         for (case, expr) in cases {
@@ -483,8 +554,15 @@ mod tests {
 
     #[test]
     fn test_expr_contexts() {
-        let expr = Expr::parse("foo.bar && abc && d.e.f").unwrap();
+        let expr = Expr::parse(
+            "foo.bar && abc && d.e.f && andThis(should.work).except.this && but().not.this",
+        )
+        .unwrap();
+
+        assert_eq!(expr.contexts(), ["foo.bar", "abc", "d.e.f", "should.work"]);
+
+        let expr = Expr::parse("fromJson(steps.runs.outputs.data).workflow_runs[0].id").unwrap();
 
-        assert_eq!(expr.contexts(), ["foo.bar", "abc", "d.e.f"]);
+        assert_eq!(expr.contexts(), ["steps.runs.outputs.data"])
     }
 }


Test patch:
diff --git a/tests/snapshot.rs b/tests/snapshot.rs
index d1c6b05..74af2a4 100644
--- a/tests/snapshot.rs
+++ b/tests/snapshot.rs
@@ -257,6 +257,12 @@ fn template_injection() -> Result<()> {
         .workflow(workflow_under_test("template-injection/static-env.yml"))
         .run()?);
 
+    insta::assert_snapshot!(zizmor()
+        .workflow(workflow_under_test(
+            "template-injection/issue-339-repro.yml"
+        ))
+        .run()?);
+
     Ok(())
 }
 
diff --git a/tests/snapshots/snapshot__template_injection-6.snap b/tests/snapshots/snapshot__template_injection-6.snap
new file mode 100644
index 0000000..130f8c2
--- /dev/null
+++ b/tests/snapshots/snapshot__template_injection-6.snap
@@ -0,0 +1,18 @@
+---
+source: tests/snapshot.rs
+expression: "zizmor().workflow(workflow_under_test(\"template-injection/issue-339-repro.yml\")).run()?"
+snapshot_kind: text
+---
+info[template-injection]: code injection via template expansion
+  --> @@INPUT@@:27:9
+   |
+27 |         - name: "Record run id"
+   |           --------------------- info: this step
+28 |           id: run-id
+29 | /         run: |
+30 | |           echo "run-id=${{ fromJson(steps.runs.outputs.data).workflow_runs[0].id }}" >> "$GITHUB_OUTPUT"
+   | |_________________________________________________________________________________________________________- info: steps.runs.outputs.data may expand into attacker-controllable code
+   |
+   = note: audit confidence → Low
+
+1 finding: 0 unknown, 1 informational, 0 low, 0 medium, 0 high
diff --git a/tests/test-data/template-injection/issue-339-repro.yml b/tests/test-data/template-injection/issue-339-repro.yml
new file mode 100644
index 0000000..e2b1f13
--- /dev/null
+++ b/tests/test-data/template-injection/issue-339-repro.yml
@@ -0,0 +1,30 @@
+# minimized from https://github.com/woodruffw/zizmor/issues/339
+
+name: "Publish"
+
+on:
+  repository_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  find-run:
+    name: "Find latest kit.yml run"
+    runs-on: "ubuntu-latest"
+    outputs:
+      run-id: ${{ steps.run-id.outputs.run-id }}
+
+    steps:
+      - name: "Find latest kit.yml run"
+        id: runs
+        uses: octokit/request-action@dad4362715b7fb2ddedf9772c8670824af564f0d # v2.4.0
+        with:
+          route: GET /repos/nedbat/coveragepy/actions/workflows/kit.yml/runs
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: "Record run id"
+        id: run-id
+        run: |
+          echo "run-id=${{ fromJson(steps.runs.outputs.data).workflow_runs[0].id }}" >> "$GITHUB_OUTPUT"