From 0cd7b49eb4786c8c0253a280a412eb94941626d4 Mon Sep 17 00:00:00 2001
From: Gabriel Paradiso <gaboparadiso@gmail.com>
Date: Tue, 14 Jan 2025 15:11:39 +0100
Subject: [PATCH] [CRE-40] Check binary size before decompression (#994)

fix: check binary size before decompression
---
 pkg/workflows/wasm/host/module.go    | 40 ++++++++++++++++++----------
 pkg/workflows/wasm/host/wasm_test.go | 40 ++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 14 deletions(-)

diff --git a/pkg/workflows/wasm/host/module.go b/pkg/workflows/wasm/host/module.go
index f8bd77d02..b2645b8a3 100644
--- a/pkg/workflows/wasm/host/module.go
+++ b/pkg/workflows/wasm/host/module.go
@@ -70,11 +70,12 @@ func (r *store) delete(id string) {
 }
 
 var (
-	defaultTickInterval     = 100 * time.Millisecond
-	defaultTimeout          = 10 * time.Second
-	defaultMinMemoryMBs     = 128
-	DefaultInitialFuel      = uint64(100_000_000)
-	defaultMaxFetchRequests = 5
+	defaultTickInterval            = 100 * time.Millisecond
+	defaultTimeout                 = 10 * time.Second
+	defaultMinMemoryMBs            = 128
+	DefaultInitialFuel             = uint64(100_000_000)
+	defaultMaxFetchRequests        = 5
+	defaultMaxCompressedBinarySize = 10 * 1024 * 1024 // 10 MB
 )
 
 type DeterminismConfig struct {
@@ -82,15 +83,16 @@ type DeterminismConfig struct {
 	Seed int64
 }
 type ModuleConfig struct {
-	TickInterval     time.Duration
-	Timeout          *time.Duration
-	MaxMemoryMBs     int64
-	MinMemoryMBs     int64
-	InitialFuel      uint64
-	Logger           logger.Logger
-	IsUncompressed   bool
-	Fetch            func(ctx context.Context, req *wasmpb.FetchRequest) (*wasmpb.FetchResponse, error)
-	MaxFetchRequests int
+	TickInterval            time.Duration
+	Timeout                 *time.Duration
+	MaxMemoryMBs            int64
+	MinMemoryMBs            int64
+	InitialFuel             uint64
+	Logger                  logger.Logger
+	IsUncompressed          bool
+	Fetch                   func(ctx context.Context, req *wasmpb.FetchRequest) (*wasmpb.FetchResponse, error)
+	MaxFetchRequests        int
+	MaxCompressedBinarySize uint64
 
 	// Labeler is used to emit messages from the module.
 	Labeler custmsg.MessageEmitter
@@ -166,6 +168,10 @@ func NewModule(modCfg *ModuleConfig, binary []byte, opts ...func(*ModuleConfig))
 		modCfg.MinMemoryMBs = int64(defaultMinMemoryMBs)
 	}
 
+	if modCfg.MaxCompressedBinarySize == 0 {
+		modCfg.MaxCompressedBinarySize = uint64(defaultMaxCompressedBinarySize)
+	}
+
 	// Take the max of the min and the configured max memory mbs.
 	// We do this because Go requires a minimum of 16 megabytes to run,
 	// and local testing has shown that with less than the min, some
@@ -183,6 +189,12 @@ func NewModule(modCfg *ModuleConfig, binary []byte, opts ...func(*ModuleConfig))
 
 	engine := wasmtime.NewEngineWithConfig(cfg)
 	if !modCfg.IsUncompressed {
+		// validate the binary size before decompressing
+		// this is to prevent decompression bombs
+		if uint64(len(binary)) > modCfg.MaxCompressedBinarySize {
+			return nil, fmt.Errorf("binary size exceeds the maximum allowed size of %d bytes", modCfg.MaxCompressedBinarySize)
+		}
+
 		rdr := brotli.NewReader(bytes.NewBuffer(binary))
 		decompedBinary, err := io.ReadAll(rdr)
 		if err != nil {
diff --git a/pkg/workflows/wasm/host/wasm_test.go b/pkg/workflows/wasm/host/wasm_test.go
index 3e5335a9d..b75170914 100644
--- a/pkg/workflows/wasm/host/wasm_test.go
+++ b/pkg/workflows/wasm/host/wasm_test.go
@@ -903,6 +903,46 @@ func TestModule_Sandbox_Memory(t *testing.T) {
 	assert.ErrorContains(t, err, "exit status 2")
 }
 
+func TestModule_CompressedBinarySize(t *testing.T) {
+	t.Parallel()
+
+	t.Run("compressed binary size is smaller than the default 10mb limit", func(t *testing.T) {
+		binary := createTestBinary(successBinaryCmd, successBinaryLocation, false, t)
+
+		_, err := NewModule(&ModuleConfig{IsUncompressed: false, Logger: logger.Test(t)}, binary)
+		require.NoError(t, err)
+	})
+
+	t.Run("compressed binary size is bigger than the default 10mb limit", func(t *testing.T) {
+		binary := make([]byte, defaultMaxCompressedBinarySize+1)
+
+		var b bytes.Buffer
+		bwr := brotli.NewWriter(&b)
+		_, err := bwr.Write(binary)
+		require.NoError(t, err)
+		require.NoError(t, bwr.Close())
+
+		_, err = NewModule(&ModuleConfig{IsUncompressed: false, Logger: logger.Test(t)}, binary)
+		default10mbLimit := fmt.Sprintf("binary size exceeds the maximum allowed size of %d bytes", defaultMaxCompressedBinarySize)
+		require.ErrorContains(t, err, default10mbLimit)
+	})
+
+	t.Run("compressed binary size is bigger than the custom limit", func(t *testing.T) {
+		customMaxCompressedBinarySize := uint64(1 * 1024 * 1024)
+		binary := make([]byte, customMaxCompressedBinarySize+1)
+
+		var b bytes.Buffer
+		bwr := brotli.NewWriter(&b)
+		_, err := bwr.Write(binary)
+		require.NoError(t, err)
+		require.NoError(t, bwr.Close())
+
+		_, err = NewModule(&ModuleConfig{IsUncompressed: false, MaxCompressedBinarySize: customMaxCompressedBinarySize, Logger: logger.Test(t)}, binary)
+		default10mbLimit := fmt.Sprintf("binary size exceeds the maximum allowed size of %d bytes", customMaxCompressedBinarySize)
+		require.ErrorContains(t, err, default10mbLimit)
+	})
+}
+
 func TestModule_Sandbox_SleepIsStubbedOut(t *testing.T) {
 	t.Parallel()
 	ctx := tests.Context(t)