OpenFlare/internal/apps/openflare/agent/waf_ip_group_test.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package agent

import (
	"context"
	"encoding/json"
	"strconv"
	"testing"

	"github.com/Rain-kl/Wavelet/internal/db"
	"github.com/Rain-kl/Wavelet/internal/model"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupWAFIPGroupTestDB(t *testing.T) func() {
	t.Helper()

	sqliteDB, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
	})
	require.NoError(t, err)
	require.NoError(t, sqliteDB.AutoMigrate(
		&model.OpenFlareWAFIPGroup{},
		&model.ConfigVersion{},
	))

	db.SetDB(sqliteDB)
	return func() {
		db.SetDB(nil)
	}
}

func seedActiveConfigWithWAFIPGroup(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":                   "agent refs",
					"enabled":                true,
					"ip_blacklist_group_ids": []uint{ipGroupID},
				},
			},
			"bindings": []any{},
		},
	}
	snapshotJSON, err := json.Marshal(snapshot)
	require.NoError(t, err)

	require.NoError(t, db.DB(ctx).Create(&model.ConfigVersion{
		Version:      "20260618-001",
		SnapshotJSON: string(snapshotJSON),
		Checksum:     "test-checksum",
		IsActive:     true,
	}).Error)
}

func TestChangedWAFIPGroupsForAgentReturnsChecksumDelta(t *testing.T) {
	cleanup := setupWAFIPGroupTestDB(t)
	defer cleanup()

	ctx := context.Background()
	ipGroup := &model.OpenFlareWAFIPGroup{
		Name:    "agent runtime group",
		Type:    "manual",
		Enabled: true,
		IPList:  `["203.0.113.44"]`,
	}
	require.NoError(t, model.CreateOpenFlareWAFIPGroup(ctx, ipGroup))
	seedActiveConfigWithWAFIPGroup(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, "203.0.113.44", groups[0].IPList[0])
	assert.NotEmpty(t, groups[0].Checksum)

	groupKey := strconv.FormatUint(uint64(ipGroup.ID), 10)
	same, err := ChangedWAFIPGroupsForAgent(ctx, nil, map[string]string{groupKey: groups[0].Checksum})
	require.NoError(t, err)
	assert.Empty(t, same)

	ipGroup.IPList = `["203.0.113.45"]`
	require.NoError(t, model.UpdateOpenFlareWAFIPGroup(ctx, ipGroup))

	delta, err := ChangedWAFIPGroupsForAgent(ctx, nil, map[string]string{groupKey: groups[0].Checksum})
	require.NoError(t, err)
	require.Len(t, delta, 1)
	assert.Equal(t, ipGroup.ID, delta[0].ID)
	assert.Equal(t, "203.0.113.45", delta[0].IPList[0])
	assert.NotEqual(t, groups[0].Checksum, delta[0].Checksum)
}

func TestSyncWAFIPGroupsReturnsChangedGroups(t *testing.T) {
	cleanup := setupWAFIPGroupTestDB(t)
	defer cleanup()

	ctx := context.Background()
	ipGroup := &model.OpenFlareWAFIPGroup{
		Name:    "sync group",
		Type:    "manual",
		Enabled: true,
		IPList:  `["198.51.100.10"]`,
	}
	require.NoError(t, model.CreateOpenFlareWAFIPGroup(ctx, ipGroup))
	seedActiveConfigWithWAFIPGroup(t, ctx, ipGroup.ID)

	result, err := SyncWAFIPGroups(ctx, WAFIPGroupSyncInput{
		IDs:       []uint{ipGroup.ID},
		Checksums: map[string]string{},
	})
	require.NoError(t, err)
	require.Len(t, result.Groups, 1)
	assert.Equal(t, ipGroup.ID, result.Groups[0].ID)
	assert.Equal(t, "198.51.100.10", result.Groups[0].IPList[0])

	result, err = SyncWAFIPGroups(ctx, WAFIPGroupSyncInput{
		IDs: []uint{ipGroup.ID},
		Checksums: map[string]string{
			strconv.FormatUint(uint64(ipGroup.ID), 10): result.Groups[0].Checksum,
		},
	})
	require.NoError(t, err)
	assert.Empty(t, result.Groups)
}

func TestChangedWAFIPGroupsForAgentDisabledGroupClearsIPList(t *testing.T) {
	cleanup := setupWAFIPGroupTestDB(t)
	defer cleanup()

	ctx := context.Background()
	ipGroup := &model.OpenFlareWAFIPGroup{
		Name:    "disabled group",
		Type:    "manual",
		Enabled: true,
		IPList:  `["203.0.113.10"]`,
	}
	require.NoError(t, model.CreateOpenFlareWAFIPGroup(ctx, ipGroup))
	ipGroup.Enabled = false
	require.NoError(t, model.UpdateOpenFlareWAFIPGroup(ctx, ipGroup))
	seedActiveConfigWithWAFIPGroup(t, ctx, ipGroup.ID)

	groups, err := ChangedWAFIPGroupsForAgent(ctx, nil, nil)
	require.NoError(t, err)
	require.Len(t, groups, 1)
	assert.False(t, groups[0].Enabled)
	assert.Empty(t, groups[0].IPList)
	assert.NotEmpty(t, groups[0].Checksum)
}