jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-05-25 13:17:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

support mcp security guard (#3295)

Browse Source
This commit is contained in:
rinfx
2026-01-29 19:25:43 +08:00
committed by GitHub
parent e824653378
commit 8eaa385a56
4 changed files with 415 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/handler.go
View File

@@ -4,6 +4,7 @@ import (
cfg "github.com/alibaba/higress/plugins/wasm-go/extensions/ai-security-guard/config"
common_text "github.com/alibaba/higress/plugins/wasm-go/extensions/ai-security-guard/lvwang/common/text"
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/image"
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/mcp"
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/text"
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
"github.com/higress-group/wasm-go/pkg/log"
@@ -28,6 +29,8 @@ func OnHttpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig, bod
log.Errorf("[on request body] image generation api don't support provider: %s", config.ProviderType)
return types.ActionContinue
}
case cfg.ApiMCP:
return mcp.HandleMcpRequestBody(ctx, config, body)
default:
log.Errorf("[on request body] multi_modal_guard don't support api: %s", config.ApiType)
return types.ActionContinue
@@ -46,6 +49,15 @@ func OnHttpResponseHeaders(ctx wrapper.HttpContext, config cfg.AISecurityConfig)
log.Errorf("[on response header] image generation api don't support provider: %s", config.ProviderType)
return types.ActionContinue
}
case cfg.ApiMCP:
if wrapper.IsApplicationJson() {
ctx.BufferResponseBody()
return types.HeaderStopIteration
} else {
ctx.SetContext("during_call", false)
ctx.NeedPauseStreamingResponse()
return types.ActionContinue
}
default:
log.Errorf("[on response header] multi_modal_guard don't support api: %s", config.ApiType)
return types.ActionContinue
@@ -56,6 +68,8 @@ func OnHttpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityC
switch config.ApiType {
case cfg.ApiTextGeneration:
return common_text.HandleTextGenerationStreamingResponseBody(ctx, config, data, endOfStream)
case cfg.ApiMCP:
return mcp.HandleMcpStreamingResponseBody(ctx, config, data, endOfStream)
default:
log.Errorf("[on streaming response body] multi_modal_guard don't support api: %s", config.ApiType)
return data
@@ -76,6 +90,8 @@ func OnHttpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig, bo
log.Errorf("[on response body] image generation api don't support provider: %s", config.ProviderType)
return types.ActionContinue
}
case cfg.ApiMCP:
return mcp.HandleMcpResponseBody(ctx, config, body)
default:
log.Errorf("[on response body] multi_modal_guard don't support api: %s", config.ApiType)
return types.ActionContinue