Repository: lestrrat-go/jwx
Pull request #1057 (2024-01-10 23:38:36)

Hardening code

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) Report
Attention: `29 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`0a15b4d`)](https://app.codecov.io/gh/lestrrat-go/jwx/commit/0a15b4d9a180b06c6924c4a03b39ae70d05514f4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) 72.00% compared to head [(`2a35494`)](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) 72.28%.
> Report is 26 commits behind head on develop/v2.

| [Files](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) | Patch % | Lines |
|---|---|---|
| [internal/keyconv/keyconv.go](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go#diff-aW50ZXJuYWwva2V5Y29udi9rZXljb252Lmdv) | 42.42% | [16 Missing and 3 partials :warning: ](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) |
| [jwt/openid/birthdate.go](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go#diff-and0L29wZW5pZC9iaXJ0aGRhdGUuZ28=) | 76.47% | [4 Missing :warning: ](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) |
| [jwe/jwe.go](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go#diff-andlL2p3ZS5nbw==) | 91.66% | [2 Missing and 1 partial :warning: ](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) |
| [jwe/internal/aescbc/aescbc.go](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go#diff-andlL2ludGVybmFsL2Flc2NiYy9hZXNjYmMuZ28=) | 85.71% | [1 Missing and 1 partial :warning: ](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) |
| [jwe/options\_gen.go](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go#diff-andlL29wdGlvbnNfZ2VuLmdv) | 92.30% | [1 Missing :warning: ](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go) |

<details><summary>Additional details and impacted files</summary>


```diff
@@              Coverage Diff               @@
##           develop/v2    #1057      +/-   ##
==============================================
+ Coverage       72.00%   72.28%   +0.28%     
==============================================
  Files              93       93              
  Lines           13795    13896     +101     
==============================================
+ Hits             9933    10045     +112     
+ Misses           3044     3036       -8     
+ Partials          818      815       -3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lestrrat-go/jwx/pull/1057?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lestrrat-go).


