feat: implement custom-response plugin in the golang version (#689)

plugins/wasm-go/extensions/custom-response/main.go

@@ -0,0 +1,135 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// 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 main
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
+	"github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm"
+	"github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm/types"
+	"github.com/tidwall/gjson"
+)
+
+func main() {
+	wrapper.SetCtx(
+		"custom-response",
+		wrapper.ParseConfigBy(parseConfig),
+		wrapper.ProcessRequestHeadersBy(onHttpRequestHeaders),
+		wrapper.ProcessResponseHeadersBy(onHttpResponseHeaders),
+	)
+}
+
+type CustomResponseConfig struct {
+	statusCode     uint32
+	headers        [][2]string
+	body           string
+	enableOnStatus []uint32
+	contentType    string
+}
+
+func parseConfig(gjson gjson.Result, config *CustomResponseConfig, log wrapper.Log) error {
+	headersArray := gjson.Get("headers").Array()
+	config.headers = make([][2]string, 0, len(headersArray))
+	for _, v := range headersArray {
+		kv := strings.SplitN(v.String(), "=", 2)
+		if len(kv) == 2 {
+			key := strings.TrimSpace(kv[0])
+			value := strings.TrimSpace(kv[1])
+			if strings.EqualFold(key, "content-type") {
+				config.contentType = value
+			} else if strings.EqualFold(key, "content-length") {
+				continue
+			} else {
+				config.headers = append(config.headers, [2]string{key, value})
+			}
+		} else {
+			return fmt.Errorf("invalid header pair format: %s", v.String())
+		}
+	}
+
+	config.body = gjson.Get("body").String()
+	if config.contentType == "" && config.body != "" {
+		if json.Valid([]byte(config.body)) {
+			config.contentType = "application/json; charset=utf-8"
+		} else {
+			config.contentType = "text/plain; charset=utf-8"
+		}
+	}
+	config.headers = append(config.headers, [2]string{"content-type", config.contentType})
+
+	config.statusCode = 200
+	if gjson.Get("status_code").Exists() {
+		statusCode := gjson.Get("status_code")
+		parsedStatusCode, err := strconv.Atoi(statusCode.String())
+		if err != nil {
+			return fmt.Errorf("invalid status code value: %s", statusCode.String())
+		}
+		config.statusCode = uint32(parsedStatusCode)
+	}
+
+	enableOnStatusArray := gjson.Get("enable_on_status").Array()
+	config.enableOnStatus = make([]uint32, 0, len(enableOnStatusArray))
+	for _, v := range enableOnStatusArray {
+		parsedEnableOnStatus, err := strconv.Atoi(v.String())
+		if err != nil {
+			return fmt.Errorf("invalid enable_on_status value: %s", v.String())
+		}
+		config.enableOnStatus = append(config.enableOnStatus, uint32(parsedEnableOnStatus))
+	}
+
+	return nil
+}
+
+func onHttpRequestHeaders(ctx wrapper.HttpContext, config CustomResponseConfig, log wrapper.Log) types.Action {
+	if len(config.enableOnStatus) != 0 {
+		return types.ActionContinue
+	}
+	err := proxywasm.SendHttpResponse(config.statusCode, config.headers, []byte(config.body), -1)
+	if err != nil {
+		log.Errorf("send http response failed: %v", err)
+	}
+
+	return types.ActionPause
+}
+
+func onHttpResponseHeaders(ctx wrapper.HttpContext, config CustomResponseConfig, log wrapper.Log) types.Action {
+	// enableOnStatus is not empty, compare the status code.
+	// if match the status code, mock the response.
+	statusCodeStr, err := proxywasm.GetHttpResponseHeader(":status")
+	if err != nil {
+		log.Errorf("get http response status code failed: %v", err)
+		return types.ActionContinue
+	}
+	statusCode, err := strconv.ParseUint(statusCodeStr, 10, 32)
+	if err != nil {
+		log.Errorf("parse http response status code failed: %v", err)
+		return types.ActionContinue
+	}
+
+	for _, v := range config.enableOnStatus {
+		if uint32(statusCode) == v {
+			err = proxywasm.SendHttpResponse(config.statusCode, config.headers, []byte(config.body), -1)
+			if err != nil {
+				log.Errorf("send http response failed: %v", err)
+			}
+		}
+	}
+
+	return types.ActionContinue
+}