extension mechanism for custom logs and span attributes (#1451)

plugins/wasm-go/pkg/wrapper/plugin_wrapper.go

@@ -15,6 +15,8 @@
 package wrapper
 
 import (
+	"encoding/json"
+	"fmt"
 	"strconv"
 	"time"
 	"unsafe"
@@ -26,6 +28,12 @@ import (
 	"github.com/alibaba/higress/plugins/wasm-go/pkg/matcher"
 )
 
+const (
+	CustomLogKey       = "custom_log"
+	AILogKey           = "ai_log"
+	TraceSpanTagPrefix = "trace_span_tag."
+)
+
 type HttpContext interface {
 	Scheme() string
 	Host() string
@@ -35,6 +43,14 @@ type HttpContext interface {
 	GetContext(key string) interface{}
 	GetBoolContext(key string, defaultValue bool) bool
 	GetStringContext(key, defaultValue string) string
+	GetUserAttribute(key string) interface{}
+	SetUserAttribute(key string, value interface{})
+	// You can call this function to set custom log
+	WriteUserAttributeToLog() error
+	// You can call this function to set custom log with your specific key
+	WriteUserAttributeToLogWithKey(key string) error
+	// You can call this function to set custom trace span attribute
+	WriteUserAttributeToTrace() error
 	// If the onHttpRequestBody handle is not set, the request body will not be read by default
 	DontReadRequestBody()
 	// If the onHttpResponseBody handle is not set, the request body will not be read by default
@@ -335,9 +351,10 @@ func (ctx *CommonPluginCtx[PluginConfig]) OnTick() {
 
 func (ctx *CommonPluginCtx[PluginConfig]) NewHttpContext(contextID uint32) types.HttpContext {
 	httpCtx := &CommonHttpCtx[PluginConfig]{
-		plugin:      ctx,
-		contextID:   contextID,
-		userContext: map[string]interface{}{},
+		plugin:        ctx,
+		contextID:     contextID,
+		userContext:   map[string]interface{}{},
+		userAttribute: map[string]interface{}{},
 	}
 	if ctx.vm.onHttpRequestBody != nil || ctx.vm.onHttpStreamingRequestBody != nil {
 		httpCtx.needRequestBody = true
@@ -367,6 +384,7 @@ type CommonHttpCtx[PluginConfig any] struct {
 	responseBodySize      int
 	contextID             uint32
 	userContext           map[string]interface{}
+	userAttribute         map[string]interface{}
 }
 
 func (ctx *CommonHttpCtx[PluginConfig]) SetContext(key string, value interface{}) {
@@ -377,6 +395,63 @@ func (ctx *CommonHttpCtx[PluginConfig]) GetContext(key string) interface{} {
 	return ctx.userContext[key]
 }
 
+func (ctx *CommonHttpCtx[PluginConfig]) SetUserAttribute(key string, value interface{}) {
+	ctx.userAttribute[key] = value
+}
+
+func (ctx *CommonHttpCtx[PluginConfig]) GetUserAttribute(key string) interface{} {
+	return ctx.userAttribute[key]
+}
+
+func (ctx *CommonHttpCtx[PluginConfig]) WriteUserAttributeToLog() error {
+	return ctx.WriteUserAttributeToLogWithKey(CustomLogKey)
+}
+
+func (ctx *CommonHttpCtx[PluginConfig]) WriteUserAttributeToLogWithKey(key string) error {
+	// e.g. {\"field1\":\"value1\",\"field2\":\"value2\"}
+	preMarshalledJsonLogStr, _ := proxywasm.GetProperty([]string{key})
+	newAttributeMap := map[string]interface{}{}
+	if string(preMarshalledJsonLogStr) != "" {
+		// e.g. {"field1":"value1","field2":"value2"}
+		preJsonLogStr := unmarshalStr(fmt.Sprintf(`"%s"`, string(preMarshalledJsonLogStr)))
+		err := json.Unmarshal([]byte(preJsonLogStr), &newAttributeMap)
+		if err != nil {
+			ctx.plugin.vm.log.Warnf("Unmarshal failed, will overwrite %s, pre value is: %s", key, string(preMarshalledJsonLogStr))
+			return err
+		}
+	}
+	// update customLog
+	for k, v := range ctx.userAttribute {
+		newAttributeMap[k] = v
+	}
+	// e.g. {"field1":"value1","field2":2,"field3":"value3"}
+	jsonStr, _ := json.Marshal(newAttributeMap)
+	// e.g. {\"field1\":\"value1\",\"field2\":2,\"field3\":\"value3\"}
+	marshalledJsonStr := marshalStr(string(jsonStr))
+	if err := proxywasm.SetProperty([]string{key}, []byte(marshalledJsonStr)); err != nil {
+		ctx.plugin.vm.log.Warnf("failed to set %s in filter state, raw is %s, err is %v", key, marshalledJsonStr, err)
+		return err
+	}
+	return nil
+}
+
+func (ctx *CommonHttpCtx[PluginConfig]) WriteUserAttributeToTrace() error {
+	for k, v := range ctx.userAttribute {
+		traceSpanTag := TraceSpanTagPrefix + k
+		traceSpanValue := fmt.Sprint(v)
+		var err error
+		if traceSpanValue != "" {
+			err = proxywasm.SetProperty([]string{traceSpanTag}, []byte(traceSpanValue))
+		} else {
+			err = fmt.Errorf("value of %s is empty", traceSpanTag)
+		}
+		if err != nil {
+			ctx.plugin.vm.log.Warnf("Failed to set trace attribute - %s: %s, error message: %v", traceSpanTag, traceSpanValue, err)
+		}
+	}
+	return nil
+}
+
 func (ctx *CommonHttpCtx[PluginConfig]) GetBoolContext(key string, defaultValue bool) bool {
 	if b, ok := ctx.userContext[key].(bool); ok {
 		return b

plugins/wasm-go/pkg/wrapper/utils.go

@@ -0,0 +1,36 @@
+package wrapper
+
+import (
+	"encoding/json"
+
+	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
+	"github.com/tidwall/gjson"
+)
+
+func unmarshalStr(marshalledJsonStr string) string {
+	// e.g. "{\"field1\":\"value1\",\"field2\":\"value2\"}"
+	var jsonStr string
+	err := json.Unmarshal([]byte(marshalledJsonStr), &jsonStr)
+	if err != nil {
+		proxywasm.LogErrorf("failed to unmarshal json string, raw string is: %s, err is: %v", marshalledJsonStr, err)
+		return ""
+	}
+	// e.g. {"field1":"value1","field2":"value2"}
+	return jsonStr
+}
+
+func marshalStr(raw string) string {
+	// e.g. {"field1":"value1","field2":"value2"}
+	helper := map[string]string{
+		"placeholder": raw,
+	}
+	marshalledHelper, _ := json.Marshal(helper)
+	marshalledRaw := gjson.GetBytes(marshalledHelper, "placeholder").Raw
+	if len(marshalledRaw) >= 2 {
+		// e.g. {\"field1\":\"value1\",\"field2\":\"value2\"}
+		return marshalledRaw[1 : len(marshalledRaw)-1]
+	} else {
+		proxywasm.LogErrorf("failed to marshal json string, raw string is: %s", raw)
+		return ""
+	}
+}