BAAS-34801: add forced mem checks on new arrays and obj props (#127)

mongodb-forks · Dec 9, 2024 · 8b50f5b · 8b50f5b
1 parent 476c6e0
commit 8b50f5b
Show file tree

Hide file tree

Showing 29 changed files with 237 additions and 71 deletions.
diff --git a/added_values_test.go b/added_values_test.go
@@ -48,8 +48,6 @@ func TestIntStringEquality(t *testing.T) {
 }
 
 func TestAddedValuesMemUsage(t *testing.T) {
-	vm := New()
-
 	for _, tc := range []struct {
 		name        string
 		val         MemUsageReporter
@@ -77,7 +75,7 @@ func TestAddedValuesMemUsage(t *testing.T) {
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			mem, err := tc.val.MemUsage(NewMemUsageContext(vm, 100, 100, 100, 100, 0.1, nil))
+			mem, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != nil {
 				t.Fatalf("Unexpected error. Actual: %v Expected: nil", err)
 			}

diff --git a/array_sparse_test.go b/array_sparse_test.go
@@ -274,7 +274,7 @@ func TestSparseArrayObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao: &sparseArrayObject{
 				items: []sparseArrayItem{
 					{
@@ -289,14 +289,14 @@ func TestSparseArrayObjectMemUsage(t *testing.T) {
 		},
 		{
 			name:        "mem is SizeEmptyStruct for nil sparse array",
-			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:          NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao:         nil,
 			expected:    SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao: &sparseArrayObject{
 				items: []sparseArrayItem{
 					{
@@ -331,7 +331,7 @@ func TestSparseArrayObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem above estimate threshold and within memory limit returns correct usage",
-			mu:   NewMemUsageContext(vm, 88, 100, 5, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 5, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao: &sparseArrayObject{
 				items: []sparseArrayItem{
 					{

diff --git a/array_test.go b/array_test.go
@@ -144,23 +144,23 @@ func TestArrayObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold given a nil slice of values",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao:   &arrayObject{},
 			// array overhead + array baseObject
 			expectedMem: SizeEmptyStruct + SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem below threshold given empty slice of values",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao:   &arrayObject{values: []Value{}},
 			// array overhead + array baseObject + values slice overhead
 			expectedMem: SizeEmptyStruct + SizeEmptyStruct + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao: &arrayObject{
 				values: []Value{
 					vm.ToValue("key0"),
@@ -179,7 +179,7 @@ func TestArrayObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "empty array with negative threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, -1, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, -1, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao: &arrayObject{
 				values: []Value{},
 			},

diff --git a/builtin_array.go b/builtin_array.go
@@ -80,7 +80,18 @@ func relToIdx(rel, l int64) int64 {
 	return max(l+rel, 0)
 }
 
-func (r *Runtime) newArrayValues(values []Value) *Object {
+func (r *Runtime) newArrayValues(values valueStack) *Object {
+	if r.shouldForceMemCheck {
+		if memCtx := newMemUsageContextClone(); memCtx != nil {
+			memUsage, err := valuesMemUsage(values, memCtx)
+			if err != nil {
+				panic(err)
+			}
+			if memCtx.MemUsageLimitExceeded(memUsage) {
+				panic(ErrMemLimitExceeded)
+			}
+		}
+	}
 	return setArrayValues(r.newArrayObject(), values).val
 }
 

diff --git a/builtin_arrray_test.go b/builtin_arrray_test.go
@@ -1,6 +1,9 @@
 package goja
 
-import "testing"
+import (
+	"context"
+	"testing"
+)
 
 func TestArrayProtoProp(t *testing.T) {
 	const SCRIPT = `
@@ -338,3 +341,53 @@ func TestArrayProto(t *testing.T) {
 	`
 	testScriptWithTestLib(SCRIPT, _undefined, t)
 }
+
+func TestNewArrayValues(t *testing.T) {
+	for _, tc := range []struct {
+		name                string
+		memLimit            uint64
+		expectedPanic       bool
+		shouldForceMemCheck bool
+	}{
+		{
+			name:                "should not panic when creating a new array under the mem limit",
+			memLimit:            1000,
+			expectedPanic:       false,
+			shouldForceMemCheck: false,
+		},
+		{
+			name:                "should panic when creating a new array over the mem limit and mem usage check is forced",
+			memLimit:            0,
+			expectedPanic:       true,
+			shouldForceMemCheck: true,
+		},
+		{
+			name:                "should not panic when creating a new array over the mem limit and mem usage check is not forced",
+			memLimit:            0,
+			expectedPanic:       false,
+			shouldForceMemCheck: false,
+		},
+		{
+			name:                "should not panic when creating a new array under the mem limit and mem usage check is forced",
+			memLimit:            1000,
+			expectedPanic:       false,
+			shouldForceMemCheck: true,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			defer func() {
+				r := recover()
+				if tc.expectedPanic && r == nil {
+					t.Error("The code is expected to panic, but it didn't")
+				}
+				if !tc.expectedPanic && r != nil {
+					t.Errorf("The code panicked, but it should not have: %v", r)
+				}
+			}()
+			vm := NewWithContext(context.Background(), tc.shouldForceMemCheck)
+			// Creating a mem usage context so it populates the package variable
+			NewMemUsageContext(100, tc.memLimit, 100, 100, 0.5, nil)
+			vm.newArrayValues([]Value{valueInt(0)})
+		})
+	}
+}
diff --git a/builtin_map_test.go b/builtin_map_test.go
@@ -286,7 +286,7 @@ func TestMapObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			mo: &mapObject{
 				m: &orderedMap{
 					hashTable: map[uint64]*mapEntry{
@@ -303,14 +303,14 @@ func TestMapObjectMemUsage(t *testing.T) {
 		},
 		{
 			name:        "mem is SizeEmptyStruct given a nil map object",
-			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:          NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			mo:          nil,
 			expectedMem: SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			mo: &mapObject{
 				m: &orderedMap{
 					hashTable: map[uint64]*mapEntry{
@@ -343,7 +343,7 @@ func TestMapObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem above estimate threshold and within memory limit returns correct mem usage",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 5, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 5, 0.1, TestNativeMemUsageChecker{}),
 			mo: &mapObject{
 				m: createOrderedMap(vm, 20),
 			},

diff --git a/builtin_proxy_test.go b/builtin_proxy_test.go
@@ -1290,7 +1290,7 @@ func TestBuiltinProxyMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != nil {
 				t.Fatalf("Unexpected error. Actual: %v Expected: nil", err)
 			}

diff --git a/builtin_set_test.go b/builtin_set_test.go
@@ -227,7 +227,7 @@ func TestSetObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			so: &setObject{
 				m: &orderedMap{
 					hashTable: map[uint64]*mapEntry{
@@ -244,14 +244,14 @@ func TestSetObjectMemUsage(t *testing.T) {
 		},
 		{
 			name:        "mem is SizeEmptyStruct given a nil map object",
-			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:          NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			so:          nil,
 			expectedMem: SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			so: &setObject{
 				m: &orderedMap{
 					hashTable: map[uint64]*mapEntry{
@@ -284,7 +284,7 @@ func TestSetObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem above estimate threshold and within memory limit returns correct mem usage",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 5, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 5, 0.1, TestNativeMemUsageChecker{}),
 			so: &setObject{
 				m: createOrderedMap(vm, 20),
 			},
@@ -298,7 +298,7 @@ func TestSetObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem above estimate threshold and within memory limit and nil values returns correct mem usage",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 1, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 100, 50, 1, 0.1, TestNativeMemUsageChecker{}),
 			so: &setObject{
 				m: createOrderedMapWithNilValues(3),
 			},
@@ -308,7 +308,7 @@ func TestSetObjectMemUsage(t *testing.T) {
 		},
 		{
 			name:        "mem is SizeEmptyStruct given a nil orderedMap object",
-			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:          NewMemUsageContext(88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			so:          &setObject{},
 			expectedMem: SizeEmptyStruct,
 			errExpected: nil,

diff --git a/compiler_test.go b/compiler_test.go
@@ -5657,7 +5657,7 @@ func TestProgramMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 50, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			p: &Program{
 				values: []Value{
 					New().newDateObject(time.Now(), true, nil),
@@ -5669,7 +5669,7 @@ func TestProgramMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(88, 50, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			p: &Program{
 				values: []Value{
 					New().newDateObject(time.Now(), true, nil),

diff --git a/date_test.go b/date_test.go
@@ -349,7 +349,7 @@ func TestDateMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}

diff --git a/destruct_test.go b/destruct_test.go
@@ -27,7 +27,7 @@ func TestDestructMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}

diff --git a/func_test.go b/func_test.go
@@ -184,7 +184,7 @@ func TestNativeFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}
@@ -234,7 +234,7 @@ func TestFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}
@@ -282,7 +282,7 @@ func TestBaseJsFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}
@@ -361,7 +361,7 @@ func TestClassFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}
@@ -421,7 +421,7 @@ func TestMethodFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}
@@ -490,7 +490,7 @@ func TestArrowFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}