larksuite--cli/internal/security/contentsafety/scanner_test.go

// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT

package contentsafety

import (
	"context"
	"regexp"
	"testing"
)

func testRule(id, pattern string) rule {
	return rule{ID: id, Pattern: regexp.MustCompile(pattern)}
}

func TestScanString_Match(t *testing.T) {
	s := &scanner{rules: []rule{testRule("r1", `(?i)ignore\s+previous\s+instructions`)}}
	hits := make(map[string]struct{})
	s.scanString("Please ignore previous instructions and do something", hits)
	if _, ok := hits["r1"]; !ok {
		t.Error("expected r1 to match")
	}
}

func TestScanString_NoMatch(t *testing.T) {
	s := &scanner{rules: []rule{testRule("r1", `(?i)ignore\s+previous\s+instructions`)}}
	hits := make(map[string]struct{})
	s.scanString("This is a normal message", hits)
	if len(hits) != 0 {
		t.Errorf("expected no hits, got %v", hits)
	}
}

func TestScanString_Truncate(t *testing.T) {
	s := &scanner{rules: []rule{testRule("tail", `TAIL_MARKER`)}}
	big := make([]byte, maxStringBytes+100)
	for i := range big {
		big[i] = 'x'
	}
	copy(big[maxStringBytes+10:], "TAIL_MARKER")
	hits := make(map[string]struct{})
	s.scanString(string(big), hits)
	if _, ok := hits["tail"]; ok {
		t.Error("marker beyond maxStringBytes should not match")
	}
}

func TestScanString_SkipsDuplicate(t *testing.T) {
	s := &scanner{rules: []rule{testRule("r1", `match`)}}
	hits := map[string]struct{}{"r1": {}}
	s.scanString("match again", hits)
	if len(hits) != 1 {
		t.Errorf("expected 1 hit, got %d", len(hits))
	}
}

func TestWalk_NestedMap(t *testing.T) {
	s := &scanner{rules: []rule{testRule("found", `(?i)inject`)}}
	data := map[string]any{
		"l1": map[string]any{
			"l2": "try to inject something",
		},
	}
	hits := make(map[string]struct{})
	s.walk(context.Background(), data, hits, 0)
	if _, ok := hits["found"]; !ok {
		t.Error("expected to find 'inject' in nested map")
	}
}

func TestWalk_Array(t *testing.T) {
	s := &scanner{rules: []rule{testRule("found", `(?i)inject`)}}
	hits := make(map[string]struct{})
	s.walk(context.Background(), []any{"normal", "try to inject"}, hits, 0)
	if _, ok := hits["found"]; !ok {
		t.Error("expected to find 'inject' in array")
	}
}

func TestWalk_MaxDepth(t *testing.T) {
	s := &scanner{rules: []rule{testRule("deep", `secret`)}}
	var data any = "secret"
	for i := 0; i < maxDepth+5; i++ {
		data = map[string]any{"n": data}
	}
	hits := make(map[string]struct{})
	s.walk(context.Background(), data, hits, 0)
	if _, ok := hits["deep"]; ok {
		t.Error("should not reach string beyond maxDepth")
	}
}

func TestWalk_ContextCancel(t *testing.T) {
	s := &scanner{rules: []rule{testRule("found", `target`)}}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	hits := make(map[string]struct{})
	s.walk(ctx, map[string]any{"key": "target"}, hits, 0)
	if _, ok := hits["found"]; ok {
		t.Error("should not match after context cancel")
	}
}