feat(waf): complete composable rule orchestration

Add the React Flow rule editor, ordered graph APIs and runtime DAG execution.\n\nPublish rules only on OpenResty reload and reconcile checksum-driven IP group snapshots in bounded shared memory.
2026-10-06 07:36:37 +08:00 · 2026-07-13 14:16:55 +08:00
parent d36409fbf9
commit a1a997bcda
72 changed files with 5897 additions and 3080 deletions
@@ -4,6 +4,7 @@
 package agent

 import (
+	"bytes"
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
@@ -11,43 +12,32 @@ import (
 	"errors"
 	"fmt"
 	"sort"
+	"strconv"
 	"strings"

 	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/Rain-kl/Wavelet/pkg/protocol"
+	openrestyrender "github.com/Rain-kl/Wavelet/pkg/render/openresty"
 )

-type snapshotWAFRuleGroupRef struct {
-	IPWhitelistGroups []uint `json:"ip_whitelist_group_ids,omitempty"`
-	IPBlacklistGroups []uint `json:"ip_blacklist_group_ids,omitempty"`
-}
-
-type snapshotWAFSection struct {
-	RuleGroups []snapshotWAFRuleGroupRef `json:"rule_groups"`
-}
-
 type activeConfigSnapshot struct {
-	WAF snapshotWAFSection `json:"waf"`
+	WAF openrestyrender.WAFDocument `json:"waf"`
+}
+
+type runtimeIPMatchConfig struct {
+	IPs        []string `json:"ips,omitempty"`
+	CIDRs      []string `json:"cidrs,omitempty"`
+	IPGroupIDs []uint   `json:"ip_group_ids,omitempty"`
 }

 // WAFIPGroupsForAgent builds agent-facing WAF IP group payloads for the given ids.
 func WAFIPGroupsForAgent(ctx context.Context, ids []uint) ([]WAFIPGroup, error) {
-	return buildAgentWAFIPGroups(ctx, ids)
+	return validatedAgentWAFIPGroups(ctx, ids, false)
 }

 // ChangedWAFIPGroupsForAgent returns WAF IP groups whose checksums differ from the agent state.
 func ChangedWAFIPGroupsForAgent(ctx context.Context, ids []uint, checksums map[string]string) ([]WAFIPGroup, error) {
-	targetIDs := uniqueUintIDs(ids)
-	if len(targetIDs) == 0 {
-		activeIDs, err := activeConfigWAFIPGroupIDs(ctx)
-		if err != nil {
-			return nil, err
-		}
-		targetIDs = activeIDs
-	}
-	if len(targetIDs) == 0 {
-		return []WAFIPGroup{}, nil
-	}
-	groups, err := buildAgentWAFIPGroups(ctx, targetIDs)
+	groups, err := validatedAgentWAFIPGroups(ctx, ids, true)
 	if err != nil {
 		return nil, err
 	}
@@ -61,6 +51,45 @@ func ChangedWAFIPGroupsForAgent(ctx context.Context, ids []uint, checksums map[s
 	return changed, nil
 }

+func validatedAgentWAFIPGroups(ctx context.Context, ids []uint, fallbackToActive bool) ([]WAFIPGroup, error) {
+	targetIDs := uniqueUintIDs(ids)
+	activeIDs, err := activeConfigWAFIPGroupIDs(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if len(targetIDs) == 0 && fallbackToActive {
+		targetIDs = activeIDs
+	}
+	if len(targetIDs) == 0 {
+		return []WAFIPGroup{}, nil
+	}
+
+	validationIDs := uniqueUintIDs(append(append([]uint{}, activeIDs...), targetIDs...))
+	allGroups, err := buildAgentWAFIPGroups(ctx, validationIDs)
+	if err != nil {
+		return nil, err
+	}
+	runtimeGroups := make(map[string]protocol.WAFIPGroup, len(allGroups))
+	for _, group := range allGroups {
+		runtimeGroups[strconv.FormatUint(uint64(group.ID), 10)] = group
+	}
+	if err = protocol.ValidateWAFIPGroupSnapshotSize(runtimeGroups); err != nil {
+		return nil, err
+	}
+
+	targetSet := make(map[uint]struct{}, len(targetIDs))
+	for _, id := range targetIDs {
+		targetSet[id] = struct{}{}
+	}
+	result := make([]WAFIPGroup, 0, len(targetIDs))
+	for _, group := range allGroups {
+		if _, ok := targetSet[group.ID]; ok {
+			result = append(result, group)
+		}
+	}
+	return result, nil
+}
+
 func buildAgentWAFIPGroups(ctx context.Context, ids []uint) ([]WAFIPGroup, error) {
 	ids = uniqueUintIDs(ids)
 	if len(ids) == 0 {
@@ -142,6 +171,8 @@ func activeConfigWAFIPGroupIDs(ctx context.Context) ([]uint, error) {
 	}
 	idSet := make(map[uint]struct{})
 	for _, group := range snapshot.WAF.RuleGroups {
+		// Retain legacy flattened references while older active snapshots may
+		// still exist during a rolling Server upgrade.
 		for _, id := range group.IPWhitelistGroups {
 			if id > 0 {
 				idSet[id] = struct{}{}
@@ -152,6 +183,20 @@ func activeConfigWAFIPGroupIDs(ctx context.Context) ([]uint, error) {
 				idSet[id] = struct{}{}
 			}
 		}
+		for nodeID, node := range group.Graph.Nodes {
+			if node.Type != "ip_match" {
+				continue
+			}
+			ids, err := runtimeIPMatchGroupIDs(node.Config)
+			if err != nil {
+				return nil, fmt.Errorf("活动配置 WAF 规则 %d 节点 %s 的 IP 匹配配置无效: %w", group.ID, nodeID, err)
+			}
+			for _, id := range ids {
+				if id > 0 {
+					idSet[id] = struct{}{}
+				}
+			}
+		}
 	}
 	ids := make([]uint, 0, len(idSet))
 	for id := range idSet {
@@ -161,6 +206,16 @@ func activeConfigWAFIPGroupIDs(ctx context.Context) ([]uint, error) {
 	return ids, nil
 }

+func runtimeIPMatchGroupIDs(raw json.RawMessage) ([]uint, error) {
+	var config runtimeIPMatchConfig
+	decoder := json.NewDecoder(bytes.NewReader(raw))
+	decoder.DisallowUnknownFields()
+	if err := decoder.Decode(&config); err != nil {
+		return nil, err
+	}
+	return config.IPGroupIDs, nil
+}
+
 func parseActiveConfigSnapshot(snapshotJSON string) (*activeConfigSnapshot, error) {
 	text := strings.TrimSpace(snapshotJSON)
 	if text == "" {
@@ -171,7 +226,7 @@ func parseActiveConfigSnapshot(snapshotJSON string) (*activeConfigSnapshot, erro
 		return nil, err
 	}
 	if snapshot.WAF.RuleGroups == nil {
-		snapshot.WAF.RuleGroups = []snapshotWAFRuleGroupRef{}
+		snapshot.WAF.RuleGroups = []openrestyrender.WAFRuleGroup{}
 	}
 	return &snapshot, nil
 }
@@ -7,10 +7,12 @@ import (
 	"context"
 	"encoding/json"
 	"strconv"
+	"strings"
 	"testing"

 	"github.com/Rain-kl/Wavelet/internal/db"
 	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/Rain-kl/Wavelet/pkg/protocol"
 	"github.com/glebarez/sqlite"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -63,6 +65,115 @@ func seedActiveConfigWithWAFIPGroup(t *testing.T, ctx context.Context, ipGroupID
 	}).Error)
 }

+func seedActiveConfigWithWAFGraphIPGroup(t *testing.T, ctx context.Context, ipGroupID uint) {
+	t.Helper()
+
+	snapshot := map[string]any{
+		"routes": []any{},
+		"waf": map[string]any{
+			"rule_groups": []map[string]any{
+				{
+					"id":      1,
+					"name":    "graph refs",
+					"enabled": true,
+					"graph": map[string]any{
+						"entry": "start",
+						"nodes": map[string]any{
+							"start": map[string]any{
+								"type":   "start",
+								"config": map[string]any{},
+								"next":   map[string]string{"next": "match"},
+							},
+							"match": map[string]any{
+								"type": "ip_match",
+								"config": map[string]any{
+									"ip_group_ids": []uint{ipGroupID},
+								},
+								"next": map[string]string{"true": "allow", "false": "allow"},
+							},
+							"allow": map[string]any{
+								"type":   "allow",
+								"config": map[string]any{},
+							},
+						},
+					},
+				},
+			},
+			"bindings": []any{},
+		},
+	}
+	snapshotJSON, err := json.Marshal(snapshot)
+	require.NoError(t, err)
+
+	require.NoError(t, db.DB(ctx).Create(&model.ConfigVersion{
+		Version:      "20260713-graph-001",
+		SnapshotJSON: string(snapshotJSON),
+		Checksum:     "graph-test-checksum",
+		IsActive:     true,
+	}).Error)
+}
+
+func TestChangedWAFIPGroupsForAgentDiscoversGraphReferences(t *testing.T) {
+	cleanup := setupWAFIPGroupTestDB(t)
+	defer cleanup()
+
+	ctx := context.Background()
+	ipGroup := &model.OpenFlareWAFIPGroup{
+		Name:    "graph runtime group",
+		Type:    "manual",
+		Enabled: true,
+		IPList:  `["192.0.2.88"]`,
+	}
+	require.NoError(t, model.CreateOpenFlareWAFIPGroup(ctx, ipGroup))
+	seedActiveConfigWithWAFGraphIPGroup(t, ctx, ipGroup.ID)
+
+	groups, err := ChangedWAFIPGroupsForAgent(ctx, nil, nil)
+	require.NoError(t, err)
+	require.Len(t, groups, 1)
+	assert.Equal(t, ipGroup.ID, groups[0].ID)
+	assert.Equal(t, []string{"192.0.2.88"}, groups[0].IPList)
+}
+
+func TestChangedWAFIPGroupsForAgentRejectsMalformedIPMatchConfig(t *testing.T) {
+	cleanup := setupWAFIPGroupTestDB(t)
+	defer cleanup()
+
+	ctx := context.Background()
+	require.NoError(t, db.DB(ctx).Create(&model.ConfigVersion{
+		Version: "20260713-malformed-001",
+		SnapshotJSON: `{"waf":{"rule_groups":[{"id":7,"graph":{"entry":"match","nodes":{` +
+			`"match":{"type":"ip_match","config":{"ip_group_ids":"not-an-array"}}}}}],"bindings":[]}}`,
+		Checksum: "malformed-test-checksum",
+		IsActive: true,
+	}).Error)
+
+	_, err := ChangedWAFIPGroupsForAgent(ctx, nil, nil)
+	require.ErrorContains(t, err, "规则 7 节点 match")
+	require.ErrorContains(t, err, "IP 匹配配置无效")
+}
+
+func TestChangedWAFIPGroupsForAgentRejectsOversizedSnapshotBeforeChecksumDelta(t *testing.T) {
+	cleanup := setupWAFIPGroupTestDB(t)
+	defer cleanup()
+
+	ctx := context.Background()
+	ipGroup := &model.OpenFlareWAFIPGroup{
+		Name:    strings.Repeat("x", protocol.MaxWAFIPGroupSnapshotBytes),
+		Type:    "manual",
+		Enabled: true,
+		IPList:  `[]`,
+	}
+	require.NoError(t, model.CreateOpenFlareWAFIPGroup(ctx, ipGroup))
+	agentGroup, err := buildAgentWAFIPGroup(ipGroup)
+	require.NoError(t, err)
+
+	_, err = ChangedWAFIPGroupsForAgent(ctx, []uint{ipGroup.ID}, map[string]string{
+		strconv.FormatUint(uint64(ipGroup.ID), 10): agentGroup.Checksum,
+	})
+	require.ErrorContains(t, err, "WAF IP 组快照大小")
+	require.ErrorContains(t, err, "超过上限")
+}
+
 func TestChangedWAFIPGroupsForAgentReturnsChecksumDelta(t *testing.T) {
 	cleanup := setupWAFIPGroupTestDB(t)
 	defer cleanup()