Replace interned `internal/jwx` with proper dependency to `lestrrat-go/jwx`

[johanfylling]

The io.jwt.* built-ins' cryptographic elements are implemented through internal/jwx. This is a modified interning of lestrrat-go/jwx. Since interning, this project has seen multiple updates (among others, support for EdDSA) that have not been downstreamed to OPA.

It should be investigated whether it's possible to drop the interned internal/jwx in favor of a proper dependency to lestrrat-go/jwx. This would make it easier to benefit from fixes and improvements made to jwx going forward.

If this is not possible, internal/jwx should be updated and then adapted to work with OPA. Preferably, an attempt should be made to upstream any necessary changes to lestrrat-go/jwx, so that it can eventually be made into a proper dependency.

[anderseknert]

Related #5835

[lestrrat]

Is there a list of actual semantic changes that you applied to internal/jwx anywhere?
Actually, from a cursory look, I think it's pre-1.0 code :D am I correct?

[anderseknert]

Most likely, yes — it's absolutely ancient 😅 I wasn't around when this was vendored, but I think in general (i.e. not just for this lib) the concern has been that even subtle changes in library behavior could potentially change the outcome of policy evaluation between OPA upgrades. And since there's no way we realistically could cover everything, the strategy was to vendor/freeze some dependencies and try and do somewhat controlled upgrades — which in many cases (and with the benefit of hindsight — obviously) never happened.

