Repository: tetratelabs/wazero
Pull request #176 (2022-02-01 04:50:27)

Add limits on number of function instances in a store
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/RATIONALE.md b/RATIONALE.md
index 2a05da2e..02afadfd 100644
--- a/RATIONALE.md
+++ b/RATIONALE.md
@@ -12,6 +12,22 @@ runtime vs interpreting Wasm directly (the `naivevm` interpreter).
 Note: `microwasm` was never specified formally, and only exists in a historical codebase of wasmtime:
 https://github.com/bytecodealliance/wasmtime/blob/v0.29.0/crates/lightbeam/src/microwasm.rs
 
+
+### Size limitations
+#### Number of functions
+
+The possible number of function instances in [a store](https://www.w3.org/TR/wasm-core-1/#store%E2%91%A0) is not specified in the WebAssembly specifications since [`funcaddr`](https://www.w3.org/TR/wasm-core-1/#syntax-funcaddr) corresponding to a function instance can be arbitrary number. 
+In wazero, we choose to use `uint32` to represent `funcaddr`. Therefore the maximum number of function instances a store can instantiate is limited to 2^32. 
+
+That is because not only we _believe_ that all use cases are fine with the limitation, but also we have no way to test wazero runtimes under these unusual circumstances.
+
+#### Number of function types
+
+There's no limitation on the number of function types in [a store](https://www.w3.org/TR/wasm-core-1/#store%E2%91%A0) according to the spec. In wazero implementation, we assign each function type to a unique ID, and choose to use `uint32` to represent the IDs.
+Therefore the maximum number of function types a store can have is limited to 2^32. 
+
+This is due to the same reason for the limitation on the number of functions above.
+
 ### JIT engine implementation
 
 See [wasm/jit/RATIONALE.md](wasm/jit/RATIONALE.md).
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index f5f52594..cd017d73 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -1039,7 +1039,7 @@ func (c *amd64Compiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) e
 	// So the target type ID lives at offset+tableElementTypeIDOffset, and we compare it
 	// with wasm.UninitializedTableElementTypeID to check if the element is initialized.
 	checkIfInitialized := c.newProg()
-	checkIfInitialized.As = x86.ACMPQ
+	checkIfInitialized.As = x86.ACMPL // 32-bit as FunctionTypeID is in 32-bit unsigned integer.
 	checkIfInitialized.From.Type = obj.TYPE_MEM
 	checkIfInitialized.From.Reg = offset.register
 	checkIfInitialized.From.Offset = tableElementFunctionTypeIDOffset
@@ -1059,7 +1059,7 @@ func (c *amd64Compiler) compileCallIndirect(o *wazeroir.OperationCallIndirect) e
 	targetFunctionType := c.f.ModuleInstance.Types[o.TypeIndex]
 	checkIfTypeMatch := c.newProg()
 	jumpIfInitialized.To.SetTarget(checkIfTypeMatch)
-	checkIfTypeMatch.As = x86.ACMPQ
+	checkIfTypeMatch.As = x86.ACMPL // 32-bit as FunctionTypeID is in 32-bit unsigned integer.
 	checkIfTypeMatch.From.Type = obj.TYPE_MEM
 	checkIfTypeMatch.From.Reg = offset.register
 	checkIfTypeMatch.From.Offset = tableElementFunctionTypeIDOffset
diff --git a/wasm/store.go b/wasm/store.go
index 0ec996ad..ed7613d7 100644
--- a/wasm/store.go
+++ b/wasm/store.go
@@ -35,6 +35,9 @@ type (
 		// do type-checks on indirect function calls.
 		TypeIDs map[string]FunctionTypeID
 
+		// maximumFunctionAddress and maximumFunctionTypes represent the limit on the number of each instance type in a store.
+		maximumFunctionAddress, maximumFunctionTypes int
+
 		// The followings fields match the definition of Store in the specification.
 
 		// Functions holds function instances (https://www.w3.org/TR/wasm-core-1/#function-instances%E2%91%A0),
@@ -138,7 +141,7 @@ type (
 		// Note: we intentionally use "[]TableElement", not "[]*TableElement",
 		// because the JIT engine accesses this slice directly from assembly.
 		// If pointer type is used, the access becomes two level indirection (two hops of pointer jumps)
-		// which is a bit costly. TableElement is 128 bit (two 64bit fields) so the cost of using value type
+		// which is a bit costly. TableElement is 96 bit (32 and 64 bit fields) so the cost of using value type
 		// would be ignorable.
 		Table []TableElement
 		Min   uint32
@@ -175,7 +178,12 @@ type (
 	// FunctionTypeID is an uniquely assigned integer for a function type.
 	// This is wazero specific runtime object and specific to a store,
 	// and used at runtime to do type-checks on indirect function calls.
-	FunctionTypeID uint64
+	FunctionTypeID uint32
+)
+
+const (
+	maximumFunctionAddress = math.MaxUint32
+	maximumFunctionTypes   = maximumFunctionAddress
 )
 
 // addExport adds and indexes the given export or errs if the name is already exported.
@@ -204,13 +212,23 @@ func (f *FunctionInstance) IsHostFunction() bool {
 }
 
 func NewStore(engine Engine) *Store {
-	return &Store{ModuleInstances: map[string]*ModuleInstance{}, TypeIDs: map[string]FunctionTypeID{}, engine: engine}
+	return &Store{
+		ModuleInstances:        map[string]*ModuleInstance{},
+		TypeIDs:                map[string]FunctionTypeID{},
+		engine:                 engine,
+		maximumFunctionAddress: maximumFunctionAddress,
+		maximumFunctionTypes:   maximumFunctionTypes,
+	}
 }
 
 func (s *Store) Instantiate(module *Module, name string) error {
 	instance := &ModuleInstance{Name: name}
 	for _, t := range module.TypeSection {
-		instance.Types = append(instance.Types, s.getTypeInstance(t))
+		typeInstance, err := s.getTypeInstance(t)
+		if err != nil {
+			return err
+		}
+		instance.Types = append(instance.Types, typeInstance)
 	}
 
 	s.ModuleInstances[name] = instance
@@ -310,9 +328,14 @@ func (s *Store) getExport(moduleName string, name string, kind ExportKind) (exp
 	return
 }
 
-func (s *Store) addFunctionInstance(f *FunctionInstance) {
+func (s *Store) addFunctionInstance(f *FunctionInstance) error {
+	l := len(s.Functions)
+	if l >= s.maximumFunctionAddress {
+		return fmt.Errorf("too many functions in a store")
+	}
 	f.Address = FunctionAddress(len(s.Functions))
 	s.Functions = append(s.Functions, f)
+	return nil
 }
 
 func (s *Store) resolveImports(module *Module, target *ModuleInstance) error {
@@ -576,9 +599,14 @@ func (s *Store) buildFunctionInstances(module *Module, target *ModuleInstance) (
 			}
 		}
 
+		typeInstace, err := s.getTypeInstance(module.TypeSection[typeIndex])
+		if err != nil {
+			return rollbackFuncs, err
+		}
+
 		f := &FunctionInstance{
 			Name:           name,
-			FunctionType:   s.getTypeInstance(module.TypeSection[typeIndex]),
+			FunctionType:   typeInstace,
 			Body:           module.CodeSection[codeIndex].Body,
 			LocalTypes:     module.CodeSection[codeIndex].LocalTypes,
 			ModuleInstance: target,
@@ -595,7 +623,10 @@ func (s *Store) buildFunctionInstances(module *Module, target *ModuleInstance) (
 		}
 
 		target.Functions = append(target.Functions, f)
-		s.addFunctionInstance(f)
+		err = s.addFunctionInstance(f)
+		if err != nil {
+			return rollbackFuncs, err
+		}
 	}
 	return rollbackFuncs, nil
 }
@@ -1757,17 +1788,24 @@ func (s *Store) AddHostFunction(moduleName, funcName string, fn reflect.Value) e
 		return fmt.Errorf("invalid signature: %w", err)
 	}
 
+	typeInstace, err := s.getTypeInstance(sig)
+	if err != nil {
+		return err
+	}
+
 	f := &FunctionInstance{
 		Name:           fmt.Sprintf("%s.%s", moduleName, funcName),
 		HostFunction:   &fn,
-		FunctionType:   s.getTypeInstance(sig),
+		FunctionType:   typeInstace,
 		ModuleInstance: m,
 	}
 
 	if err = s.engine.Compile(f); err != nil {
 		return fmt.Errorf("failed to compile %s: %v", f.Name, err)
 	}
-	s.addFunctionInstance(f)
+	if err = s.addFunctionInstance(f); err != nil {
+		return err
+	}
 	if err = m.addExport(funcName, &ExportInstance{Kind: ExportKindFunc, Function: f}); err != nil {
 		s.Functions = s.Functions[:len(s.Functions)-1] // revert the add on conflict
 		return err
@@ -1805,14 +1843,18 @@ func (s *Store) AddMemoryInstance(moduleName, name string, min uint32, max *uint
 	return m.addExport(name, &ExportInstance{Kind: ExportKindMemory, Memory: memory})
 }
 
-func (s *Store) getTypeInstance(t *FunctionType) *TypeInstance {
+func (s *Store) getTypeInstance(t *FunctionType) (*TypeInstance, error) {
 	key := t.String()
 	id, ok := s.TypeIDs[key]
 	if !ok {
+		l := len(s.TypeIDs)
+		if l >= s.maximumFunctionTypes {
+			return nil, fmt.Errorf("too many function types in a store")
+		}
 		id = FunctionTypeID(len(s.TypeIDs))
 		s.TypeIDs[key] = id
 	}
-	return &TypeInstance{Type: t, TypeID: id}
+	return &TypeInstance{Type: t, TypeID: id}, nil
 }
 
 // getModuleInstance returns an existing ModuleInstance if exists, or assigns a new one.
@@ -1841,4 +1883,4 @@ func newTableInstance(min uint32, max *uint32) *TableInstance {
 }
 
 // UninitializedTableElementTypeID math.MaxUint64 to represent the uninitialized elements.
-var UninitializedTableElementTypeID FunctionTypeID = math.MaxUint64
+var UninitializedTableElementTypeID FunctionTypeID = math.MaxUint32


Test patch:
diff --git a/wasm/store_test.go b/wasm/store_test.go
index 89d54fa6..dc58e5d5 100644
--- a/wasm/store_test.go
+++ b/wasm/store_test.go
@@ -2,6 +2,7 @@ package wasm
 
 import (
 	"reflect"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -95,3 +96,64 @@ func (e *nopEngine) Call(_ *FunctionInstance, _ ...uint64) (results []uint64, er
 func (e *nopEngine) Compile(_ *FunctionInstance) error {
 	return nil
 }
+
+func TestStore_addFunctionInstance(t *testing.T) {
+	t.Run("too many functions", func(t *testing.T) {
+		s := NewStore(nopEngineInstance)
+		const max = 10
+		s.maximumFunctionAddress = max
+		s.Functions = make([]*FunctionInstance, max)
+		err := s.addFunctionInstance(nil)
+		require.Error(t, err)
+	})
+	t.Run("ok", func(t *testing.T) {
+		s := NewStore(nopEngineInstance)
+		for i := 0; i < 10; i++ {
+			f := &FunctionInstance{}
+			require.Len(t, s.Functions, i)
+
+			err := s.addFunctionInstance(f)
+			require.NoError(t, err)
+
+			// After the addition, one instance is added.
+			require.Len(t, s.Functions, i+1)
+
+			// The added function instance must have i for its funcaddr.
+			require.Equal(t, FunctionAddress(i), f.Address)
+		}
+	})
+}
+
+func TestStore_getTypeInstance(t *testing.T) {
+	t.Run("too many functions", func(t *testing.T) {
+		s := NewStore(nopEngineInstance)
+		const max = 10
+		s.maximumFunctionTypes = max
+		s.TypeIDs = make(map[string]FunctionTypeID)
+		for i := 0; i < max; i++ {
+			s.TypeIDs[strconv.Itoa(i)] = 0
+		}
+		_, err := s.getTypeInstance(&FunctionType{})
+		require.Error(t, err)
+	})
+	t.Run("ok", func(t *testing.T) {
+		for _, tc := range []*FunctionType{
+			{Params: []ValueType{}},
+			{Params: []ValueType{ValueTypeF32}},
+			{Results: []ValueType{ValueTypeF64}},
+			{Params: []ValueType{ValueTypeI32}, Results: []ValueType{ValueTypeI64}},
+		} {
+			tc := tc
+			t.Run(tc.String(), func(t *testing.T) {
+				s := NewStore(nopEngineInstance)
+				actual, err := s.getTypeInstance(tc)
+				require.NoError(t, err)
+
+				expectedTypeID, ok := s.TypeIDs[tc.String()]
+				require.True(t, ok)
+				require.Equal(t, expectedTypeID, actual.TypeID)
+				require.Equal(t, tc, actual.Type)
+			})
+		}
+	})
+}