add PluginState to share state between plugins (#357)

Signed-off-by: spacewander <[email protected]>

internal/plugin_state/plugin_state.go

@@ -0,0 +1,45 @@
+// Copyright The HTNN Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package plugin_state
+
+import (
+	"mosn.io/htnn/pkg/filtermanager/api"
+)
+
+type pluginState struct {
+	store map[string]map[string]any
+}
+
+func NewPluginState() api.PluginState {
+	return &pluginState{
+		store: make(map[string]map[string]any),
+	}
+}
+
+func (p *pluginState) Get(namespace string, key string) any {
+	if pluginStore, ok := p.store[namespace]; ok {
+		return pluginStore[key]
+	}
+	return nil
+}
+
+func (p *pluginState) Set(namespace string, key string, value any) {
+	pluginStore, ok := p.store[namespace]
+	if !ok {
+		pluginStore = make(map[string]any)
+		p.store[namespace] = pluginStore
+	}
+	pluginStore[key] = value
+}

internal/plugin_state/plugin_state_test.go

@@ -0,0 +1,33 @@
+// Copyright The HTNN Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package plugin_state
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPluginStateSetGet(t *testing.T) {
+	p := NewPluginState()
+	assert.Nil(t, p.Get("unknown", "key"))
+
+	p.Set("plugin", "key", "value")
+	assert.Nil(t, p.Get("plugin", "unknown"))
+	assert.Equal(t, "value", p.Get("plugin", "key"))
+
+	p.Set("plugin", "key2", "value")
+	assert.Equal(t, "value", p.Get("plugin", "key2"))
+}

pkg/filtermanager/api/api.go

@@ -202,13 +202,18 @@ type FilterCallbackHandler interface {
 	// * ErrValueNotFound
 	GetProperty(key string) (string, error)
 
+	// Methods added by HTNN
+
 	// LookupConsumer is used in the Authn plugins to fetch the corresponding consumer, with
 	// the plugin name and plugin specific key. We return a 'fat' Consumer so that additional
 	// info like `Name` can be retrieved.
 	LookupConsumer(pluginName, key string) (Consumer, bool)
 	// SetConsumer is used in the Authn plugins to set the corresponding consumer after authentication.
 	SetConsumer(c Consumer)
 	GetConsumer() Consumer
+
+	// PluginState returns the PluginState associated to this request.
+	PluginState() PluginState
 }
 
 type FilterFactory func(config interface{}, callbacks FilterCallbackHandler) Filter
@@ -219,6 +224,18 @@ type DynamicMetadata = api.DynamicMetadata
 // FilterState operates the Envoy's filter state
 type FilterState = api.FilterState
 
+// PluginState stores the plugin level state shared between Go plugins. Unlike DynamicMetadata,
+// it doesn't do any serialization/deserialization. So:
+// 1. modifying the returned state can affect the internal structure.
+// 2. fields can't be marshalled can be kept in the state.
+// 3. one can't access the state outside the current Envoy Go filter.
+type PluginState interface {
+	// Get the value. Returns nil if the value doesn't exist.
+	Get(namespace string, key string) any
+	// Set the value.
+	Set(namespace string, key string, value any)
+}
+
 // ConfigCallbackHandler provides API that is used during initializing configuration
 type ConfigCallbackHandler interface {
 	// The ConfigCallbackHandler from Envoy is only available when the plugin is

pkg/filtermanager/filtermanager.go

@@ -33,6 +33,7 @@ import (
 
 	internalConsumer "mosn.io/htnn/internal/consumer"
 	"mosn.io/htnn/internal/cookie"
+	"mosn.io/htnn/internal/plugin_state"
 	"mosn.io/htnn/pkg/filtermanager/api"
 	"mosn.io/htnn/pkg/filtermanager/model"
 	pkgPlugins "mosn.io/htnn/pkg/plugins"
@@ -304,8 +305,9 @@ func (s *filterManagerStreamInfo) DownstreamRemoteAddress() string {
 type filterManagerCallbackHandler struct {
 	capi.FilterCallbackHandler
 
-	namespace string
-	consumer  api.Consumer
+	namespace   string
+	consumer    api.Consumer
+	pluginState api.PluginState
 
 	streamInfo *filterManagerStreamInfo
 }
@@ -344,6 +346,13 @@ func (cb *filterManagerCallbackHandler) SetConsumer(c api.Consumer) {
 	cb.consumer = c
 }
 
+func (cb *filterManagerCallbackHandler) PluginState() api.PluginState {
+	if cb.pluginState == nil {
+		cb.pluginState = plugin_state.NewPluginState()
+	}
+	return cb.pluginState
+}
+
 type phase int
 
 const (

pkg/filtermanager/filtermanager_test.go

@@ -430,3 +430,58 @@ func TestFiltersFromConsumer(t *testing.T) {
 		assert.True(t, ok)
 	}
 }
+
+func setPluginStateFilterFactory(c interface{}, callbacks api.FilterCallbackHandler) api.Filter {
+	return &setPluginStateFilter{
+		callbacks: callbacks,
+	}
+}
+
+type setPluginStateFilter struct {
+	api.PassThroughFilter
+	callbacks api.FilterCallbackHandler
+}
+
+func (f *setPluginStateFilter) DecodeHeaders(headers api.RequestHeaderMap, endStream bool) api.ResultAction {
+	f.callbacks.PluginState().Set("test", "key", "value")
+	return api.Continue
+}
+
+func getPluginStateFilterFactory(c interface{}, callbacks api.FilterCallbackHandler) api.Filter {
+	return &getPluginStateFilter{
+		callbacks: callbacks,
+	}
+}
+
+type getPluginStateFilter struct {
+	api.PassThroughFilter
+	callbacks api.FilterCallbackHandler
+}
+
+func (f *getPluginStateFilter) DecodeHeaders(headers api.RequestHeaderMap, endStream bool) api.ResultAction {
+	v := f.callbacks.PluginState().Get("test", "key")
+	headers.Set("x-htnn-v", v.(string))
+	return api.Continue
+}
+
+func TestPluginState(t *testing.T) {
+	cb := envoy.NewCAPIFilterCallbackHandler()
+	config := initFilterManagerConfig("ns")
+	config.parsed = []*model.ParsedFilterConfig{
+		{
+			Name:    "alice",
+			Factory: setPluginStateFilterFactory,
+		},
+		{
+			Name:    "bob",
+			Factory: getPluginStateFilterFactory,
+		},
+	}
+	m := FilterManagerFactory(config, cb).(*filterManager)
+	h := http.Header{}
+	hdr := envoy.NewRequestHeaderMap(h)
+	m.DecodeHeaders(hdr, true)
+	cb.WaitContinued()
+	v, _ := hdr.Get("x-htnn-v")
+	assert.Equal(t, "value", v)
+}

plugins/tests/pkg/envoy/capi.go

@@ -27,6 +27,7 @@ import (
 	capi "github.com/envoyproxy/envoy/contrib/golang/common/go/api"
 
 	"mosn.io/htnn/internal/cookie"
+	"mosn.io/htnn/internal/plugin_state"
 	"mosn.io/htnn/pkg/filtermanager/api"
 )
 
@@ -421,10 +422,11 @@ type filterCallbackHandler struct {
 	// add lock to the test helper to satisfy -race check
 	lock *sync.RWMutex
 
-	streamInfo api.StreamInfo
-	resp       LocalResponse
-	consumer   api.Consumer
-	ch         chan struct{}
+	streamInfo  api.StreamInfo
+	resp        LocalResponse
+	consumer    api.Consumer
+	pluginState api.PluginState
+	ch          chan struct{}
 }
 
 func NewFilterCallbackHandler() *filterCallbackHandler {
@@ -498,6 +500,13 @@ func (i *filterCallbackHandler) SetConsumer(c api.Consumer) {
 	i.consumer = c
 }
 
+func (i *filterCallbackHandler) PluginState() api.PluginState {
+	if i.pluginState == nil {
+		i.pluginState = plugin_state.NewPluginState()
+	}
+	return i.pluginState
+}
+
 var _ api.FilterCallbackHandler = (*filterCallbackHandler)(nil)
 
 type capiFilterCallbackHandler struct {