Files
LinkPort/internal/store/audit_shards.go
T
曾志威 9da95beb42
continuous-integration/drone Build is failing
feat: refine management UI, connection diagnostics and local startup
Migrate the workspace to React and Ant Design, improve account forms and responsive layouts, support outbound diagnostics and SSH fingerprint confirmation, and persist the management startup port. Preserve database migrations, remove S3 settings storage, and improve gateway audit handling.
2026-10-13 19:43:23 +08:00

312 lines
8.5 KiB
Go

package store
import (
"database/sql"
"errors"
"fmt"
"os"
"path/filepath"
"sort"
"strconv"
"strings"
"time"
)
const (
auditDayStride int64 = 1 << 40
auditShardStride int64 = 1 << 28
auditMaxParts = int(auditDayStride/auditShardStride) - 2
)
const auditSchema = `PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000; PRAGMA foreign_keys=ON;
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 TABLE IF NOT EXISTS command_results(command_id TEXT PRIMARY KEY REFERENCES commands(id) ON DELETE CASCADE,content BLOB NOT NULL);
CREATE TABLE IF NOT EXISTS command_outputs(seq INTEGER PRIMARY KEY AUTOINCREMENT,id TEXT NOT NULL UNIQUE,command_id TEXT NOT NULL REFERENCES commands(id) ON DELETE CASCADE,stream TEXT NOT NULL,content BLOB NOT NULL,size INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS command_outputs_command ON command_outputs(command_id,seq);
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 auditShardDate(key string) (time.Time, error) {
base := key
if index := strings.Index(base, ".part"); index >= 0 {
part := base[index+len(".part"):]
if part == "" {
return time.Time{}, errors.New("invalid audit shard part")
}
value, err := strconv.Atoi(part)
if err != nil {
return time.Time{}, err
}
if value < 1 || value > auditMaxParts {
return time.Time{}, errors.New("audit shard part is out of range")
}
base = base[:index]
}
return time.Parse("2006-01-02", base)
}
func auditShardPath(directory, key string) string {
return filepath.Join(directory, key+".sqlite")
}
func auditShardPart(key string) int {
if index := strings.Index(key, ".part"); index >= 0 {
part, _ := strconv.Atoi(key[index+len(".part"):])
return part
}
return 0
}
func auditShardBases(keys []string) map[string]int64 {
maxPart := map[string]int{}
for _, key := range keys {
date, err := auditShardDate(key)
if err != nil {
continue
}
day := date.Format("2006-01-02")
if part := auditShardPart(key); part > maxPart[day] {
maxPart[day] = part
}
}
bases := make(map[string]int64, len(keys))
for _, key := range keys {
date, err := auditShardDate(key)
if err != nil {
continue
}
part := auditShardPart(key)
if part == 0 {
part = maxPart[date.Format("2006-01-02")] + 1
}
bases[key] = date.Unix()/86400*auditDayStride + int64(part)*auditShardStride
}
return bases
}
func (s *Store) auditDB(day string) (*sql.DB, error) {
if _, err := auditShardDate(day); err != nil {
return nil, err
}
s.auditMu.Lock()
defer s.auditMu.Unlock()
if s.closed {
return nil, errors.New("store is closed")
}
return s.openAuditDBLocked(day)
}
func (s *Store) openAuditDBLocked(day string) (*sql.DB, error) {
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", auditShardPath(s.auditDir, day))
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
}
// auditDBForWrite rotates the active day shard before it grows beyond the
// configured limit. A single record may exceed the limit; it is never split
// across databases, so command/output writes remain atomic.
func (s *Store) auditDBForWrite(day string, estimate int64) (*sql.DB, error) {
if _, err := auditShardDate(day); err != nil {
return nil, err
}
s.auditMu.Lock()
defer s.auditMu.Unlock()
if s.closed {
return nil, errors.New("store is closed")
}
db, err := s.openAuditDBLocked(day)
if err != nil {
return nil, err
}
if estimate <= 0 || s.auditMaxBytes <= 0 {
return db, nil
}
current, err := s.auditFileBytesLocked(day)
if err != nil {
return nil, err
}
if current == 0 || current+estimate <= s.auditMaxBytes {
return db, nil
}
if err = s.rotateAuditLocked(day); err != nil {
return nil, err
}
return s.openAuditDBLocked(day)
}
func (s *Store) auditFileBytesLocked(day string) (int64, error) {
var total int64
for _, suffix := range []string{"", "-wal", "-shm"} {
info, err := os.Stat(auditShardPath(s.auditDir, day) + suffix)
if errors.Is(err, os.ErrNotExist) {
continue
}
if err != nil {
return 0, err
}
total += info.Size()
}
return total, nil
}
func (s *Store) rotateAuditLocked(day string) error {
db := s.audits[day]
if db != nil {
if _, err := db.Exec("PRAGMA wal_checkpoint(TRUNCATE)"); err != nil {
return err
}
var journalMode string
if err := db.QueryRow("PRAGMA journal_mode=DELETE").Scan(&journalMode); err != nil {
return err
}
if !strings.EqualFold(journalMode, "delete") {
return fmt.Errorf("close audit WAL before rotation: unexpected journal mode %q", journalMode)
}
if err := db.Close(); err != nil {
return err
}
delete(s.audits, day)
}
path := auditShardPath(s.auditDir, day)
if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
return nil
} else if err != nil {
return err
}
// DELETE mode leaves the main file self-contained. Remove any stale WAL
// sidecars before renaming so a crash cannot pair them with a new shard.
for _, suffix := range []string{"-wal", "-shm"} {
if err := os.Remove(path + suffix); err != nil && !errors.Is(err, os.ErrNotExist) {
return err
}
}
part, err := s.nextAuditPartLocked(day)
if err != nil {
return err
}
archived := auditShardPath(s.auditDir, day+".part"+formatAuditPart(part))
if err := os.Rename(path, archived); err != nil {
return err
}
return nil
}
func (s *Store) nextAuditPartLocked(day string) (int, error) {
entries, err := os.ReadDir(s.auditDir)
if err != nil {
return 0, err
}
prefix := day + ".part"
maxPart := 0
for _, entry := range entries {
if entry.IsDir() || !strings.HasPrefix(entry.Name(), prefix) || !strings.HasSuffix(entry.Name(), ".sqlite") {
continue
}
key := strings.TrimSuffix(entry.Name(), ".sqlite")
if _, err = auditShardDate(key); err != nil {
continue
}
if part := auditShardPart(key); part > maxPart {
maxPart = part
}
}
if maxPart >= auditMaxParts {
return 0, errors.New("daily audit shard limit exceeded")
}
return maxPart + 1, nil
}
func formatAuditPart(part int) string {
return fmt.Sprintf("%04d", part)
}
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 = auditShardDate(day); err == nil {
days = append(days, day)
}
}
sort.Slice(days, func(i, j int) bool {
left, _ := auditShardDate(days[i])
right, _ := auditShardDate(days[j])
if left.Equal(right) {
leftPart, rightPart := auditShardPart(days[i]), auditShardPart(days[j])
if leftPart == 0 {
return rightPart != 0
}
if rightPart == 0 {
return false
}
return leftPart > rightPart
}
return left.After(right)
})
return days, nil
}
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 {
date, parseErr := auditShardDate(day)
if parseErr != nil {
continue
}
dayDate := date.Format("2006-01-02")
if dayDate < 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{"", "-wal", "-shm"} {
_ = os.Remove(auditShardPath(s.auditDir, day) + suffix)
}
s.auditMu.Unlock()
} else if dayDate == 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))
}
}
}
}