auth-proxy/src/db.go

package main

import (
	"database/sql"
	"errors"
	"log/slog"
	"net"
	"strings"
	"time"

	sqlite "modernc.org/sqlite"
)

func InitDB(dbPath string) (*sql.DB, error) {
	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, err
	}

	if _, err = db.Exec(`PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;`); err != nil {
		return nil, err
	}

	schema := `
CREATE TABLE IF NOT EXISTS perm_whitelist (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	ip TEXT NOT NULL UNIQUE,
	created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS temp_whitelist (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	ip TEXT NOT NULL UNIQUE,
	expires_at DATETIME NOT NULL,
	reason TEXT,
	created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_temp_expires ON temp_whitelist(expires_at);

CREATE TABLE IF NOT EXISTS audit_log (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	timestamp DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
	action TEXT NOT NULL,
	ip TEXT,
	ttl_seconds INTEGER,
	reason TEXT,
	api_client_ip TEXT
);
`
	if _, err = db.Exec(schema); err != nil {
		return nil, err
	}

	return db, nil
}

// isUniqueConstraintError checks if the error is a SQLite UNIQUE constraint violation (code 19).
func isUniqueConstraintError(err error) bool {
	if err == nil {
		return false
	}
	var sqliteErr *sqlite.Error
	return errors.As(err, &sqliteErr) && sqliteErr.Code() == 19
}

// IsIP returns true if s is a valid IP address (v4 or v6).
func IsIP(s string) bool {
	return net.ParseIP(s) != nil
}

// IsCIDR returns true if s is a valid CIDR notation.
func IsCIDR(s string) bool {
	_, _, err := net.ParseCIDR(s)
	return err == nil
}

// ipMatchesEntry checks whether the given IP matches a whitelist entry.
// It handles both exact IP matches and CIDR ranges.
func ipMatchesEntry(ip, entry string) (bool, error) {
	if entry == ip {
		return true, nil
	}
	if !strings.Contains(entry, "/") {
		return false, nil
	}
	_, ipnet, err := net.ParseCIDR(entry)
	if err != nil {
		return false, err
	}
	return ipnet.Contains(net.ParseIP(ip)), nil
}

// AddPermanentEntry inserts a permanent whitelist entry.
// Returns true if the entry was added, false if it already existed.
func AddPermanentEntry(db *sql.DB, ip, apiClientIP string) (bool, error) {
	_, err := db.Exec(`INSERT INTO perm_whitelist(ip) VALUES(?)`, ip)
	if err != nil {
		if isUniqueConstraintError(err) {
			_, logErr := db.Exec(`INSERT INTO audit_log(action, ip, api_client_ip) VALUES('add_perm_duplicate', ?, ?)`, ip, apiClientIP)
			if logErr != nil {
				return false, logErr
			}
			return false, nil
		}
		return false, err
	}
	_, err = db.Exec(`INSERT INTO audit_log(action, ip, api_client_ip) VALUES('add_perm', ?, ?)`, ip, apiClientIP)
	if err != nil {
		return true, err
	}
	return true, nil
}

// DeletePermanentEntry removes a permanent whitelist entry.
func DeletePermanentEntry(db *sql.DB, ip, apiClientIP string) error {
	result, err := db.Exec(`DELETE FROM perm_whitelist WHERE ip = ?`, ip)
	if err != nil {
		return err
	}
	if rows, _ := result.RowsAffected(); rows == 0 {
		slog.Warn("tried to delete non-existent perm ip", "ip", ip)
		return nil
	}
	_, err = db.Exec(`INSERT INTO audit_log(action, ip, api_client_ip) VALUES('delete_perm', ?, ?)`, ip, apiClientIP)
	if err != nil {
		return err
	}
	return nil
}

// IsPermanentWhitelisted checks whether the given IP is in the permanent whitelist.
// It iterates all entries to support CIDR matching.
func IsPermanentWhitelisted(db *sql.DB, ip string) (bool, error) {
	rows, err := db.Query(`SELECT ip FROM perm_whitelist`)
	if err != nil {
		return false, err
	}
	defer rows.Close()

	for rows.Next() {
		var entry string
		if err := rows.Scan(&entry); err != nil {
			continue
		}
		if matches, _ := ipMatchesEntry(ip, entry); matches {
			return true, nil
		}
	}
	return false, nil
}

// AddTempEntry inserts or updates a temporary whitelist entry (upsert).
func AddTempEntry(db *sql.DB, ip string, ttlSeconds int, reason, apiClientIP string) error {
	expiresAt := time.Now().Add(time.Duration(ttlSeconds) * time.Second).Format(time.RFC3339)
	_, err := db.Exec(`
		INSERT INTO temp_whitelist(ip, expires_at, reason)
		VALUES(?, ?, ?)
		ON CONFLICT(ip) DO UPDATE SET expires_at=excluded.expires_at, reason=excluded.reason
	`, ip, expiresAt, reason)
	if err != nil {
		return err
	}
	_, err = db.Exec(`INSERT INTO audit_log(action, ip, ttl_seconds, reason, api_client_ip) VALUES('add_temp', ?, ?, ?, ?)`, ip, ttlSeconds, reason, apiClientIP)
	if err != nil {
		return err
	}
	return nil
}

