From f9aa8c5febcb31f1139fb9ae824f004054f210f9 Mon Sep 17 00:00:00 2001
From: Mike Jensen <jentfoo@users.noreply.github.com>
Date: Thu, 22 Jun 2023 15:22:13 -0600
Subject: [PATCH] GODRIVER-2869 Two protocol validations to reduce client
 denial of service risks (#1291)

Co-authored-by: Alan Parra <alan.parra@goteleport.com>
Co-authored-by: Qingyang Hu <103950869+qingyang-hu@users.noreply.github.com>
---
 x/bsonx/bsoncore/bsoncore.go       |  3 +++
 x/bsonx/bsoncore/bsoncore_test.go  |  9 +++++++++
 x/mongo/driver/compression.go      |  6 ++++++
 x/mongo/driver/compression_test.go | 17 +++++++++++++++++
 4 files changed, 35 insertions(+)

diff --git a/x/bsonx/bsoncore/bsoncore.go b/x/bsonx/bsoncore/bsoncore.go
index 94d479428f4..e52674aacf9 100644
--- a/x/bsonx/bsoncore/bsoncore.go
+++ b/x/bsonx/bsoncore/bsoncore.go
@@ -825,6 +825,9 @@ func readLengthBytes(src []byte) ([]byte, []byte, bool) {
 	if !ok {
 		return nil, src, false
 	}
+	if l < 4 {
+		return nil, src, false
+	}
 	if len(src) < int(l) {
 		return nil, src, false
 	}
diff --git a/x/bsonx/bsoncore/bsoncore_test.go b/x/bsonx/bsoncore/bsoncore_test.go
index b7d91a77153..ba2688ebe46 100644
--- a/x/bsonx/bsoncore/bsoncore_test.go
+++ b/x/bsonx/bsoncore/bsoncore_test.go
@@ -14,6 +14,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+
 	"go.mongodb.org/mongo-driver/bson/bsontype"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/internal/assert"
@@ -943,6 +944,14 @@ func TestNullBytes(t *testing.T) {
 	})
 }
 
+func TestInvalidBytes(t *testing.T) {
+	t.Run("read length less than 4 int bytes", func(t *testing.T) {
+		_, src, ok := readLengthBytes([]byte{0x00, 0x00, 0x00, 0x01})
+		assert.False(t, ok, "expected not ok response for invalid length read")
+		assert.Equal(t, 4, len(src), "expected src to contain the size parameter still")
+	})
+}
+
 func compareDecimal128(d1, d2 primitive.Decimal128) bool {
 	d1H, d1L := d1.GetBytes()
 	d2H, d2L := d2.GetBytes()
diff --git a/x/mongo/driver/compression.go b/x/mongo/driver/compression.go
index c474714ff48..7f355f61a44 100644
--- a/x/mongo/driver/compression.go
+++ b/x/mongo/driver/compression.go
@@ -111,6 +111,12 @@ func DecompressPayload(in []byte, opts CompressionOpts) (uncompressed []byte, er
 	case wiremessage.CompressorNoOp:
 		return in, nil
 	case wiremessage.CompressorSnappy:
+		l, err := snappy.DecodedLen(in)
+		if err != nil {
+			return nil, fmt.Errorf("decoding compressed length %w", err)
+		} else if int32(l) != opts.UncompressedSize {
+			return nil, fmt.Errorf("unexpected decompression size, expected %v but got %v", opts.UncompressedSize, l)
+		}
 		uncompressed = make([]byte, opts.UncompressedSize)
 		return snappy.Decode(uncompressed, in)
 	case wiremessage.CompressorZLib:
diff --git a/x/mongo/driver/compression_test.go b/x/mongo/driver/compression_test.go
index 5c74bef9aee..8a37de65c27 100644
--- a/x/mongo/driver/compression_test.go
+++ b/x/mongo/driver/compression_test.go
@@ -41,6 +41,23 @@ func TestCompression(t *testing.T) {
 	}
 }
 
+func TestDecompressFailures(t *testing.T) {
+	t.Run("snappy decompress huge size", func(t *testing.T) {
+		opts := CompressionOpts{
+			Compressor:       wiremessage.CompressorSnappy,
+			UncompressedSize: 100, // reasonable size
+		}
+		// Compressed data is twice as large as declared above.
+		// In test we use actual compression so that the decompress action would pass without fix (thus failing test).
+		// When decompression starts it allocates a buffer of the defined size, regardless of a valid compressed body following.
+		compressedData, err := CompressPayload(make([]byte, opts.UncompressedSize*2), opts)
+		assert.NoError(t, err, "premature error making compressed example")
+
+		_, err = DecompressPayload(compressedData, opts)
+		assert.Error(t, err)
+	})
+}
+
 func BenchmarkCompressPayload(b *testing.B) {
 	payload := func() []byte {
 		buf, err := os.ReadFile("compression.go")