git log --pretty=format:"%h%x09%an%x09%ad%x09%s" -- internal/jwx
e43ef0a97       Anders Eknert   Mon May 12 13:57:48 2025 +0200  Use `any` in place of `interface{}` (#7566)
afb30d3f9       Anders Eknert   Mon Feb 24 16:28:41 2025 +0100  Add gocritic linter, fix a bunch of stuff (#7377)
55e87e79a       Anders Eknert   Fri Jan 31 20:24:05 2025 +0100  Add perfsprint linter (#7334)
2131da3d9       Iceber Gu       Wed Jan 11 12:23:17 2023 +0800  remove github.com/pkg/errors
50d4e31d6       Anders Eknert   Thu Oct 27 17:44:56 2022 +0200  chore: Use t.Setenv in tests (#5321)
95708108f       Anders Eknert   Thu Oct 27 13:35:39 2022 +0200  linters: add unconvert (#5318)
cb6a4c0b7       Anders Eknert   Thu Jun 2 12:07:57 2022 +0200   Ignore keys of unknown alg when verifying JWTs with JWKS (#4725)
4dd7fb1c0       Jason Hall      Wed May 18 05:29:35 2022 -0400  Remove use of github.com/pkg/errors (#4696)
22d505fd9       Cris He Fri Sep 17 01:03:16 2021 -0400  built-ins: decode a RSA private key into JWK format (#3783)
e732b0b48       Stephan Renatus Thu Aug 19 07:43:16 2021 +0200  topdown/buitins: io.jwt.encode_sign uses BuiltinContext random source (#3738)
3be1d08b8       Will Beason     Wed May 19 00:52:02 2021 -0500  Change check-lint to use golangci-lint (#3465)
134ed9b50       André Håland    Fri Mar 19 11:43:15 2021 +0100  Allow PKCS8 encoded EC private keys (#3288)
fe97f335f       Anders Eknert   Wed Feb 3 20:16:14 2021 +0100   Allow PKCS8 encoded private keys (#3117)
ccac1370c       Stephan Renatus Fri Dec 18 15:39:29 2020 +0100  Fix const types in internal/jwx/jwk/key_ops.go and repl/repl.go (#3021)
338583c18       Ashutosh Narkar Tue Jun 9 16:15:27 2020 -0700   Add support for OPA bundle signatures

Many recent changes looks like global search and replace stuff. These look like the ones we'll want to zoom in on:

cb6a4c0b7       Anders Eknert   Thu Jun 2 12:07:57 2022 +0200   Ignore keys of unknown alg when verifying JWTs with JWKS (#4725)
22d505fd9       Cris He Fri Sep 17 01:03:16 2021 -0400  built-ins: decode a RSA private key into JWK format (#3783)
e732b0b48       Stephan Renatus Thu Aug 19 07:43:16 2021 +0200  topdown/buitins: io.jwt.encode_sign uses BuiltinContext random source (#3738)
134ed9b50       André Håland    Fri Mar 19 11:43:15 2021 +0100  Allow PKCS8 encoded EC private keys (#3288)
fe97f335f       Anders Eknert   Wed Feb 3 20:16:14 2021 +0100   Allow PKCS8 encoded private keys (#3117)
ccac1370c       Stephan Renatus Fri Dec 18 15:39:29 2020 +0100  Fix const types in internal/jwx/jwk/key_ops.go and repl/repl.go (#3021)

Thanks for dropping by, btw :)

[johanfylling]

Some added context: one of the reasons for why this dependency was interned was to minimize the number of active dependencies on the eval hot path; to make integrations easier, with minimal dependency versioning conflicts. If this goal remains unchanged, the preferred course forward is to downstream the updated code, but to keep it interned.
However, doing so will eventually end us up in the same situation in the future, with outdated code that might need manual updating.

[lestrrat]

Huh, I let Cursor do the initial cleanup, and started digging, but hmmm.
It looks like you are using the jwx packages more so for building blocks of your own jwt/jws implementation, which is fine, but jwx has been evolving more towards a high-level API, batteries-included type of everything JOSE, and I'm kind of confused how much modern jwx-ism I can put into opa without causing major problems.

For example, if I replace some of the JWS code with the modern jwx/jws stuff, it will affect many of the tests in the slightest parts -- while opa takes string headers like {"typ":"text/plain"\n "alg":"HS256"} (notice the newline and space, also the order of keys) and convert them literally, when you go through jwx this is normalized to {"alg":"HS256","typ":"text/plain"} and then the signature is produced, so outcome is different. I know jwx's signing is valid, but since jwx doesn't expect to take raw header payloads to sign, there's just no way to "fix" it to match the old behavior right now. And I can't tell if this is important to opa or not.

So I stopped for now.

---

On a tangent, would you rather if jwx provided low level building blocks rather than a high-level API for these things?

e.g.

jwt.Sign(header []byte, payload []byte, alg jwa.SignatureAlgorithm, key any) ([]byte, error)

Instead of

token, _ := jwt.NewBuilder().Issuer("foo")./* other types of initialization... */Build()
jwt.Sign(token, jwt.WithKey(alg, key)

[anderseknert]

@lestrrat thanks for taking the time! If we have tests that break only because a signed JWT is different from one version to another, those aren't great to begin with, and we should rewrite them, IMHO. While we're pretty strict about not breaking things, I think the contract for JWT signing will have to be "is it a valid JWT with the payload the user provided?" and not that its string/byte representation is identical to some past version.

With that said, the contract/description for io.jwt.encode_sign_raw is pretty much that we leave control over all details to the user, and a low-level function as you suggested would likely be needed.

Where we use jwx internally, we can probably use the high-level API — the built-in functions related to signing and verification are really what matters the most.

[lestrrat]

Okay, if upgrading jwx isn't a top priority, you may want to wait for me conjure up the low level APIs for jwx. I'm starting to think of something like this

package jwsbb // bb for building blocks -- needed something non-generic so that it doesn't have a chance to clash with some emerging new JOSE standard :)

func SignECDSA(payload, encodedHeader []byte, h crypto.Hash, key *ecdsa.PrivateKey) ([]byte, error)
func SignEdDSA(payload, encodedHeader []byte, key ed25519.PrivateKey) ([]byte, error)
func SignHMAC(payload, encodedHeader []byte, hfunc func() hash.Hash, key []byte) ([]byte, error)
func SignRSA(payload, encodedHeader []byte, h crypto.Hash, pss bool, key *rsa.PrivateKey) ([]byte, error)

No promises yet, though!

(It turns out that I kind of wanted it myself when trying to optimize jwx this past week)

[anderseknert]

Awesome! No, there’s no urgency here whatsoever. We’ve evidently done OK with this ancient version for the past five years, so what’s another week/month or two? :)

[lestrrat]

@anderseknert

I have semi-working code in my jwx3 branch.

• I've been focusing on v1/bundle directory. I haven't really checked other parts.
• I haven't run go mod vendor to add/delete files, as I thought it would clutter the diff.
• Tests that check error contents are failing

The test failures are due to the tests checking for matches in the error messages. I obviously changed a lot, so the error messages got all screwed. I didn't know if I should fix the tests or try to match the original expected results, so I'm letting them be for now. At least they are all erroring when the should be erroring :)

One general assessment is that you seem to be implementing a lot of the JWT/JWS workflow yourself. If you were up to it, you could just delegate most of it to jwx -- but again, that would change minor things like error messages and such, so I'm not sure that's what you want.

Anyway, please see if this general direction works for you?

[anderseknert]

Hey, and sorry for not getting back to you sooner! I’ll set some time aside to check this out next week and will come back with my thoughts. Thanks! 🙏

[anderseknert]

Changes all look good to me. Thanks a lot @lestrrat 😃

The tests that fail on matching exact output can just be updated to the new output for now, and we can consider later how to better test signing/verification. Likewise for the error messages.

With that out of the way, all that remains is OPA's built-in functions for handling JWTs. I checked out your repo and figured I would give it a shot, and the jwsbb package is great! Now doing this for signing raw payloads:

payload := jwsbb.SignBuffer(nil, []byte(inputHeaders), []byte(jwsPayload), base64.RawURLEncoding, true)

signature, err := jwsbb.Sign(key, algStr, payload)
if err != nil {
	return err
}

jwsCompact := string(payload) + "." + base64.RawURLEncoding.EncodeToString(signature)

And all but one of the failing tests are related to error messages, which should be easy to fix. The one not related to that however seems a little more tricky. That's this test, which asserts that a provided seed is used for signing with ECDSA. That's needed in order to be able to replay a policy decision later, given the same input data. The currently used jws.SignLiteral would let us pass an io.Reader for this purpose, but that's AFAICS not an option using jwsbb. What do you think the best way for us to accomplish that now would be?

[anderseknert]

I got that working borrowing some code from your implementation, modified to allow passing a reader:

	if ecdsaKey, ok := key.(*ecdsa.PrivateKey); ok {
		signature, err = signECDSA(ecdsaKey, payload, tokenAlgorithms[algStr].hash, bctx.Seed)
		if err != nil {
			return fmt.Errorf("failed to sign payload using ECDSA: %w", err)
		}
	} else {
		signature, err = jwsbb.Sign(key, algStr, payload)
		if err != nil {
			return err
		}
	}

This works for me, but let me know if there's a better way.

[lestrrat]

@anderseknert Do you have the branch you're working on available? I'm not too familiar with the internal workings of opa, so I'm having to trace where the mysterious seed is being passed to/from