feat(ai-security-guard): enhance risk action resolution and support sensitive data masking (#3690)

Co-authored-by: rinfx <yucheng.lxr@alibaba-inc.com>
2026-05-28 06:37:26 +08:00 · 2026-04-15 11:14:56 +08:00
parent e2beb6cd45
commit b1187cc14d
13 changed files with 5019 additions and 214 deletions
--- a/plugins/wasm-go/extensions/ai-security-guard/config/action_resolver_test.go
+++ b/plugins/wasm-go/extensions/ai-security-guard/config/action_resolver_test.go
@@ -0,0 +1,365 @@
+// Copyright (c) 2024 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 config
+
+import (
+	"os"
+	"testing"
+
+	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/proxytest"
+	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
+	"github.com/stretchr/testify/require"
+)
+
+// testVMContext is a minimal VMContext for setting up the proxy-wasm mock host.
+type testVMContext struct {
+	types.DefaultVMContext
+}
+
+// TestMain sets up the proxy-wasm mock host for all tests in the config package.
+// This is required because functions like enforceMaskBoundary call proxywasm.LogWarnf.
+func TestMain(m *testing.M) {
+	opt := proxytest.NewEmulatorOption().WithVMContext(&testVMContext{})
+	_, reset := proxytest.NewHostEmulator(opt)
+	defer reset()
+	os.Exit(m.Run())
+}
+
+// =============================================================================
+// TC-RESOLVE: 动作解析优先级测试（ResolveRiskActionByType）
+// =============================================================================
+
+// TestTC_RESOLVE_001 仅全局 riskAction=mask，无维度动作
+// => sensitiveData 返回 mask，非 sensitiveData 维度降级为 block
+func TestTC_RESOLVE_001(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "mask",
+	}
+
+	// sensitiveData 返回 mask（source=global_global）
+	action, source := config.ResolveRiskActionByType("", SensitiveDataType)
+	require.Equal(t, "mask", action)
+	require.Equal(t, "global_global", source)
+
+	// promptAttack 降级为 block（source=global_global）
+	action, source = config.ResolveRiskActionByType("", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_global", source)
+
+	// contentModeration 降级为 block
+	action, source = config.ResolveRiskActionByType("", ContentModerationType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_global", source)
+
+	// maliciousUrl 降级为 block
+	action, source = config.ResolveRiskActionByType("", MaliciousUrlDataType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_global", source)
+}
+
+// TestTC_RESOLVE_002 全局 riskAction=mask + 全局 promptAttackAction=block
+// => promptAttack 返回 block，sensitiveData 返回 mask
+func TestTC_RESOLVE_002(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction:         "mask",
+		PromptAttackAction: "block",
+	}
+
+	// promptAttack 返回 block（source=global_dimension）
+	action, source := config.ResolveRiskActionByType("", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_dimension", source)
+
+	// sensitiveData 返回 mask（source=global_global）
+	action, source = config.ResolveRiskActionByType("", SensitiveDataType)
+	require.Equal(t, "mask", action)
+	require.Equal(t, "global_global", source)
+}
+
+// TestTC_RESOLVE_003 consumer 规则含 riskAction=block，全局 sensitiveDataAction=mask
+// => sensitiveData 返回 block（consumer_global 优先于 global_dimension）
+func TestTC_RESOLVE_003(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction:          "mask",
+		SensitiveDataAction: "mask",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":    Matcher{Exact: "user-a"},
+				"riskAction": "block",
+			},
+		},
+	}
+
+	// consumer_global(block) 优先于 global_dimension(mask)
+	action, source := config.ResolveRiskActionByType("user-a", SensitiveDataType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_global", source)
+
+	// 未命中 consumer 规则时，回退到 global_dimension
+	action, source = config.ResolveRiskActionByType("user-b", SensitiveDataType)
+	require.Equal(t, "mask", action)
+	require.Equal(t, "global_dimension", source)
+}
+
+// TestTC_RESOLVE_004 consumer 规则含 sensitiveDataAction=mask 且 riskAction=block
+// => sensitiveData 返回 mask（consumer_dimension 优先）
+func TestTC_RESOLVE_004(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":             Matcher{Exact: "user-a"},
+				"riskAction":          "block",
+				"sensitiveDataAction": "mask",
+			},
+		},
+	}
+
+	// consumer_dimension(mask) 优先于 consumer_global(block)
+	action, source := config.ResolveRiskActionByType("user-a", SensitiveDataType)
+	require.Equal(t, "mask", action)
+	require.Equal(t, "consumer_dimension", source)
+
+	// promptAttack 无 consumer_dimension，回退到 consumer_global(block)
+	action, source = config.ResolveRiskActionByType("user-a", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_global", source)
+}
+
+// TestTC_RESOLVE_005 都未配置 => 返回 block（source=default）
+func TestTC_RESOLVE_005(t *testing.T) {
+	config := AISecurityConfig{}
+
+	action, source := config.ResolveRiskActionByType("", SensitiveDataType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "default", source)
+
+	action, source = config.ResolveRiskActionByType("", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "default", source)
+}
+
+// =============================================================================
+// TC-MATCH: first-match 语义测试（getMatchedConsumerRiskRule）
+// =============================================================================
+
+// TestTC_MATCH_001 两条规则都可命中（prefix + exact），prefix 在前 => 命中 prefix
+func TestTC_MATCH_001(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":             Matcher{Prefix: "user-"},
+				"sensitiveDataAction": "mask",
+			},
+			{
+				"matcher":             Matcher{Exact: "user-a"},
+				"sensitiveDataAction": "block",
+			},
+		},
+	}
+
+	// "user-a" 同时匹配 prefix("user-") 和 exact("user-a")，但 prefix 在前
+	action, source := config.ResolveRiskActionByType("user-a", SensitiveDataType)
+	require.Equal(t, "mask", action)
+	require.Equal(t, "consumer_dimension", source)
+}
+
+// TestTC_MATCH_002 首条命中但未配置某维度动作，第二条配置了 => 不读取第二条，回退全局
+func TestTC_MATCH_002(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction:         "mask",
+		PromptAttackAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":    Matcher{Prefix: "user-"},
+				"riskAction": "mask",
+				// 未配置 promptAttackAction
+			},
+			{
+				"matcher":            Matcher{Exact: "user-a"},
+				"promptAttackAction": "block",
+			},
+		},
+	}
+
+	// "user-a" 命中首条 prefix 规则，promptAttackAction 未配置
+	// 回退到 consumer_global(mask)，然后 enforceMaskBoundary 降级为 block
+	action, source := config.ResolveRiskActionByType("user-a", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_global", source)
+}
+
+// TestTC_MATCH_003 无规则命中 => 回退全局
+func TestTC_MATCH_003(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction:          "mask",
+		SensitiveDataAction: "mask",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":             Matcher{Exact: "vip-user"},
+				"sensitiveDataAction": "block",
+			},
+		},
+	}
+
+	// "other-user" 不匹配任何规则，回退到 global_dimension
+	action, source := config.ResolveRiskActionByType("other-user", SensitiveDataType)
+	require.Equal(t, "mask", action)
+	require.Equal(t, "global_dimension", source)
+
+	// promptAttack 无 global_dimension，回退到 global_global(mask)，降级为 block
+	action, source = config.ResolveRiskActionByType("other-user", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_global", source)
+}
+
+// =============================================================================
+// 补充边界测试
+// =============================================================================
+
+// TestTC_RESOLVE_006 consumer 规则中 promptAttackAction=mask => 降级为 block
+func TestTC_RESOLVE_006(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":            Matcher{Exact: "user-a"},
+				"promptAttackAction": "mask", // 非 sensitiveData 维度配置 mask
+			},
+		},
+	}
+
+	// consumer_dimension(mask) 降级为 block
+	action, source := config.ResolveRiskActionByType("user-a", PromptAttackType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_dimension", source)
+}
+
+// TestTC_RESOLVE_007 consumer 规则中 contentModerationAction=mask => 降级为 block
+func TestTC_RESOLVE_007(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":                 Matcher{Exact: "user-a"},
+				"contentModerationAction": "mask",
+			},
+		},
+	}
+
+	action, source := config.ResolveRiskActionByType("user-a", ContentModerationType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_dimension", source)
+}
+
+// TestTC_RESOLVE_008 consumer 规则中 maliciousUrlAction=mask => 降级为 block
+func TestTC_RESOLVE_008(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":            Matcher{Exact: "user-a"},
+				"maliciousUrlAction": "mask",
+			},
+		},
+	}
+
+	action, source := config.ResolveRiskActionByType("user-a", MaliciousUrlDataType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_dimension", source)
+}
+
+// TestTC_RESOLVE_009 未知 detailType => dimensionActionKey 返回空，跳过 consumer_dimension
+func TestTC_RESOLVE_009(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":    Matcher{Exact: "user-a"},
+				"riskAction": "mask",
+			},
+		},
+	}
+
+	// 未知 Type，dimKey 为空，跳过 consumer_dimension，回退到 consumer_global(mask)
+	// 非 sensitiveData 维度的 mask 降级为 block
+	action, source := config.ResolveRiskActionByType("user-a", "unknownType")
+	require.Equal(t, "block", action)
+	require.Equal(t, "consumer_global", source)
+}
+
+// TestTC_RESOLVE_010 未知 detailType + 无 consumer 匹配 => 回退到 global_global
+func TestTC_RESOLVE_010(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction: "mask",
+	}
+
+	// 未知 Type，无 consumer 匹配，回退到 global_global(mask)
+	// 非 sensitiveData 维度的 mask 降级为 block
+	action, source := config.ResolveRiskActionByType("", "unknownType")
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_global", source)
+}
+
+// TestTC_RESOLVE_011 所有 6 个维度的 global dimension action 正确映射
+func TestTC_RESOLVE_011(t *testing.T) {
+	config := AISecurityConfig{
+		ContentModerationAction:  "block",
+		PromptAttackAction:       "block",
+		SensitiveDataAction:      "mask",
+		MaliciousUrlAction:       "block",
+		ModelHallucinationAction: "block",
+		CustomLabelAction:        "block",
+	}
+
+	tests := []struct {
+		detailType     string
+		expectedAction string
+		expectedSource string
+	}{
+		{ContentModerationType, "block", "global_dimension"},
+		{PromptAttackType, "block", "global_dimension"},
+		{SensitiveDataType, "mask", "global_dimension"},
+		{MaliciousUrlDataType, "block", "global_dimension"},
+		{ModelHallucinationDataType, "block", "global_dimension"},
+		{CustomLabelType, "block", "global_dimension"},
+	}
+
+	for _, tt := range tests {
+		action, source := config.ResolveRiskActionByType("", tt.detailType)
+		require.Equal(t, tt.expectedAction, action, "detailType=%s", tt.detailType)
+		require.Equal(t, tt.expectedSource, source, "detailType=%s", tt.detailType)
+	}
+}
+
+// TestTC_MATCH_004 空 consumer 不匹配 exact/prefix 规则 => 回退全局
+func TestTC_MATCH_004(t *testing.T) {
+	config := AISecurityConfig{
+		RiskAction:          "mask",
+		SensitiveDataAction: "block",
+		ConsumerRiskLevel: []map[string]interface{}{
+			{
+				"matcher":             Matcher{Exact: "vip"},
+				"sensitiveDataAction": "mask",
+			},
+		},
+	}
+
+	// 空 consumer 不匹配 exact("vip")，回退到 global_dimension(block)
+	action, source := config.ResolveRiskActionByType("", SensitiveDataType)
+	require.Equal(t, "block", action)
+	require.Equal(t, "global_dimension", source)
+}