|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + |
| 6 | + "github.com/proxy-wasm/proxy-wasm-go-sdk/proxywasm" |
| 7 | + "github.com/proxy-wasm/proxy-wasm-go-sdk/proxywasm/types" |
| 8 | + "github.com/tidwall/gjson" |
| 9 | +) |
| 10 | + |
| 11 | +func main() {} |
| 12 | +func init() { |
| 13 | + // SetVMContext is the entrypoint for setting up this entire Wasm VM. |
| 14 | + // Please make sure that this entrypoint be called during "main()" function, otherwise |
| 15 | + // this VM would fail. |
| 16 | + proxywasm.SetVMContext(&vmContext{}) |
| 17 | +} |
| 18 | + |
| 19 | +// vmContext implements types.VMContext. |
| 20 | +type vmContext struct { |
| 21 | + // Embed the default VM context here, |
| 22 | + // so that we don't need to reimplement all the methods. |
| 23 | + types.DefaultVMContext |
| 24 | +} |
| 25 | + |
| 26 | +// NewPluginContext implements types.VMContext. |
| 27 | +func (*vmContext) NewPluginContext(contextID uint32) types.PluginContext { |
| 28 | + return &pluginContext{} |
| 29 | +} |
| 30 | + |
| 31 | +// pluginContext implements types.PluginContext. |
| 32 | +type pluginContext struct { |
| 33 | + // Embed the default plugin context here, |
| 34 | + // so that we don't need to reimplement all the methods. |
| 35 | + types.DefaultPluginContext |
| 36 | + configuration pluginConfiguration |
| 37 | +} |
| 38 | + |
| 39 | +// pluginConfiguration is a type to represent an example configuration for this wasm plugin. |
| 40 | +type pluginConfiguration struct { |
| 41 | + // Example configuration field. |
| 42 | + // The plugin will validate if those fields exist in the json payload. |
| 43 | + requiredKeys []string |
| 44 | +} |
| 45 | + |
| 46 | +// OnPluginStart implements types.PluginContext. |
| 47 | +func (ctx *pluginContext) OnPluginStart(pluginConfigurationSize int) types.OnPluginStartStatus { |
| 48 | + data, err := proxywasm.GetPluginConfiguration() |
| 49 | + if err != nil && err != types.ErrorStatusNotFound { |
| 50 | + proxywasm.LogCriticalf("error reading plugin configuration: %v", err) |
| 51 | + return types.OnPluginStartStatusFailed |
| 52 | + } |
| 53 | + config, err := parsePluginConfiguration(data) |
| 54 | + if err != nil { |
| 55 | + proxywasm.LogCriticalf("error parsing plugin configuration: %v", err) |
| 56 | + return types.OnPluginStartStatusFailed |
| 57 | + } |
| 58 | + ctx.configuration = config |
| 59 | + return types.OnPluginStartStatusOK |
| 60 | +} |
| 61 | + |
| 62 | +// parsePluginConfiguration parses the json plugin configuration data and returns pluginConfiguration. |
| 63 | +// Note that this parses the json data by gjson, since TinyGo doesn't support encoding/json. |
| 64 | +// You can also try https://github.com/mailru/easyjson, which supports decoding to a struct. |
| 65 | +func parsePluginConfiguration(data []byte) (pluginConfiguration, error) { |
| 66 | + if len(data) == 0 { |
| 67 | + return pluginConfiguration{}, nil |
| 68 | + } |
| 69 | + |
| 70 | + config := &pluginConfiguration{} |
| 71 | + if !gjson.ValidBytes(data) { |
| 72 | + return pluginConfiguration{}, fmt.Errorf("the plugin configuration is not a valid json: %q", string(data)) |
| 73 | + } |
| 74 | + |
| 75 | + jsonData := gjson.ParseBytes(data) |
| 76 | + requiredKeys := jsonData.Get("requiredKeys").Array() |
| 77 | + for _, requiredKey := range requiredKeys { |
| 78 | + config.requiredKeys = append(config.requiredKeys, requiredKey.Str) |
| 79 | + } |
| 80 | + |
| 81 | + return *config, nil |
| 82 | +} |
| 83 | + |
| 84 | +// NewHttpContext implements types.PluginContext. |
| 85 | +func (ctx *pluginContext) NewHttpContext(contextID uint32) types.HttpContext { |
| 86 | + return &payloadValidationContext{requiredKeys: ctx.configuration.requiredKeys} |
| 87 | +} |
| 88 | + |
| 89 | +// payloadValidationContext implements types.HttpContext. |
| 90 | +type payloadValidationContext struct { |
| 91 | + // Embed the default root http context here, |
| 92 | + // so that we don't need to reimplement all the methods. |
| 93 | + types.DefaultHttpContext |
| 94 | + totalRequestBodySize int |
| 95 | + requiredKeys []string |
| 96 | +} |
| 97 | + |
| 98 | +// OnHttpRequestHeaders implements types.HttpContext. |
| 99 | +func (*payloadValidationContext) OnHttpRequestHeaders(numHeaders int, _ bool) types.Action { |
| 100 | + contentType, err := proxywasm.GetHttpRequestHeader("content-type") |
| 101 | + if err != nil || contentType != "application/json" { |
| 102 | + // If the header doesn't have the expected content value, send the 403 response, |
| 103 | + if err := proxywasm.SendHttpResponse(403, nil, []byte("content-type must be provided"), -1); err != nil { |
| 104 | + panic(err) |
| 105 | + } |
| 106 | + // and terminates the further processing of this traffic by ActionPause. |
| 107 | + return types.ActionPause |
| 108 | + } |
| 109 | + |
| 110 | + // ActionContinue lets the host continue the processing the body. |
| 111 | + return types.ActionContinue |
| 112 | +} |
| 113 | + |
| 114 | +// OnHttpRequestBody implements types.HttpContext. |
| 115 | +func (ctx *payloadValidationContext) OnHttpRequestBody(bodySize int, endOfStream bool) types.Action { |
| 116 | + ctx.totalRequestBodySize += bodySize |
| 117 | + if !endOfStream { |
| 118 | + // OnHttpRequestBody may be called each time a part of the body is received. |
| 119 | + // Wait until we see the entire body to replace. |
| 120 | + return types.ActionPause |
| 121 | + } |
| 122 | + |
| 123 | + body, err := proxywasm.GetHttpRequestBody(0, ctx.totalRequestBodySize) |
| 124 | + if err != nil { |
| 125 | + proxywasm.LogErrorf("failed to get request body: %v", err) |
| 126 | + return types.ActionContinue |
| 127 | + } |
| 128 | + if !ctx.validatePayload(body) { |
| 129 | + // If the validation fails, send the 403 response, |
| 130 | + if err := proxywasm.SendHttpResponse(403, nil, []byte("invalid payload"), -1); err != nil { |
| 131 | + proxywasm.LogErrorf("failed to send the 403 response: %v", err) |
| 132 | + } |
| 133 | + // and terminates this traffic. |
| 134 | + return types.ActionPause |
| 135 | + } |
| 136 | + |
| 137 | + return types.ActionContinue |
| 138 | +} |
| 139 | + |
| 140 | +// validatePayload validates the given json payload. |
| 141 | +// Note that this function parses the json data by gjson, since TinyGo doesn't support encoding/json. |
| 142 | +func (ctx *payloadValidationContext) validatePayload(body []byte) bool { |
| 143 | + if !gjson.ValidBytes(body) { |
| 144 | + proxywasm.LogErrorf("body is not a valid json: %q", string(body)) |
| 145 | + return false |
| 146 | + } |
| 147 | + jsonData := gjson.ParseBytes(body) |
| 148 | + |
| 149 | + // Do any validation on the json. Check if required keys exist here as an example. |
| 150 | + // The required keys are configurable via the plugin configuration. |
| 151 | + for _, requiredKey := range ctx.requiredKeys { |
| 152 | + if !jsonData.Get(requiredKey).Exists() { |
| 153 | + proxywasm.LogErrorf("required key (%v) is missing: %v", requiredKey, jsonData) |
| 154 | + return false |
| 155 | + } |
| 156 | + } |
| 157 | + |
| 158 | + return true |
| 159 | +} |
0 commit comments