LinkPort/internal/store/audit_shards.go

package store

import (
	"database/sql"
	"errors"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

const shardStride int64 = 1 << 32
const auditSchema = `PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;
CREATE TABLE IF NOT EXISTS commands(id TEXT PRIMARY KEY,proxy_id TEXT NOT NULL,target_id TEXT NOT NULL,source TEXT NOT NULL,protocol TEXT NOT NULL,kind TEXT NOT NULL,decision TEXT NOT NULL,content BLOB NOT NULL,at TEXT NOT NULL);
CREATE INDEX IF NOT EXISTS commands_at ON commands(at);
CREATE INDEX IF NOT EXISTS commands_proxy ON commands(proxy_id);
CREATE INDEX IF NOT EXISTS commands_target ON commands(target_id);
CREATE INDEX IF NOT EXISTS commands_filter ON commands(protocol,decision,kind);
CREATE TABLE IF NOT EXISTS events(id INTEGER PRIMARY KEY AUTOINCREMENT,proxy_id TEXT NOT NULL,source TEXT NOT NULL,category TEXT NOT NULL,detail TEXT NOT NULL,at TEXT NOT NULL);
CREATE INDEX IF NOT EXISTS events_at ON events(at);
CREATE INDEX IF NOT EXISTS events_proxy ON events(proxy_id);`

func (s *Store) auditDB(day string) (*sql.DB, error) {
	if _, err := time.Parse("2006-01-02", day); err != nil {
		return nil, err
	}
	s.auditMu.Lock()
	defer s.auditMu.Unlock()
	if s.closed {
		return nil, errors.New("store is closed")
	}
	if db := s.audits[day]; db != nil {
		return db, nil
	}
	if err := os.MkdirAll(s.auditDir, 0700); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", filepath.Join(s.auditDir, day+".sqlite"))
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(1)
	if _, err = db.Exec(auditSchema); err != nil {
		db.Close()
		return nil, err
	}
	s.audits[day] = db
	return db, nil
}
func (s *Store) auditDays() ([]string, error) {
	entries, err := os.ReadDir(s.auditDir)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	days := []string{}
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sqlite") {
			continue
		}
		day := strings.TrimSuffix(entry.Name(), ".sqlite")
		if _, err = time.Parse("2006-01-02", day); err == nil {
			days = append(days, day)
		}
	}
	sort.Sort(sort.Reverse(sort.StringSlice(days)))
	return days, nil
}
func (s *Store) migrateCommands() error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	rows, err := tx.Query("SELECT id,proxy_id,target_id,source,protocol,kind,decision,content,at FROM commands ORDER BY rowid")
	if err != nil {
		return err
	}
	migrated := false
	var shardTx *sql.Tx
	var shardDay string
	batchSize := 0
	defer func() {
		if shardTx != nil {
			shardTx.Rollback()
		}
	}()
	for rows.Next() {
		var id, proxy, target, source, protocol, kind, decision, at string
		var content []byte
		if err = rows.Scan(&id, &proxy, &target, &source, &protocol, &kind, &decision, &content, &at); err != nil {
			rows.Close()
			return err
		}
		date, e := time.Parse(time.RFC3339Nano, at)
		if e != nil {
			rows.Close()
			return e
		}
		day := date.UTC().Format("2006-01-02")
		if shardTx != nil && (day != shardDay || batchSize >= 1000) {
			if e = shardTx.Commit(); e != nil {
				rows.Close()
				return e
			}
			shardTx = nil
		}
		if shardTx == nil {
			db, e := s.auditDB(day)
			if e != nil {
				rows.Close()
				return e
			}
			shardTx, e = db.Begin()
			if e != nil {
				rows.Close()
				return e
			}
			shardDay, batchSize = day, 0
		}
		if _, e = shardTx.Exec("INSERT OR IGNORE INTO commands(id,proxy_id,target_id,source,protocol,kind,decision,content,at) VALUES(?,?,?,?,?,?,?,?,?)", id, proxy, target, source, protocol, kind, decision, content, at); e != nil {
			rows.Close()
			return e
		}
		// A previous interrupted migration may already have copied this row.
		// Only discard the legacy copy when all persisted values match exactly.
		var matches int
		e = shardTx.QueryRow("SELECT COUNT(*) FROM commands WHERE id=? AND proxy_id=? AND target_id=? AND source=? AND protocol=? AND kind=? AND decision=? AND content=? AND at=?", id, proxy, target, source, protocol, kind, decision, content, at).Scan(&matches)
		if e != nil || matches != 1 {
			rows.Close()
			if e != nil {
				return e
			}
			return errors.New("audit migration found conflicting record " + id)
		}
		migrated = true
		batchSize++
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return err
	}
	if migrated {
		if err = shardTx.Commit(); err != nil {
			return err
		}
		shardTx = nil
		if _, err = tx.Exec("DELETE FROM commands"); err != nil {
			return err
		}
	}
	if err = tx.Commit(); err != nil {
		return err
	}
	if migrated {
		_, err = s.db.Exec("VACUUM")
	}
	return err
}
func (s *Store) pruneCommands(cutoff time.Time) {
	s.auditOps.Lock()
	defer s.auditOps.Unlock()
	if s.closed {
		return
	}
	cutoff = cutoff.UTC()
	days, err := s.auditDays()
	if err != nil {
		return
	}
	for _, day := range days {
		if day < cutoff.Format("2006-01-02") {
			s.auditMu.Lock()
			if db := s.audits[day]; db != nil {
				db.Close()
				delete(s.audits, day)
			}
			for _, suffix := range []string{".sqlite", ".sqlite-wal", ".sqlite-shm"} {
				_ = os.Remove(filepath.Join(s.auditDir, day+suffix))
			}
			s.auditMu.Unlock()
		} else if day == cutoff.Format("2006-01-02") {
			db, err := s.auditDB(day)
			if err == nil {
				db.Exec("DELETE FROM commands WHERE julianday(at)<julianday(?)", cutoff.Format(time.RFC3339Nano))
				db.Exec("DELETE FROM events WHERE julianday(at)<julianday(?)", cutoff.Format(time.RFC3339Nano))
			}
		}
	}
}