Huge performance gap of `JwtPublicKeyVerify` depending on the key within a keyset used for signature generation

Help us help you

We'd like to know more about
your Tink deployment.

Describe the bug:

JwtPublicKeyVerify primitive shows a huge performance gap in verifying a JWT depending on the key, within a keyset, used
to generate the JWT signature.

• JHM Benchmark Result (BM Code)

  # JMH version: 1.35
  # VM version: JDK 11.0.17, OpenJDK 64-Bit Server VM, 11.0.17+8
  # VM invoker: /Users/ks-yim/.sdkman/candidates/java/11.0.17-tem/bin/java
  # VM options: -Dfile.encoding=UTF-8 -Djava.io.tmpdir=/Users/ks-yim/IdeaProjects/tink-bm/build/tmp/jmh -Duser.country=KR -Duser.language=en -Duser.variant
  # Blackhole mode: full + dont-inline hint (auto-detected, use -Djmh.blackhole.autoDetect=false to disable)
  # Warmup: 1 iterations, 10 s each
  # Measurement: 1 iterations, 10 s each
  # Timeout: 10 min per iteration
  # Threads: 8 threads, will synchronize iterations
  # Benchmark mode: Throughput, ops/time
  
  Benchmark                                                               Mode  Cnt      Score   Error  Units
  JwtPublicKeyVerifyBenchmark.verifyJwt_signedWith1stKey                 thrpt       17936.100          ops/s
  JwtPublicKeyVerifyBenchmark.verifyJwt_signedWith2ndKey                 thrpt       21440.300          ops/s
  JwtPublicKeyVerifyBenchmark.verifyJwt_signedWith3rdKey                 thrpt       21041.441          ops/s
  JwtPublicKeyVerifyBenchmark.verifyJwt_signedWith4thKey                 thrpt       43127.601          ops/s
  JwtPublicKeyVerifyBenchmark.verifyJwt_signedWith5thKey                 thrpt       39966.294          ops/s
  

What was the expected behavior?

No significant signature verification performance gap

How can we reproduce the bug?

Clone the BM code repository and hit the following command:

./gradlew --no-daemon clean jmh -Pjmh.includes=^com.example.tink.bm.JwtPublicKeyVerifyBenchmark

Do you have any debugging information?

WrappedJwtPublicKeyVerify#verifyAndDecode doesn't use TINK's conventional way of determining the right key within a keyset by parsing the keyId prefix which is done at O(1) time.
Instead, it iterates over PrimitiveSet<JwtPublicKeyVerifyInternal> until it finds a matching key which generates the same signature as the one in the JWT, which makes the worst-case time complexity O(n).

But, it's not enough to explain the huge perf gap shown in the benchmark.

The real problem comes from the fact that verifyAndDecode generates a token signature with all tried keys within a keyset until it finds a match and uses GeneralSecurityException to control the code execution flow, which is very costly due to stacktrace generation.

A CPU profiling result for the worst case BM reveals that CPU cycles consumed for signature verification & token decoding process can potentially be optimized away. (For interactive version, hit curl -O https://raw.githubusercontent.com/ks-yim/tink-bm/master/src/jmh/java/com/example/tink/bm/jwt-public-key-verify-cpu-profiling.svg and open the downloaded file on your local)

• Signature Verification Part
  
• Token Decoding Part
  

What version of Tink are you using?

1.7.0

Can you tell us more about your development environment?

JDK 11, MacBook Pro(16-inch, 2021) Apple M1 Pro 32 GB

Is there anything else you'd like to add?

I'd like to suggest to move JwtFormat#splitSignedCompact from JwtPublicKeyVerificationInternal to JwtPublicKeyVerification so that we can have access to the parsed keyID upfront and find the right keyset at O(1) and do the token decoding & signature verification only once no matter how many keys exist in the keyset.

This will introduce a change to JwtPublicKeyVerifyInternal#verifyAndDecodeWithKid's method signature so I'd like to get some advices from the maintainers about whether it is the right direction.

Reproduce the behavior with the linked JwtPublicKeyVerifyBenchmark and inspect WrappedJwtPublicKeyVerify#verifyAndDecode, JwtPublicKeyVerificationInternal, JwtPublicKeyVerification, and JwtFormat#splitSignedCompact. Compare key selection and decoding costs across key positions; done means verification uses the available key ID without repeated signature attempts and the benchmark no longer shows a substantial position-dependent gap.