// DeleteTempEntry removes a temporary whitelist entry.
func DeleteTempEntry(db *sql.DB, ip, apiClientIP string) error {
	result, err := db.Exec(`DELETE FROM temp_whitelist WHERE ip = ?`, ip)
	if err != nil {
		return err
	}
	if rows, _ := result.RowsAffected(); rows == 0 {
		slog.Warn("tried to delete non-existent temp ip", "ip", ip)
		return nil
	}
	_, err = db.Exec(`INSERT INTO audit_log(action, ip, api_client_ip) VALUES('delete_temp', ?, ?)`, ip, apiClientIP)
	if err != nil {
		return err
	}
	return nil
}

// IsTempWhitelisted checks whether the given IP has a valid temporary whitelist entry.
func IsTempWhitelisted(db *sql.DB, ip string) (bool, error) {
	var expiresAt string
	err := db.QueryRow(`SELECT expires_at FROM temp_whitelist WHERE ip = ?`, ip).Scan(&expiresAt)
	if err == sql.ErrNoRows {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	expires, parseErr := time.Parse(time.RFC3339, expiresAt)
	if parseErr != nil {
		slog.Error("failed to parse expires_at", "ip", ip, "error", parseErr)
		return false, nil
	}
	return expires.After(time.Now()), nil
}

// CleanupExpiredTemp removes expired temporary whitelist entries and logs them.
func CleanupExpiredTemp(db *sql.DB) error {
	rows, err := db.Query(`SELECT ip FROM temp_whitelist WHERE expires_at <= ?`, time.Now().Format(time.RFC3339))
	if err != nil {
		return err
	}

	var expiredIPs []string
	for rows.Next() {
		var ip string
		if err := rows.Scan(&ip); err == nil {
			expiredIPs = append(expiredIPs, ip)
		}
	}
	if err := rows.Err(); err != nil {
		return err
	}

	_, err = db.Exec(`DELETE FROM temp_whitelist WHERE expires_at <= ?`, time.Now().Format(time.RFC3339))
	if err != nil {
		return err
	}

	for _, ip := range expiredIPs {
		_, logErr := db.Exec(`INSERT INTO audit_log(action, ip, reason) VALUES('expire_temp', ?, ?)`, ip, "expired")
		if logErr != nil {
			slog.Warn("failed to log expired temp ip", "ip", ip, "error", logErr)
		}
		slog.Info("expired temporary whitelist", "ip", ip)
	}
	return nil
}

// ListTempEntriesPaginated returns a page of valid temporary whitelist entries.
func ListTempEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) {
	rows, err := db.Query(`
		SELECT ip, expires_at, reason
		FROM temp_whitelist
		WHERE expires_at > ?
		ORDER BY expires_at
		LIMIT ? OFFSET ?
	`, time.Now().Format(time.RFC3339), limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanTempEntries(rows)
}

// ListPermEntriesPaginated returns a page of permanent whitelist entries.
func ListPermEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) {
	rows, err := db.Query(`SELECT ip FROM perm_whitelist ORDER BY created_at LIMIT ? OFFSET ?`, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanPermEntries(rows)
}

func scanTempEntries(rows *sql.Rows) ([]map[string]interface{}, error) {
	list := make([]map[string]interface{}, 0)
	for rows.Next() {
		var ip string
		var expiresAt string
		var reason sql.NullString
		if err := rows.Scan(&ip, &expiresAt, &reason); err != nil {
			continue
		}
		list = append(list, map[string]interface{}{
			"ip":      ip,
			"expires": expiresAt,
			"reason":  reason.String,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return list, nil
}

func scanPermEntries(rows *sql.Rows) ([]map[string]interface{}, error) {
	list := make([]map[string]interface{}, 0)
	for rows.Next() {
		var ip string
		if err := rows.Scan(&ip); err != nil {
			continue
		}
		list = append(list, map[string]interface{}{"ip": ip})
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return list, nil
}

// RotateAuditLog trims the audit log if it exceeds maxEntries entries.
func RotateAuditLog(db *sql.DB, maxEntries int) error {
	if maxEntries <= 0 {
		maxEntries = 200000
	}
	var count int
	if err := db.QueryRow(`SELECT COUNT(*) FROM audit_log`).Scan(&count); err != nil {
		return err
	}
	if count <= maxEntries {
		return nil
	}
	_, err := db.Exec(`
		DELETE FROM audit_log
		WHERE id IN (
			SELECT id FROM audit_log
			ORDER BY timestamp ASC
			LIMIT ?
		)
	`, count-maxEntries)
	return err
}