Implement AWS ALB User Claims parser (#1328)

* [jws][jwt] Implement WithBase64Encoder

* remove unused nolint

* tweak for bazel

* move test to openid

this is easier to handle than doing it in jwt, because of how the
dependencies work

* revert jws.Signature API change, workaround it

* Tweak comments

* Update Changes

* Update jws/options_gen.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update jws/options.yaml

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: lestrrat

Changes

@@ -4,6 +4,15 @@ Changes
 v3 has many incompatibilities with v2. To see the full list of differences between
 v2 and v3, please read the Changes-v3.md file (https://github.com/lestrrat-go/jwx/blob/develop/v3/Changes-v3.md)
 
+v3.0.0-beta2 UNRELEASED
+  * [jwk] Fix a bug where `jwk.Set`'s `Keys()` method did not return the proper
+    non-standard fields. (#1322)
+  * [jws][jwt] Implement `WithBase64Encoder()` options to pass base64 encoders
+    to use during signing/verifying signatures. This useful when the token
+    provider generates JWTs that don't follow the specification and uses base64
+    encoding other than raw url encoding (no padding), such as, apparently,
+    AWS ALB. (#1324, #1328)
+
 v3.0.0-beta1 15 Mar 2025
   * [jwt] Token validation no longer truncates time based fields by default.
     To restore old behavior, you can either change the global settings by

internal/base64/base64.go

@@ -35,6 +35,10 @@ func getEncoder() Encoder {
 	return encoder
 }
 
+func DefaultEncoder() Encoder {
+	return getEncoder()
+}
+
 func SetDecoder(dec Decoder) {
 	muDecoder.Lock()
 	defer muDecoder.Unlock()

jws/interface.go

@@ -1,9 +1,20 @@
 package jws
 
 import (
+	"github.com/lestrrat-go/jwx/v3/internal/base64"
 	"github.com/lestrrat-go/jwx/v3/jwa"
 )
 
+// Base64Encoder is an interface that can be used when encoding JWS message
+// components to base64. This is useful when you want to use a non-standard
+// base64 encoder while generating or verifying signatures. By default JWS
+// uses raw url base64 encoding (without padding), but there are apparently
+// some cases where you may want to use a base64 encoders that uses padding.
+//
+// For example, apparently AWS ALB User Claims is provided in JWT format,
+// but it uses a base64 encoding with padding.
+type Base64Encoder = base64.Encoder
+
 type DecodeCtx interface {
 	CollectRaw() bool
 }
@@ -55,6 +66,7 @@ type Message struct {
 }
 
 type Signature struct {
+	encoder   Base64Encoder
 	dc        DecodeCtx
 	headers   Headers // Unprotected Headers
 	protected Headers // Protected Headers

jws/jws.go

@@ -171,6 +171,7 @@ func Sign(payload []byte, options ...SignOption) ([]byte, error) {
 	var detached bool
 	var noneSignature *payloadSigner
 	var validateKey bool
+	var encoder Base64Encoder = base64.DefaultEncoder()
 	for _, option := range options {
 		//nolint:forcetypeassert
 		switch option.Ident() {
@@ -209,6 +210,8 @@ func Sign(payload []byte, options ...SignOption) ([]byte, error) {
 			payload = option.Value().([]byte)
 		case identValidateKey{}:
 			validateKey = option.Value().(bool)
+		case identBase64Encoder{}:
+			encoder = option.Value().(Base64Encoder)
 		}
 	}
 
@@ -260,6 +263,7 @@ func Sign(payload []byte, options ...SignOption) ([]byte, error) {
 		sig := &Signature{
 			headers:   signer.PublicHeader(),
 			protected: protected,
+			encoder:   encoder,
 			// cheat. FIXXXXXXMEEEEEE
 			detached: detached,
 		}
@@ -289,6 +293,11 @@ func Sign(payload []byte, options ...SignOption) ([]byte, error) {
 		if detached {
 			compactOpts = append(compactOpts, WithDetached(detached))
 		}
+		for _, option := range options {
+			if copt, ok := option.(CompactOption); ok {
+				compactOpts = append(compactOpts, copt)
+			}
+		}
 		return Compact(&result, compactOpts...)
 	default:
 		return nil, signerr(`invalid serialization format`)
@@ -328,6 +337,7 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 	var keyProviders []KeyProvider
 	var keyUsed interface{}
 	var validateKey bool
+	var encoder Base64Encoder = base64.DefaultEncoder()
 
 	ctx := context.Background()
 
@@ -359,6 +369,8 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 			validateKey = option.Value().(bool)
 		case identSerialization{}:
 			parseOptions = append(parseOptions, option.(ParseOption))
+		case identBase64Encoder{}:
+			encoder = option.Value().(Base64Encoder)
 		default:
 			return nil, verifyerr(`invalid jws.VerifyOption %q passed`, `With`+strings.TrimPrefix(fmt.Sprintf(`%T`, option.Ident()), `jws.ident`))
 		}
@@ -385,7 +397,7 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 	// Pre-compute the base64 encoded version of payload
 	var payload string
 	if msg.b64 {
-		payload = base64.EncodeToString(msg.payload)
+		payload = encoder.EncodeToString(msg.payload)
 	} else {
 		payload = string(msg.payload)
 	}
@@ -400,7 +412,7 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 		var encodedProtectedHeader string
 		if rbp, ok := sig.protected.(interface{ rawBuffer() []byte }); ok {
 			if raw := rbp.rawBuffer(); raw != nil {
-				encodedProtectedHeader = base64.EncodeToString(raw)
+				encodedProtectedHeader = encoder.EncodeToString(raw)
 			}
 		}
 
@@ -410,7 +422,7 @@ func Verify(buf []byte, options ...VerifyOption) ([]byte, error) {
 				return nil, verifyerr(`failed to marshal "protected" for signature #%d: %w`, i+1, err)
 			}
 
-			encodedProtectedHeader = base64.EncodeToString(protected)
+			encodedProtectedHeader = encoder.EncodeToString(protected)
 		}
 
 		verifyBuf.WriteString(encodedProtectedHeader)

jws/message.go

@@ -127,13 +127,17 @@ func (s *Signature) Sign(payload []byte, signer Signer, key interface{}) ([]byte
 	buf := pool.GetBytesBuffer()
 	defer pool.ReleaseBytesBuffer(buf)
 
-	buf.WriteString(base64.EncodeToString(hdrbuf))
+	encoder := s.encoder
+	if encoder == nil {
+		encoder = base64.DefaultEncoder()
+	}
+	buf.WriteString(encoder.EncodeToString(hdrbuf))
 	buf.WriteByte('.')
 
 	var plen int
 	b64 := getB64Value(hdrs)
 	if b64 {
-		encoded := base64.EncodeToString(payload)
+		encoded := encoder.EncodeToString(payload)
 		plen = len(encoded)
 		buf.WriteString(encoded)
 	} else {
@@ -158,7 +162,7 @@ func (s *Signature) Sign(payload []byte, signer Signer, key interface{}) ([]byte
 	}
 
 	buf.WriteByte('.')
-	buf.WriteString(base64.EncodeToString(signature))
+	buf.WriteString(encoder.EncodeToString(signature))
 	ret := make([]byte, buf.Len())
 	copy(ret, buf.Bytes())
 
@@ -466,11 +470,14 @@ func Compact(msg *Message, options ...CompactOption) ([]byte, error) {
 	}
 
 	var detached bool
+	var encoder Base64Encoder = base64.DefaultEncoder()
 	for _, option := range options {
 		//nolint:forcetypeassert
 		switch option.Ident() {
 		case identDetached{}:
 			detached = option.Value().(bool)
+		case identBase64Encoder{}:
+			encoder = option.Value().(Base64Encoder)
 		}
 	}
 
@@ -486,12 +493,12 @@ func Compact(msg *Message, options ...CompactOption) ([]byte, error) {
 	buf := pool.GetBytesBuffer()
 	defer pool.ReleaseBytesBuffer(buf)
 
-	buf.WriteString(base64.EncodeToString(hdrbuf))
+	buf.WriteString(encoder.EncodeToString(hdrbuf))
 	buf.WriteByte('.')
 
 	if !detached {
 		if getB64Value(hdrs) {
-			encoded := base64.EncodeToString(msg.payload)
+			encoded := encoder.EncodeToString(msg.payload)
 			buf.WriteString(encoded)
 		} else {
 			if bytes.Contains(msg.payload, []byte{'.'}) {
@@ -502,7 +509,7 @@ func Compact(msg *Message, options ...CompactOption) ([]byte, error) {
 	}
 
 	buf.WriteByte('.')
-	buf.WriteString(base64.EncodeToString(s.signature))
+	buf.WriteString(encoder.EncodeToString(s.signature))
 	ret := make([]byte, buf.Len())
 	copy(ret, buf.Bytes())
 	return ret, nil

jws/options.yaml

@@ -20,6 +20,15 @@ interfaces:
       - parseOption
     comment: |
       SignVerifyOption describes options that can be passed to either `jws.Verify` or `jws.Sign`
+  - name: SignVerifyCompactOption
+    methods:
+      - signOption
+      - verifyOption
+      - compactOption
+      - parseOption
+    comment: |
+      SignVerifyCompactOption describes options that can be passed to either `jws.Verify`,
+      `jws.Sign`, or `jws.Compact`
   - name: WithJSONSuboption
     concrete_type: withJSONSuboption
     comment: |
@@ -78,6 +87,13 @@ options:
        must be set to `nil`.
        
        If you have to verify using this option, you should know exactly how and why this works.
+  - ident: Base64Encoder
+    interface: SignVerifyCompactOption
+    argument_type: Base64Encoder
+    comment: |
+      WithBase64Encoder specifies the base64 encoder to be used while signing or
+      verifying the JWS message. By default, the raw URL base64 encoding (no padding)
+      is used.
   - ident: Message
     interface: VerifyOption
     argument_type: '*Message'

jws/options_gen.go

@@ -201,7 +201,7 @@ func parseBytes(data []byte, options ...ParseOption) (Token, error) {
 
 		//nolint:forcetypeassert
 		switch o.Ident() {
-		case identKey{}, identKeySet{}, identVerifyAuto{}, identKeyProvider{}:
+		case identKey{}, identKeySet{}, identVerifyAuto{}, identKeyProvider{}, identBase64Encoder{}:
 			verifyOpts = append(verifyOpts, o)
 		case identToken{}:
 			token, ok := o.Value().(Token)