profiling jwt verification using io.jwt.decode_verify shows multiple invocations

[alam-chime]

I'm seeing inconsistent policy evaluation latencies with with jwt verification using io.jwt.decode_verify.
Below is a simple example to test io.jwt.decode_verify call.

policy.rego

package authz

default allow := false

jwks := `{
  "keys": [
    {
      "alg": "RS256",
      "e": "AQAB",
      "kid": "f6b0b265-a619-45c5-a252-b3c39a6bcf2f",
      "kty": "RSA",
      "n": "xJ5vz6hhY1ESZQkB0tfyxN8PGwJJygfgqUApITCsw1RpWZKjGu4u0_yBwkJJvk_-GpWDwE8vQnMBGwatPMRlHsvbrOQ3Ltfk2dFP9-M-PbTlB_da3IZO1pJLQAR0cMeFrd_zIgeaRJyrn945xYMOR2WehxXRhn69oAfwy4eODqFervJo0TETAEEJWbt1EDA-4wR8Vto4lnwdiZH6MTwxwVAE9fKiQUBuRIRG2jf0uP1Z68MpnyXCU-IK6_urqta0oLiG6QRVccc8omVEwTQo1ZYPWXUWcRPD_73ClcfBJb8ocHILGPjgh0etmB9NY8s4GDrW2L39FxYvSrTVB9Ckfw",
      "use": "sig"
    }
  ]
}
`

allow {
	input.message == "hello"
	[valid, p, s] := io.jwt.decode_verify(input.token, {"cert": jwks, "aud": "some-audience"})
}

input.json

{
    "message": "hello",
    "token": "eyJhbGciOiJSUzI1NiIsImtpZCI6ImY2YjBiMjY1LWE2MTktNDVjNS1hMjUyLWIzYzM5YTZiY2YyZiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoic29tZS11c2VyIiwiYWNjb3VudCI6InNvbWUtYWNjb3VudCIsImFybiI6InNvbWUtYXJuIiwic2NwIjpbIm1vbm9jbGUuYWJvdXQuZ2V0Il0sImlhdCI6MTY2MzA5NTgzNiwiZXhwIjoxNjk0NjMxODM2LCJpc3MiOiJzb21lLWlzc3VlciIsImF1ZCI6InNvbWUtYXVkaWVuY2UiLCJzdmMiOiJzb21lc2VydmljZSJ9.NAmreK4b9sRHDerCt6s_WDonHEjRnz5n6U96Hj_ZBLIPAyL740C2m3NHlO0FeRGOMzp5IxD3Xq6VTrP7RzYt3WR_vu3RoAyqLIA7vbqwquQLWIbpu0KB3EObY3Ue0d2Yv5KXiCi0BPbof19D4Xy_5NBXUJ_fAgexITRIihgaRjrxYicaDbF1Xr_WUt3rDD_VV8oUrLdkJ9tWQqcIIFfOFkUebe2TtSyrB34c-onvm4IuyAKjCR3UdGC9zXXDF2mDu5b0xfmOpZPFNZImX-viTP-vSVRlBqHYuICV35kh3I2LIHTazn-J2NmwB93e_ttRDqMFAdwzgoQCNsj5DnQM1A=="
}

profiling the above policy using
opa eval 'data.authz.allow' --data policy.rego --input input.json --profile --format=pretty --explain=full results in the following output

query:1            | Eval data.authz.allow = _
query:1            | Index data.authz.allow (matched 1 rule, early exit)
policy.rego:19     | Enter data.authz.allow
policy.rego:20     | | Eval input.message = "hello"
policy.rego:21     | | Eval __local4__ = input.token
policy.rego:21     | | Eval __local5__ = data.authz.jwks
policy.rego:21     | | Index data.authz.jwks (matched 1 rule, early exit)
policy.rego:5      | | Enter data.authz.jwks
policy.rego:5      | | | Eval true
policy.rego:5      | | | Exit data.authz.jwks early
policy.rego:21     | | Eval io.jwt.decode_verify(__local4__, {"aud": "some-audience", "cert": __local5__}, __local3__)
policy.rego:21     | | Eval [valid, p, s] = __local3__
policy.rego:19     | | Exit data.authz.allow early
query:1            | Exit data.authz.allow = _
query:1            Redo data.authz.allow = _
query:1            | Redo data.authz.allow = _
policy.rego:19     | Redo data.authz.allow
policy.rego:21     | | Redo [valid, p, s] = __local3__
policy.rego:21     | | Redo io.jwt.decode_verify(__local4__, {"aud": "some-audience", "cert": __local5__}, __local3__)
policy.rego:21     | | Redo __local5__ = data.authz.jwks
policy.rego:5      | | | Redo true
policy.rego:21     | | Redo __local4__ = input.token
policy.rego:20     | | Redo input.message = "hello"
true
+------------------------------+--------+
|            METRIC            | VALUE  |
+------------------------------+--------+
| timer_rego_load_files_ns     | 426378 |
| timer_rego_module_compile_ns | 812491 |
| timer_rego_module_parse_ns   | 268172 |
| timer_rego_query_compile_ns  | 58428  |
| timer_rego_query_eval_ns     | 465356 |
| timer_rego_query_parse_ns    | 63158  |
+------------------------------+--------+
+-----------+----------+----------+------------------+
|   TIME    | NUM EVAL | NUM REDO |     LOCATION     |
+-----------+----------+----------+------------------+
| 370.894µs | 4        | 4        | policy.rego:21   |
| 57.816µs  | 1        | 1        | data.authz.allow |
| 19.057µs  | 1        | 1        | policy.rego:20   |
| 9.948µs   | 1        | 1        | policy.rego:5    |
+-----------+----------+----------+------------------+

I'm unclear about policy.rego:21 being evaluated 4 times with 4 redos. I'm also curious if this can be optimized to save more time in policy evaluation. Would appreciate if anyone has thoughts/comments.

[srenatus]

So here's what is notoriously misleading about the NUM EVAL and NUM REDO: it accumulates the numbers of the expressions that the source line expanded into.

In your example, allow is rewritten in the compiler to

allow = true {
    input.message = "hello"
    __local4__ = input.token
    __local5__ = data.authz.jwks
    io.jwt.decode_verify(__local4__, {"aud": "some-audience", "cert": __local5__}, __local3__)
    [__local0__, __local1__, __local2__] = __local3__
}

and the last 4 lines of the body come from the one line you had there in policy.rego:21. They're each evaluated (and redone) once, but aggregated, you'll get a 4 on both counts.

Does that help? I'm afraid there's no simple way to inspect the compiler's intermediate format there, short of adding a log.Printf to ast/compiler.go.