Patch:
diff --git a/Changes b/Changes
index 385e45a7..7c6ae61b 100644
--- a/Changes
+++ b/Changes
@@ -4,6 +4,17 @@ Changes
 v2 has many incompatibilities with v1. To see the full list of differences between
 v1 and v2, please read the Changes-v2.md file (https://github.com/lestrrat-go/jwx/blob/develop/v2/Changes-v2.md)
 
+v2.0.20 UNRELEASED
+[New Features]
+  * [jwe] Added jwe.Settings(WithMaxBufferSize(int64)) to set the maximum size of
+    internal buffers. The default value is 256MB. Most users do not need to change
+    this value.
+
+[Miscellaneous]
+  * Internal key conversions should now allow private keys to be used in place of
+    public keys. This would allow you to pass private keys where public keys are
+    expected.
+
 v2.0.19 09 Jan 2024
 [New Features]
   * [jws] Added jws.IsVerificationError to check if the error returned by `jws.Verify`
diff --git a/internal/keyconv/keyconv.go b/internal/keyconv/keyconv.go
index 807da1de..f0ecbae1 100644
--- a/internal/keyconv/keyconv.go
+++ b/internal/keyconv/keyconv.go
@@ -18,7 +18,7 @@ func RSAPrivateKey(dst, src interface{}) error {
 	if jwkKey, ok := src.(jwk.Key); ok {
 		var raw rsa.PrivateKey
 		if err := jwkKey.Raw(&raw); err != nil {
-			return fmt.Errorf(`failed to produce rsa.PrivateKey from %T: %w`, src, err)
+			return fmt.Errorf(`keyconv: failed to produce rsa.PrivateKey from %T: %w`, src, err)
 		}
 		src = &raw
 	}
@@ -30,7 +30,7 @@ func RSAPrivateKey(dst, src interface{}) error {
 	case *rsa.PrivateKey:
 		ptr = src
 	default:
-		return fmt.Errorf(`expected rsa.PrivateKey or *rsa.PrivateKey, got %T`, src)
+		return fmt.Errorf(`keyconv: expected rsa.PrivateKey or *rsa.PrivateKey, got %T`, src)
 	}
 
 	return blackmagic.AssignIfCompatible(dst, ptr)
@@ -41,21 +41,25 @@ func RSAPrivateKey(dst, src interface{}) error {
 // `src` may be rsa.PublicKey, *rsa.PublicKey, or a jwk.Key
 func RSAPublicKey(dst, src interface{}) error {
 	if jwkKey, ok := src.(jwk.Key); ok {
-		var raw rsa.PublicKey
-		if err := jwkKey.Raw(&raw); err != nil {
-			return fmt.Errorf(`failed to produce rsa.PublicKey from %T: %w`, src, err)
+		pk, err := jwk.PublicRawKeyOf(jwkKey)
+		if err != nil {
+			return fmt.Errorf(`keyconv: failed to produce public key from %T: %w`, src, err)
 		}
-		src = &raw
+		src = pk
 	}
 
 	var ptr *rsa.PublicKey
 	switch src := src.(type) {
+	case rsa.PrivateKey:
+		ptr = &src.PublicKey
+	case *rsa.PrivateKey:
+		ptr = &src.PublicKey
 	case rsa.PublicKey:
 		ptr = &src
 	case *rsa.PublicKey:
 		ptr = src
 	default:
-		return fmt.Errorf(`expected rsa.PublicKey or *rsa.PublicKey, got %T`, src)
+		return fmt.Errorf(`keyconv: expected rsa.PublicKey/rsa.PrivateKey or *rsa.PublicKey/*rsa.PrivateKey, got %T`, src)
 	}
 
 	return blackmagic.AssignIfCompatible(dst, ptr)
@@ -67,7 +71,7 @@ func ECDSAPrivateKey(dst, src interface{}) error {
 	if jwkKey, ok := src.(jwk.Key); ok {
 		var raw ecdsa.PrivateKey
 		if err := jwkKey.Raw(&raw); err != nil {
-			return fmt.Errorf(`failed to produce ecdsa.PrivateKey from %T: %w`, src, err)
+			return fmt.Errorf(`keyconv: failed to produce ecdsa.PrivateKey from %T: %w`, src, err)
 		}
 		src = &raw
 	}
@@ -79,7 +83,7 @@ func ECDSAPrivateKey(dst, src interface{}) error {
 	case *ecdsa.PrivateKey:
 		ptr = src
 	default:
-		return fmt.Errorf(`expected ecdsa.PrivateKey or *ecdsa.PrivateKey, got %T`, src)
+		return fmt.Errorf(`keyconv: expected ecdsa.PrivateKey or *ecdsa.PrivateKey, got %T`, src)
 	}
 	return blackmagic.AssignIfCompatible(dst, ptr)
 }
@@ -88,21 +92,25 @@ func ECDSAPrivateKey(dst, src interface{}) error {
 // non-pointer to a pointer
 func ECDSAPublicKey(dst, src interface{}) error {
 	if jwkKey, ok := src.(jwk.Key); ok {
-		var raw ecdsa.PublicKey
-		if err := jwkKey.Raw(&raw); err != nil {
-			return fmt.Errorf(`failed to produce ecdsa.PublicKey from %T: %w`, src, err)
+		pk, err := jwk.PublicRawKeyOf(jwkKey)
+		if err != nil {
+			return fmt.Errorf(`keyconv: failed to produce public key from %T: %w`, src, err)
 		}
-		src = &raw
+		src = pk
 	}
 
 	var ptr *ecdsa.PublicKey
 	switch src := src.(type) {
+	case ecdsa.PrivateKey:
+		ptr = &src.PublicKey
+	case *ecdsa.PrivateKey:
+		ptr = &src.PublicKey
 	case ecdsa.PublicKey:
 		ptr = &src
 	case *ecdsa.PublicKey:
 		ptr = src
 	default:
-		return fmt.Errorf(`expected ecdsa.PublicKey or *ecdsa.PublicKey, got %T`, src)
+		return fmt.Errorf(`keyconv: expected ecdsa.PublicKey/ecdsa.PrivateKey or *ecdsa.PublicKey/*ecdsa.PrivateKey, got %T`, src)
 	}
 	return blackmagic.AssignIfCompatible(dst, ptr)
 }
@@ -111,13 +119,13 @@ func ByteSliceKey(dst, src interface{}) error {
 	if jwkKey, ok := src.(jwk.Key); ok {
 		var raw []byte
 		if err := jwkKey.Raw(&raw); err != nil {
-			return fmt.Errorf(`failed to produce []byte from %T: %w`, src, err)
+			return fmt.Errorf(`keyconv: failed to produce []byte from %T: %w`, src, err)
 		}
 		src = raw
 	}
 
 	if _, ok := src.([]byte); !ok {
-		return fmt.Errorf(`expected []byte, got %T`, src)
+		return fmt.Errorf(`keyconv: expected []byte, got %T`, src)
 	}
 	return blackmagic.AssignIfCompatible(dst, src)
 }
@@ -145,11 +153,18 @@ func Ed25519PrivateKey(dst, src interface{}) error {
 
 func Ed25519PublicKey(dst, src interface{}) error {
 	if jwkKey, ok := src.(jwk.Key); ok {
-		var raw ed25519.PublicKey
-		if err := jwkKey.Raw(&raw); err != nil {
-			return fmt.Errorf(`failed to produce ed25519.PublicKey from %T: %w`, src, err)
+		pk, err := jwk.PublicRawKeyOf(jwkKey)
+		if err != nil {
+			return fmt.Errorf(`keyconv: failed to produce public key from %T: %w`, src, err)
 		}
-		src = &raw
+		src = pk
+	}
+
+	switch key := src.(type) {
+	case ed25519.PrivateKey:
+		src = key.Public()
+	case *ed25519.PrivateKey:
+		src = key.Public()
 	}
 
 	var ptr *ed25519.PublicKey
diff --git a/jwe/BUILD.bazel b/jwe/BUILD.bazel
index a5e3f4e0..d0a06f0d 100644
--- a/jwe/BUILD.bazel
+++ b/jwe/BUILD.bazel
@@ -25,6 +25,7 @@ go_library(
         "//internal/keyconv",
         "//internal/pool",
         "//jwa",
+        "//jwe/internal/aescbc",
         "//jwe/internal/cipher",
         "//jwe/internal/content_crypt",
         "//jwe/internal/keyenc",
diff --git a/jwe/internal/aescbc/aescbc.go b/jwe/internal/aescbc/aescbc.go
index acb5a83a..749277b9 100644
--- a/jwe/internal/aescbc/aescbc.go
+++ b/jwe/internal/aescbc/aescbc.go
@@ -10,18 +10,42 @@ import (
 	"errors"
 	"fmt"
 	"hash"
+	"sync/atomic"
 )
 
 const (
 	NonceSize = 16
 )
 
+const defaultBufSize int64 = 256 * 1024 * 1024
+
+// Grr, we would like to use atomic.Int64, but that's only available
+// from Go 1.19. Yes, we will cut support for Go 1.19 at some point,
+// but not today (probably going to up the minimum required Go version
+// some time after 1.22 is released)
+var maxBufSize int64
+
+func init() {
+	atomic.StoreInt64(&maxBufSize, defaultBufSize)
+}
+
+func SetMaxBufferSize(siz int64) {
+	if siz <= 0 {
+		siz = defaultBufSize
+	}
+	atomic.StoreInt64(&maxBufSize, siz)
+}
+
 func pad(buf []byte, n int) []byte {
 	rem := n - len(buf)%n
 	if rem == 0 {
 		return buf
 	}
 
+	mbs := atomic.LoadInt64(&maxBufSize)
+	if int64(len(buf)+rem) > mbs {
+		panic(fmt.Errorf("failed to allocate buffer"))
+	}
 	newbuf := make([]byte, len(buf)+rem)
 	copy(newbuf, buf)
 
@@ -174,6 +198,11 @@ func ensureSize(dst []byte, n int) []byte {
 // Seal fulfills the crypto.AEAD interface
 func (c Hmac) Seal(dst, nonce, plaintext, data []byte) []byte {
 	ctlen := len(plaintext)
+	bufsiz := ctlen + c.Overhead()
+	mbs := atomic.LoadInt64(&maxBufSize)
+	if int64(bufsiz) > mbs {
+		panic(fmt.Errorf("failed to allocate buffer"))
+	}
 	ciphertext := make([]byte, ctlen+c.Overhead())[:ctlen]
 	copy(ciphertext, plaintext)
 	ciphertext = pad(ciphertext, c.blockCipher.BlockSize())
diff --git a/jwe/jwe.go b/jwe/jwe.go
index 7c2905a0..ae1b8e3a 100644
--- a/jwe/jwe.go
+++ b/jwe/jwe.go
@@ -19,6 +19,7 @@ import (
 	"github.com/lestrrat-go/jwx/v2/jwk"
 
 	"github.com/lestrrat-go/jwx/v2/jwa"
+	"github.com/lestrrat-go/jwx/v2/jwe/internal/aescbc"
 	"github.com/lestrrat-go/jwx/v2/jwe/internal/content_crypt"
 	"github.com/lestrrat-go/jwx/v2/jwe/internal/keyenc"
 	"github.com/lestrrat-go/jwx/v2/jwe/internal/keygen"
@@ -36,6 +37,8 @@ func Settings(options ...GlobalOption) {
 		switch option.Ident() {
 		case identMaxPBES2Count{}:
 			maxPBES2Count = option.Value().(int)
+		case identMaxBufferSize{}:
+			aescbc.SetMaxBufferSize(option.Value().(int64))
 		}
 	}
 }
diff --git a/jwe/options.yaml b/jwe/options.yaml
index bf7e0a01..2de7fdd1 100644
--- a/jwe/options.yaml
+++ b/jwe/options.yaml
@@ -138,4 +138,15 @@ options:
     comment: |
       WithMaxPBES2Count specifies the maximum number of PBES2 iterations
       to use when decrypting a message. If not specified, the default
-      value of 10,000 is used.
\ No newline at end of file
+      value of 10,000 is used.
+
+      This option has a global effect.
+  - ident: MaxBufferSize
+    interface: GlobalOption
+    argument_type: int64
+    comment: |
+      WithMaxBufferSize specifies the maximum buffer size for internal
+      calculations, such as when AES-CBC is performed. The default value is 256MB.
+      If set to an invalid value, the default value is used.
+
+      This option has a global effect.
\ No newline at end of file
diff --git a/jwe/options_gen.go b/jwe/options_gen.go
index b3de1320..3b74a8a9 100644
--- a/jwe/options_gen.go
+++ b/jwe/options_gen.go
@@ -129,6 +129,7 @@ type identFS struct{}
 type identKey struct{}
 type identKeyProvider struct{}
 type identKeyUsed struct{}
+type identMaxBufferSize struct{}
 type identMaxPBES2Count struct{}
 type identMergeProtectedHeaders struct{}
 type identMessage struct{}
@@ -166,6 +167,10 @@ func (identKeyUsed) String() string {
 	return "WithKeyUsed"
 }
 
+func (identMaxBufferSize) String() string {
+	return "WithMaxBufferSize"
+}
+
 func (identMaxPBES2Count) String() string {
 	return "WithMaxPBES2Count"
 }
@@ -245,9 +250,20 @@ func WithKeyUsed(v interface{}) DecryptOption {
 	return &decryptOption{option.New(identKeyUsed{}, v)}
 }
 
+// WithMaxBufferSize specifies the maximum buffer size for internal
+// calculations, such as when AES-CBC is performed. The default value is 256MB.
+// If set to an invalid value, the default value is used.
+//
+// This option has a global effect.
+func WithMaxBufferSize(v int64) GlobalOption {
+	return &globalOption{option.New(identMaxBufferSize{}, v)}
+}
+
 // WithMaxPBES2Count specifies the maximum number of PBES2 iterations
 // to use when decrypting a message. If not specified, the default
 // value of 10,000 is used.
+//
+// This option has a global effect.
 func WithMaxPBES2Count(v int) GlobalOption {
 	return &globalOption{option.New(identMaxPBES2Count{}, v)}
 }
diff --git a/jwt/openid/birthdate.go b/jwt/openid/birthdate.go
index c356d7f0..a193b403 100644
--- a/jwt/openid/birthdate.go
+++ b/jwt/openid/birthdate.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"math"
 	"regexp"
 	"strconv"
 
@@ -57,9 +58,22 @@ func (b *BirthdateClaim) UnmarshalJSON(data []byte) error {
 	return nil
 }
 
-func tointptr(v int64) *int {
-	i := int(v)
-	return &i
+var intSize int
+
+func init() {
+	switch math.MaxInt {
+	case math.MaxInt16:
+		intSize = 16
+	case math.MaxInt32:
+		intSize = 32
+	case math.MaxInt64:
+		intSize = 64
+	}
+}
+
+func parseBirthdayInt(s string) int {
+	i, _ := strconv.ParseInt(s, 10, intSize)
+	return int(i)
 }
 
 var birthdateRx = regexp.MustCompile(`^(\d{4})-(\d{2})-(\d{2})$`)
@@ -100,23 +114,23 @@ func (b *BirthdateClaim) Accept(v interface{}) error {
 		// we can assume that strconv.ParseInt always succeeds.
 		// strconv.ParseInt (and strconv.ParseUint that it uses internally)
 		// only returns range errors, so we should be safe.
-		year, _ := strconv.ParseInt(v[indices[2]:indices[3]], 10, 64)
+		year := parseBirthdayInt(v[indices[2]:indices[3]])
 		if year <= 0 {
 			return fmt.Errorf(`failed to parse birthdate year`)
 		}
-		tmp.year = tointptr(year)
+		tmp.year = &year
 
-		month, _ := strconv.ParseInt(v[indices[4]:indices[5]], 10, 64)
+		month := parseBirthdayInt(v[indices[4]:indices[5]])
 		if month <= 0 {
 			return fmt.Errorf(`failed to parse birthdate month`)
 		}
-		tmp.month = tointptr(month)
+		tmp.month = &month
 
-		day, _ := strconv.ParseInt(v[indices[6]:indices[7]], 10, 64)
+		day := parseBirthdayInt(v[indices[6]:indices[7]])
 		if day <= 0 {
 			return fmt.Errorf(`failed to parse birthdate day`)
 		}
-		tmp.day = tointptr(day)
+		tmp.day = &day
 
 		*b = tmp
 		return nil


Test patch:
diff --git a/jwe/jwe_test.go b/jwe/jwe_test.go
index 9559abcd..5f9a3bcb 100644
--- a/jwe/jwe_test.go
+++ b/jwe/jwe_test.go
@@ -959,3 +959,15 @@ func TestGHSA_7f9x_gw85_8grf(t *testing.T) {
 		}
 	}
 }
+
+func TestMaxBufferSize(t *testing.T) {
+	// NOTE: This has GLOBAL EFFECT
+	jwe.Settings(jwe.WithMaxBufferSize(1))
+	defer jwe.Settings(jwe.WithMaxBufferSize(0))
+
+	key, err := jwxtest.GenerateRsaJwk()
+	require.NoError(t, err, `jwxtest.GenerateRsaJwk should succeed`)
+
+	_, err = jwe.Encrypt([]byte("Lorem Ipsum"), jwe.WithContentEncryption(jwa.A128CBC_HS256), jwe.WithKey(jwa.RSA_OAEP, key))
+	require.Error(t, err, `jwe.Encrypt should fail`)
+}
diff --git a/jwe/options_gen_test.go b/jwe/options_gen_test.go
index d36b9765..b58d2362 100644
--- a/jwe/options_gen_test.go
+++ b/jwe/options_gen_test.go
@@ -16,6 +16,7 @@ func TestOptionIdent(t *testing.T) {
 	require.Equal(t, "WithKey", identKey{}.String())
 	require.Equal(t, "WithKeyProvider", identKeyProvider{}.String())
 	require.Equal(t, "WithKeyUsed", identKeyUsed{}.String())
+	require.Equal(t, "WithMaxBufferSize", identMaxBufferSize{}.String())
 	require.Equal(t, "WithMaxPBES2Count", identMaxPBES2Count{}.String())
 	require.Equal(t, "WithMergeProtectedHeaders", identMergeProtectedHeaders{}.String())
 	require.Equal(t, "WithMessage", identMessage{}.String())