github.com/tetratelabs/proxy-wasm-go-sdk@v0.23.1-0.20240517021853-021aa9cf78e8/proxywasm/internal/abi_callback_lifecycle_test.go (about)

     1  // Copyright 2021 Tetrate
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package internal
    16  
    17  import (
    18  	"testing"
    19  
    20  	"github.com/stretchr/testify/require"
    21  
    22  	"github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm/types"
    23  )
    24  
    25  type testOnContextCreatePluginVMContext struct{}
    26  
    27  func (*testOnContextCreatePluginVMContext) OnVMStart(vmConfigurationSize int) types.OnVMStartStatus {
    28  	return types.OnVMStartStatusOK
    29  }
    30  
    31  func (*testOnContextCreatePluginVMContext) NewPluginContext(contextID uint32) types.PluginContext {
    32  	return &testOnContextCreatePluginContext{cnt: 1}
    33  }
    34  
    35  type testOnContextCreatePluginContext struct {
    36  	types.DefaultPluginContext
    37  	cnt int
    38  }
    39  
    40  func (ctx *testOnContextCreatePluginContext) NewTcpContext(contextID uint32) types.TcpContext {
    41  	if contextID == 100 {
    42  		ctx.cnt += 100
    43  		return &types.DefaultTcpContext{}
    44  	}
    45  	return nil
    46  }
    47  
    48  func (ctx *testOnContextCreatePluginContext) NewHttpContext(contextID uint32) types.HttpContext {
    49  	if contextID == 1000 {
    50  		ctx.cnt += 1000
    51  		return &types.DefaultHttpContext{}
    52  	}
    53  	return nil
    54  }
    55  
    56  func Test_proxyOnContextCreate(t *testing.T) {
    57  	currentStateMux.Lock()
    58  	defer currentStateMux.Unlock()
    59  
    60  	vmContext := &testOnContextCreatePluginVMContext{}
    61  	currentState = &state{
    62  		pluginContexts:    map[uint32]*pluginContextState{},
    63  		httpContexts:      map[uint32]types.HttpContext{},
    64  		tcpContexts:       map[uint32]types.TcpContext{},
    65  		contextIDToRootID: map[uint32]uint32{},
    66  	}
    67  
    68  	// Set the VM context.
    69  	SetVMContext(vmContext)
    70  
    71  	// Create Plugin context.
    72  	proxyOnContextCreate(1, 0)
    73  	require.Contains(t, currentState.pluginContexts, uint32(1))
    74  	pluginContext := currentState.pluginContexts[1].context.(*testOnContextCreatePluginContext)
    75  	require.Equal(t, 1, pluginContext.cnt)
    76  
    77  	// Create Http contexts.
    78  	proxyOnContextCreate(100, 1)
    79  	require.Equal(t, 101, pluginContext.cnt)
    80  	proxyOnContextCreate(1000, 1)
    81  	require.Equal(t, 1101, pluginContext.cnt)
    82  }
    83  
    84  type lifecycleContext struct {
    85  	types.DefaultPluginContext
    86  	types.DefaultHttpContext
    87  	types.DefaultTcpContext
    88  	onDoneCalled bool
    89  }
    90  
    91  func (ctx *lifecycleContext) OnPluginDone() bool {
    92  	ctx.onDoneCalled = true
    93  	return true
    94  }
    95  
    96  func (ctx *lifecycleContext) OnStreamDone() {
    97  	ctx.onDoneCalled = true
    98  }
    99  
   100  func (ctx *lifecycleContext) OnHttpStreamDone() {
   101  	ctx.onDoneCalled = true
   102  }
   103  
   104  func Test_onDone_or_onLog(t *testing.T) {
   105  	currentStateMux.Lock()
   106  	defer currentStateMux.Unlock()
   107  
   108  	currentState = &state{
   109  		pluginContexts: map[uint32]*pluginContextState{},
   110  		httpContexts:   map[uint32]types.HttpContext{},
   111  		tcpContexts:    map[uint32]types.TcpContext{},
   112  	}
   113  
   114  	// Stream Contexts are only called on on_log, not on on_done.
   115  	var id uint32 = 1
   116  	ctx := &lifecycleContext{}
   117  	currentState.httpContexts[id] = ctx
   118  	proxyOnLog(id)
   119  	require.True(t, ctx.onDoneCalled)
   120  	require.Equal(t, id, currentState.activeContextID)
   121  
   122  	id = 2
   123  	ctx = &lifecycleContext{}
   124  	currentState.tcpContexts[id] = ctx
   125  	proxyOnLog(id)
   126  	require.True(t, ctx.onDoneCalled)
   127  	require.Equal(t, id, currentState.activeContextID)
   128  
   129  	// Root Contexts are only called on on_done, not on on_log.
   130  	id = 3
   131  	ctx = &lifecycleContext{}
   132  	currentState.pluginContexts[id] = &pluginContextState{context: ctx}
   133  	proxyOnDone(id)
   134  	require.True(t, ctx.onDoneCalled)
   135  	require.Equal(t, id, currentState.activeContextID)
   136  }