LinkPort/internal/store/audit_shards_test.go

package store

import (
	"agentdb/internal/model"
	"agentdb/internal/security"
	"bytes"
	"context"
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

func shardFixture(t *testing.T) (*Store, string, *security.Vault) {
	t.Helper()
	v, err := security.NewVault(base64.StdEncoding.EncodeToString(make([]byte, 32)))
	if err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(t.TempDir(), "agentdb.sqlite")
	s, err := Open(path, v)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.Close() })
	return s, path, v
}

func TestAuditShardsCrossDayPaginationFilteringAndWalk(t *testing.T) {
	s, _, _ := shardFixture(t)
	now := time.Date(2026, 9, 9, 12, 0, 0, 0, time.UTC)
	p := model.Proxy{ID: "gateway", TargetID: "target", Type: "mysql"}
	for day := 0; day < 3; day++ {
		for i := 0; i < 137; i++ {
			text := fmt.Sprintf("SELECT %d, %d", day, i)
			if day == 2 && i < 3 {
				text = "SELECT 'old-page-match'"
			}
			err := s.recordCommandAt(p, "192.0.2.1", model.Command{Kind: "query", Decision: "allowed", Text: text}, now.AddDate(0, 0, -day).Add(time.Duration(i)*time.Second))
			if err != nil {
				t.Fatal(err)
			}
		}
	}
	seen := map[string]bool{}
	var cursor int64
	for page, expected := range []int{200, 200, 11, 0} {
		recs, err := s.Commands("gateway", "mysql", cursor)
		if err != nil || len(recs) != expected {
			t.Fatalf("page %d got %d want %d: %v", page, len(recs), expected, err)
		}
		for _, rec := range recs {
			if seen[rec.ID] || (cursor != 0 && rec.RowID >= cursor) {
				t.Fatalf("duplicate or unordered record %+v", rec)
			}
			seen[rec.ID] = true
			cursor = rec.RowID
		}
	}
	if len(seen) != 411 {
		t.Fatal("missing records", len(seen))
	}
	f := CommandFilter{Gateway: "gateway", Target: "target", Source: "192.0.2.1", Protocol: "mysql", Kind: "query", Decision: "allowed", Keyword: "OLD-PAGE-MATCH"}
	recs, err := s.FilterCommands(context.Background(), f, 0)
	if err != nil || len(recs) != 3 {
		t.Fatalf("older keyword matches: %d %v", len(recs), err)
	}
	f.From = now.AddDate(0, 0, -1).Format(time.RFC3339Nano)
	recs, err = s.FilterCommands(context.Background(), f, 0)
	if err != nil || len(recs) != 0 {
		t.Fatalf("time filter: %d %v", len(recs), err)
	}
	exported := map[string]bool{}
	err = s.WalkCommands(context.Background(), CommandFilter{}, 0, 0, func(rec CommandRecord) error {
		if exported[rec.ID] {
			t.Fatal("duplicate export row", rec.ID)
		}
		exported[rec.ID] = true
		return nil
	})
	if err != nil || len(exported) != 411 {
		t.Fatalf("cross-day export got %d: %v", len(exported), err)
	}
	stop := errors.New("export disconnected")
	if err = s.WalkCommands(context.Background(), CommandFilter{}, 0, 0, func(CommandRecord) error { return stop }); !errors.Is(err, stop) {
		t.Fatal("visitor error lost", err)
	}
}

