From f4c051a0653d78199a053892f7619ebf96339502 Mon Sep 17 00:00:00 2001
From: Jacob Hoffman-Andrews <github@hoffman-andrews.com>
Date: Thu, 7 Mar 2024 11:46:13 -0800
Subject: [PATCH] jwe: limit maximum output from decompressing (#106)

---
 CHANGELOG.md     | 11 +++++++++++
 crypter.go       |  6 ++++++
 encoding.go      | 18 ++++++++++++++----
 encoding_test.go | 34 ++++++++++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 08e43ab..28bdd2f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,14 @@
+# v4.0.1
+
+## Fixed
+
+ - An attacker could send a JWE containing compressed data that used large
+   amounts of memory and CPU when decompressed by `Decrypt` or `DecryptMulti`.
+   Those functions now return an error if the decompressed data would exceed
+   250kB or 10x the compressed size (whichever is larger). Thanks to
+   Enze Wang@Alioth and Jianjun Chen@Zhongguancun Lab (@zer0yu and @chenjj)
+   for reporting.
+
 # v4.0.0
 
 This release makes some breaking changes in order to more thoroughly
diff --git a/crypter.go b/crypter.go
index 073767f..aba0842 100644
--- a/crypter.go
+++ b/crypter.go
@@ -440,6 +440,9 @@ func (ctx *genericEncrypter) Options() EncrypterOptions {
 //
 // Note that ed25519 is only available for signatures, not encryption, so is
 // not an option here.
+//
+// Automatically decompresses plaintext, but returns an error if the decompressed
+// data would be >250kB or >10x the size of the compressed data, whichever is larger.
 func (obj JSONWebEncryption) Decrypt(decryptionKey interface{}) ([]byte, error) {
 	headers := obj.mergedHeaders(nil)
 
@@ -511,6 +514,9 @@ func (obj JSONWebEncryption) Decrypt(decryptionKey interface{}) ([]byte, error)
 //
 // The decryptionKey argument must have one of the types allowed for the
 // decryptionKey argument of Decrypt().
+//
+// Automatically decompresses plaintext, but returns an error if the decompressed
+// data would be >250kB or >3x the size of the compressed data, whichever is larger.
 func (obj JSONWebEncryption) DecryptMulti(decryptionKey interface{}) (int, Header, []byte, error) {
 	globalHeaders := obj.mergedHeaders(nil)
 
diff --git a/encoding.go b/encoding.go
index bc71404..4f6e0d4 100644
--- a/encoding.go
+++ b/encoding.go
@@ -21,6 +21,7 @@ import (
 	"compress/flate"
 	"encoding/base64"
 	"encoding/binary"
+	"fmt"
 	"io"
 	"math/big"
 	"strings"
@@ -85,7 +86,7 @@ func decompress(algorithm CompressionAlgorithm, input []byte) ([]byte, error) {
 	}
 }
 
-// Compress with DEFLATE
+// deflate compresses the input.
 func deflate(input []byte) ([]byte, error) {
 	output := new(bytes.Buffer)
 
@@ -97,15 +98,24 @@ func deflate(input []byte) ([]byte, error) {
 	return output.Bytes(), err
 }
 
-// Decompress with DEFLATE
+// inflate decompresses the input.
+//
+// Errors if the decompressed data would be >250kB or >10x the size of the
+// compressed data, whichever is larger.
 func inflate(input []byte) ([]byte, error) {
 	output := new(bytes.Buffer)
 	reader := flate.NewReader(bytes.NewBuffer(input))
 
-	_, err := io.Copy(output, reader)
-	if err != nil {
+	maxCompressedSize := max(250_000, 10*int64(len(input)))
+
+	limit := maxCompressedSize + 1
+	n, err := io.CopyN(output, reader, limit)
+	if err != nil && err != io.EOF {
 		return nil, err
 	}
+	if n == limit {
+		return nil, fmt.Errorf("uncompressed data would be too large (>%d bytes)", maxCompressedSize)
+	}
 
 	err = reader.Close()
 	return output.Bytes(), err
diff --git a/encoding_test.go b/encoding_test.go
index fc48685..4b1d451 100644
--- a/encoding_test.go
+++ b/encoding_test.go
@@ -18,6 +18,8 @@ package jose
 
 import (
 	"bytes"
+	"crypto/rand"
+	"io"
 	"strings"
 	"testing"
 )
@@ -57,6 +59,38 @@ func TestInvalidCompression(t *testing.T) {
 	}
 }
 
+// TestLargeZip tests that we can decompress a large input, so long as its
+// compression ratio is reasonable.
+func TestLargeZip(t *testing.T) {
+	input := new(bytes.Buffer)
+	_, err := io.CopyN(input, rand.Reader, 251_000)
+	if err != nil {
+		t.Fatalf("generating input: %s", err)
+	}
+	compressed, err := compress(DEFLATE, input.Bytes())
+	if err != nil {
+		t.Errorf("compressing: %s", err)
+	}
+	t.Logf("compression ratio: %g", float64(len(input.Bytes()))/float64(len(compressed)))
+	_, err = decompress(DEFLATE, compressed)
+	if err != nil {
+		t.Errorf("decompressing large input with low compression ratio: %s", err)
+	}
+}
+
+func TestZipBomb(t *testing.T) {
+	input := strings.Repeat("a", 251_000)
+	compressed, err := compress(DEFLATE, []byte(input))
+	if err != nil {
+		t.Errorf("compressing: %s", err)
+	}
+	t.Logf("compression ratio: %d %g", len(compressed), float64(len(input))/float64(len(compressed)))
+	out, err := decompress(DEFLATE, compressed)
+	if err == nil {
+		t.Errorf("expected error decompressing zip bomb, got none. output size %d", len(out))
+	}
+}
+
 func TestByteBufferTrim(t *testing.T) {
 	buf := newBufferFromInt(1)
 	if !bytes.Equal(buf.data, []byte{1}) {