Support set buffer limit (#1153)

This commit is contained in:
澄潭
2024-07-25 20:42:39 +08:00
committed by GitHub
parent 95edce024d
commit 29079f4e2a
4 changed files with 75 additions and 0 deletions

View File

@@ -0,0 +1,13 @@
diff --git a/source/common/http/headers.h b/source/common/http/headers.h
index a7a8a3393e..6af4a2852d 100644
--- a/source/common/http/headers.h
+++ b/source/common/http/headers.h
@@ -123,7 +123,7 @@ public:
const LowerCaseString TriCostTime{"req-cost-time"};
const LowerCaseString TriStartTime{"req-start-time"};
const LowerCaseString TriRespStartTime{"resp-start-time"};
- const LowerCaseString EnvoyOriginalHost{"original-host"};
+ const LowerCaseString EnvoyOriginalHost{"x-envoy-original-host"};
const LowerCaseString HigressOriginalService{"x-higress-original-service"};
} AliExtendedValues;
#endif

View File

@@ -0,0 +1,43 @@
diff --git a/source/extensions/common/wasm/context.cc b/source/extensions/common/wasm/context.cc
index 9642d8abd3..410baa856f 100644
--- a/source/extensions/common/wasm/context.cc
+++ b/source/extensions/common/wasm/context.cc
@@ -62,6 +62,21 @@ constexpr absl::string_view CelStateKeyPrefix = "wasm.";
#if defined(ALIMESH)
constexpr std::string_view ClearRouteCacheKey = "clear_route_cache";
constexpr std::string_view DisableClearRouteCache = "off";
+constexpr std::string_view SetDecoderBufferLimit = "set_decoder_buffer_limit";
+constexpr std::string_view SetEncoderBufferLimit = "set_encoder_buffer_limit";
+
+bool stringViewToUint32(std::string_view str, uint32_t& out_value) {
+ try {
+ unsigned long temp = std::stoul(std::string(str));
+ if (temp <= std::numeric_limits<uint32_t>::max()) {
+ out_value = static_cast<uint32_t>(temp);
+ return true;
+ }
+ } catch (const std::exception& e) {
+ ENVOY_LOG_MISC(critical, "stringToUint exception '{}'", e.what());
+ }
+ return false;
+}
#endif
using HashPolicy = envoy::config::route::v3::RouteAction::HashPolicy;
@@ -1280,6 +1295,16 @@ WasmResult Context::setProperty(std::string_view path, std::string_view value) {
} else {
disable_clear_route_cache_ = false;
}
+ } else if (path == SetDecoderBufferLimit && decoder_callbacks_) {
+ uint32_t buffer_limit;
+ if (stringViewToUint32(value, buffer_limit)) {
+ decoder_callbacks_->setDecoderBufferLimit(buffer_limit);
+ }
+ } else if (path == SetEncoderBufferLimit && encoder_callbacks_) {
+ uint32_t buffer_limit;
+ if (stringViewToUint32(value, buffer_limit)) {
+ encoder_callbacks_->setEncoderBufferLimit(buffer_limit);
+ }
}
#endif
if (!state->setValue(toAbslStringView(value))) {

View File

@@ -21,6 +21,8 @@ const (
pluginName = "ai-proxy"
ctxKeyApiName = "apiKey"
defaultMaxBodyBytes uint32 = 10 * 1024 * 1024
)
func main() {
@@ -76,6 +78,7 @@ func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConf
action, err := handler.OnRequestHeaders(ctx, apiName, log)
if err == nil {
if contentType, err := proxywasm.GetHttpRequestHeader("Content-Type"); err == nil && contentType != "" {
ctx.SetRequestBodyBufferLimit(defaultMaxBodyBytes)
// Always return types.HeaderStopIteration to support fallback routing,
// as long as onHttpRequestBody can be called.
return types.HeaderStopIteration

View File

@@ -15,6 +15,7 @@
package wrapper
import (
"strconv"
"time"
"unsafe"
@@ -49,7 +50,12 @@ type HttpContext interface {
// If the onHttpStreamingResponseBody handle is not set, and the onHttpResponseBody handle is set, the response body will be buffered by default
BufferResponseBody()
// If any request header is changed in onHttpRequestHeaders, envoy will re-calculate the route. Call this function to disable the re-routing.
// You need to call this before making any header modification operations.
DisableReroute()
// Note that this parameter affects the gateway's memory usageSupport setting a maximum buffer size for each request body individually in request phase.
SetRequestBodyBufferLimit(size uint32)
// Note that this parameter affects the gateway's memory usage! Support setting a maximum buffer size for each response body individually in response phase.
SetResponseBodyBufferLimit(size uint32)
}
type ParseConfigFunc[PluginConfig any] func(json gjson.Result, config *PluginConfig, log Log) error
@@ -353,6 +359,16 @@ func (ctx *CommonHttpCtx[PluginConfig]) DisableReroute() {
_ = proxywasm.SetProperty([]string{"clear_route_cache"}, []byte("off"))
}
func (ctx *CommonHttpCtx[PluginConfig]) SetRequestBodyBufferLimit(size uint32) {
ctx.plugin.vm.log.Infof("SetRequestBodyBufferLimit: %d", size)
_ = proxywasm.SetProperty([]string{"set_decoder_buffer_limit"}, []byte(strconv.Itoa(int(size))))
}
func (ctx *CommonHttpCtx[PluginConfig]) SetResponseBodyBufferLimit(size uint32) {
ctx.plugin.vm.log.Infof("SetResponseBodyBufferLimit: %d", size)
_ = proxywasm.SetProperty([]string{"set_encoder_buffer_limit"}, []byte(strconv.Itoa(int(size))))
}
func (ctx *CommonHttpCtx[PluginConfig]) OnHttpRequestHeaders(numHeaders int, endOfStream bool) types.Action {
config, err := ctx.plugin.GetMatchConfig()
if err != nil {