func TestAuditShardsMigrationRestartAndEncryption(t *testing.T) {
	s, path, v := shardFixture(t)
	for i := 0; i < 3; i++ {
		id := fmt.Sprintf("legacy-%d", i)
		at := time.Date(2026, 9, 9-i, 12, 0, 0, 0, time.UTC)
		blob, err := s.seal(model.Command{Kind: "query", Decision: "allowed", Text: "SELECT 'legacy-secret'"}, "command:"+id)
		if err != nil {
			t.Fatal(err)
		}
		args := []any{id, "gateway", "target", "192.0.2.1", "mysql", "query", "allowed", blob, at.Format(time.RFC3339Nano)}
		const insert = "INSERT INTO commands(id,proxy_id,target_id,source,protocol,kind,decision,content,at) VALUES(?,?,?,?,?,?,?,?,?)"
		if _, err = s.db.Exec(insert, args...); err != nil {
			t.Fatal(err)
		}
		if i == 0 { // Simulate interruption after copying one row, before legacy deletion.
			db, e := s.auditDB(at.Format("2006-01-02"))
			if e != nil {
				t.Fatal(e)
			}
			if _, e = db.Exec(insert, args...); e != nil {
				t.Fatal(e)
			}
		}
	}
	s.Close()
	for attempt := 0; attempt < 2; attempt++ {
		reopened, err := Open(path, v)
		if err != nil {
			t.Fatal(err)
		}
		recs, err := reopened.Commands("", "", 0)
		if err != nil || len(recs) != 3 {
			reopened.Close()
			t.Fatalf("restart %d: %d %v", attempt, len(recs), err)
		}
		for _, rec := range recs {
			if rec.Text != "SELECT 'legacy-secret'" {
				t.Fatal("migration changed content")
			}
		}
		var count int
		if err = reopened.db.QueryRow("SELECT COUNT(*) FROM commands").Scan(&count); err != nil || count != 0 {
			t.Fatal("legacy data not cleared", count, err)
		}
		reopened.Close()
	}
	files, err := filepath.Glob(filepath.Join(filepath.Dir(path), "audit", "*"))
	if err != nil {
		t.Fatal(err)
	}
	if len(files) != 3 {
		t.Fatal("daily shard count", len(files))
	}
	for _, file := range files {
		data, err := os.ReadFile(file)
		if err != nil {
			t.Fatal(err)
		}
		if bytes.Contains(data, []byte("legacy-secret")) {
			t.Fatal("plaintext persisted in", file)
		}
	}
}

func TestAuditShardsMigrationConflictPreservesLegacy(t *testing.T) {
	s, path, v := shardFixture(t)
	at := time.Date(2026, 9, 9, 12, 0, 0, 0, time.UTC)
	blob, err := s.seal(model.Command{Text: "SELECT 'original'"}, "command:collision")
	if err != nil {
		t.Fatal(err)
	}
	args := []any{"collision", "gateway", "target", "source", "mysql", "query", "allowed", blob, at.Format(time.RFC3339Nano)}
	const insert = "INSERT INTO commands(id,proxy_id,target_id,source,protocol,kind,decision,content,at) VALUES(?,?,?,?,?,?,?,?,?)"
	if _, err = s.db.Exec(insert, args...); err != nil {
		t.Fatal(err)
	}
	db, err := s.auditDB(at.Format("2006-01-02"))
	if err != nil {
		t.Fatal(err)
	}
	args[3] = "different-source"
	if _, err = db.Exec(insert, args...); err != nil {
		t.Fatal(err)
	}
	s.Close()
	if reopened, err := Open(path, v); err == nil {
		reopened.Close()
		t.Fatal("conflicting migration accepted")
	}
	legacy, err := sql.Open("sqlite", path)
	if err != nil {
		t.Fatal(err)
	}
	defer legacy.Close()
	var count int
	if err = legacy.QueryRow("SELECT COUNT(*) FROM commands").Scan(&count); err != nil || count != 1 {
		t.Fatal("legacy row lost", count, err)
	}
}

func TestAuditShardsFailedMigrationCanResume(t *testing.T) {
	s, path, v := shardFixture(t)
	// Put the invalid timestamp after a committed batch to exercise partial progress.
	for i := 0; i < 1002; i++ {
		id := fmt.Sprintf("restart-%d", i)
		blob, err := s.seal(model.Command{Text: "SELECT 1"}, "command:"+id)
		if err != nil {
			t.Fatal(err)
		}
		at := "2026-09-09T12:00:00Z"
		if i == 1001 {
			at = "invalid"
		}
		if _, err = s.db.Exec("INSERT INTO commands(id,proxy_id,target_id,source,protocol,kind,decision,content,at) VALUES(?,?,?,?,?,?,?,?,?)", id, "gateway", "target", "source", "mysql", "query", "allowed", blob, at); err != nil {
			t.Fatal(err)
		}
	}
	s.Close()
	if reopened, err := Open(path, v); err == nil {
		reopened.Close()
		t.Fatal("invalid migration timestamp accepted")
	}
	legacy, err := sql.Open("sqlite", path)
	if err != nil {
		t.Fatal(err)
	}
	var count int
	if err = legacy.QueryRow("SELECT COUNT(*) FROM commands").Scan(&count); err != nil || count != 1002 {
		t.Fatal("legacy rows lost after interruption", count, err)
	}
	if _, err = legacy.Exec("UPDATE commands SET at='2026-09-09T12:00:00Z' WHERE at='invalid'"); err != nil {
		t.Fatal(err)
	}
	legacy.Close()
	reopened, err := Open(path, v)
	if err != nil {
		t.Fatal(err)
	}
	defer reopened.Close()
	seen := map[string]bool{}
	err = reopened.WalkCommands(context.Background(), CommandFilter{}, 0, 0, func(rec CommandRecord) error {
		if seen[rec.ID] {
			t.Fatal("duplicate after migration retry")
		}
		seen[rec.ID] = true
		return nil
	})
	if err != nil || len(seen) != 1002 {
		t.Fatal("incomplete migration retry", len(seen), err)
	}
}

