From a016461d897d3df6d0c8f0e033fd5cd43f436f1c Mon Sep 17 00:00:00 2001
From: lestrrat <49281+lestrrat@users.noreply.github.com>
Date: Fri, 30 Jul 2021 13:06:14 +0900
Subject: [PATCH] Rfc7797 prohibit in jwt (#424)

* Do away with draining, discard error if write fails

* Prohibit b64 = false for JWTs
---
 jwt/jwt_test.go  | 19 +++++++++++++++++++
 jwt/serialize.go | 10 ++++++++++
 2 files changed, 29 insertions(+)

diff --git a/jwt/jwt_test.go b/jwt/jwt_test.go
index 8dc55e88..0827e3ae 100644
--- a/jwt/jwt_test.go
+++ b/jwt/jwt_test.go
@@ -1060,3 +1060,22 @@ func TestNested(t *testing.T) {
 	}
 	_ = parsed
 }
+
+func TestRFC7797(t *testing.T) {
+	key, err := jwxtest.GenerateRsaKey()
+	if !assert.NoError(t, err, `jwxtest.GenerateRsaKey should succeed`) {
+		return
+	}
+
+	hdrs := jws.NewHeaders()
+	hdrs.Set("b64", false)
+	hdrs.Set("crit", "b64")
+
+	token := jwt.New()
+	token.Set(jwt.AudienceKey, `foo`)
+
+	_, err = jwt.Sign(token, jwa.RS256, key, jwt.WithJwsHeaders(hdrs))
+	if !assert.Error(t, err, `jwt.Sign should fail`) {
+		return
+	}
+}
diff --git a/jwt/serialize.go b/jwt/serialize.go
index 00427d32..a3665ce6 100644
--- a/jwt/serialize.go
+++ b/jwt/serialize.go
@@ -147,6 +147,16 @@ func (s *jwsSerializer) Serialize(ctx SerializeCtx, v interface{}) (interface{},
 	if err := setTypeOrCty(ctx, hdrs); err != nil {
 		return nil, err // this is already wrapped
 	}
+
+	// JWTs MUST NOT use b64 = false
+	// https://datatracker.ietf.org/doc/html/rfc7797#section-7
+	if v, ok := hdrs.Get("b64"); ok {
+		if bval, bok := v.(bool); bok {
+			if !bval { // b64 = false
+				return nil, errors.New(`b64 cannot be false for JWTs`)
+			}
+		}
+	}
 	return jws.Sign(payload, s.alg, s.key, jws.WithHeaders(hdrs))
 }