func TestAuditShardsRetentionAndCancellation(t *testing.T) {
	s, _, _ := shardFixture(t)
	now := time.Date(2026, 9, 9, 12, 0, 0, 0, time.UTC)
	p := model.Proxy{ID: "gateway", TargetID: "target", Type: "mysql"}
	for _, at := range []time.Time{now.Add(-8 * 24 * time.Hour), now.Add(-7*24*time.Hour - time.Second), now.Add(-7 * 24 * time.Hour), now} {
		if err := s.recordCommandAt(p, "source", model.Command{Text: "SELECT 1"}, at); err != nil {
			t.Fatal(err)
		}
	}
	for _, name := range []string{"keep.txt", "not-a-date.sqlite", "2026-09-01.sqlite.backup"} {
		if err := os.WriteFile(filepath.Join(s.auditDir, name), []byte("preserve"), 0600); err != nil {
			t.Fatal(err)
		}
	}
	s.pruneCommands(now.Add(-7 * 24 * time.Hour).In(time.FixedZone("offset", -13*3600)))
	recs, err := s.Commands("", "", 0)
	if err != nil || len(recs) != 2 {
		t.Fatal("retention boundary", len(recs), err)
	}
	if _, err = os.Stat(filepath.Join(s.auditDir, "2026-09-01.sqlite")); !errors.Is(err, os.ErrNotExist) {
		t.Fatal("old shard exists", err)
	}
	for _, name := range []string{"keep.txt", "not-a-date.sqlite", "2026-09-01.sqlite.backup"} {
		if _, err = os.Stat(filepath.Join(s.auditDir, name)); err != nil {
			t.Fatal("unrelated file removed", name, err)
		}
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err = s.FilterCommands(ctx, CommandFilter{}, 0); !errors.Is(err, context.Canceled) {
		t.Fatal("cancellation lost", err)
	}
	ctx, cancel = context.WithCancel(context.Background())
	err = s.WalkCommands(ctx, CommandFilter{}, 0, 0, func(CommandRecord) error { cancel(); return nil })
	if !errors.Is(err, context.Canceled) {
		t.Fatal("mid-export cancellation lost", err)
	}
}

func TestAuditShardsConcurrentReadWritePrune(t *testing.T) {
	s, _, _ := shardFixture(t)
	now := time.Now().UTC()
	p := model.Proxy{ID: "gateway", Type: "mysql"}
	if err := s.recordCommandAt(p, "source", model.Command{Text: "old"}, now.Add(-8*24*time.Hour)); err != nil {
		t.Fatal(err)
	}
	var wg sync.WaitGroup
	errCh := make(chan error, 100)
	for worker := 0; worker < 3; worker++ {
		wg.Add(1)
		go func(worker int) {
			defer wg.Done()
			for i := 0; i < 20; i++ {
				var err error
				switch worker {
				case 0:
					err = s.RecordCommand(p, "source", model.Command{Text: "current"})
				case 1:
					_, err = s.Commands("", "", 0)
				case 2:
					s.pruneCommands(now.Add(-7 * 24 * time.Hour))
				}
				if err != nil {
					errCh <- err
				}
			}
		}(worker)
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
		t.Error(err)
	}
	recs, err := s.Commands("", "", 0)
	if err != nil || len(recs) != 20 {
		t.Fatal("concurrent writes lost", len(recs), err)
	}
	s.Close()
	if err = s.RecordCommand(p, "source", model.Command{Text: "closed"}); err == nil {
		t.Fatal("write after close accepted")
	}
	if _, err = s.Commands("", "", 0); err == nil {
		t.Fatal("read after close accepted")
	}
}