auth-proxy/db.go

package main

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

	sqlite "modernc.org/sqlite"
)

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

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

	schema := `
	CREATE TABLE IF NOT EXISTS perm_whitelist (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		entry 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
	);
	`
	_, err = db.Exec(schema)
	if err != nil {
		return nil, err
	}

	return db, nil
}

func AddPermanentEntry(db *sql.DB, entry, apiClientIP string) (bool, error) {
	_, err := db.Exec(`INSERT INTO perm_whitelist(entry) VALUES(?)`, entry)
	if err != nil {
		if isUniqueConstraintError(err) {
			_, err = db.Exec(`
				INSERT INTO audit_log(action, ip, api_client_ip)
				VALUES('add_perm_duplicate', ?, ?)
			`, entry, apiClientIP)
			if err != nil {
				return false, err
			}
			return false, nil
		}
		return false, err
	}
	_, err = db.Exec(`
		INSERT INTO audit_log(action, ip, api_client_ip)
		VALUES('add_perm', ?, ?)
	`, entry, apiClientIP)
	if err != nil {
		return true, err
	}
	return true, nil
}

func DeletePermanentEntry(db *sql.DB, entry, apiClientIP string) error {
	_, err := db.Exec(`
		DELETE FROM perm_whitelist
		WHERE entry = ?
	`, entry)
	if err != nil {
		return err
	}
	_, err = db.Exec(`
		INSERT INTO audit_log(action, ip, api_client_ip)
		VALUES('delete_perm', ?, ?)
	`, entry, apiClientIP)
	if err != nil {
		return err
	}
	return nil
}

func IsPermanentWhitelisted(db *sql.DB, ip string) (bool, error) {
	rows, err := db.Query(`
		SELECT entry
		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
		}
		matches, err := ipMatchesEntry(ip, entry)
		if err != nil {
			continue
		}
		if matches {
			return true, nil
		}
	}
	return false, nil
}

func AddTempEntry(db *sql.DB, ip string, ttlSeconds int, reason, apiClientIP string) error {
	expiresAt := time.Now().Add(time.Duration(ttlSeconds) * time.Second)
	_, 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
}

func DeleteTempEntry(db *sql.DB, ip, apiClientIP string) error {
	_, err := db.Exec(`
		DELETE FROM temp_whitelist
		WHERE ip = ?
	`, ip)
	if err != nil {
		return err
	}
	_, err = db.Exec(`
		INSERT INTO audit_log(action, ip, api_client_ip)
		VALUES('delete_temp', ?, ?)
	`, ip, apiClientIP)
	if err != nil {
		return err
	}
	return nil
}

func IsTempWhitelisted(db *sql.DB, ip string) (bool, error) {
	var expiresAt time.Time
	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
	}
	return expiresAt.After(time.Now()), nil
}

func CleanupExpiredTemp(db *sql.DB) error {
	rows, err := db.Query(`
		SELECT ip
		FROM temp_whitelist
		WHERE expires_at <= ?
	`, time.Now())
	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)
		}
	}
	rows.Close()

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

	var reason string
	reason = "expired"

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

func ListTempEntries(db *sql.DB) ([]map[string]interface{}, error) {
	rows, err := db.Query(`
		SELECT ip, expires_at, reason
		FROM temp_whitelist
		WHERE expires_at > ?
		ORDER BY expires_at
	`, time.Now())
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var list []map[string]interface{} = make([]map[string]interface{}, 0)
	for rows.Next() {
		var ip string
		var expiresAt time.Time
		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,
		})
	}
	return list, nil
}

func ListPermEntries(db *sql.DB) ([]map[string]interface{}, error) {
	rows, err := db.Query(`
		SELECT entry
		FROM perm_whitelist
		ORDER BY created_at
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var list []map[string]interface{} = make([]map[string]interface{}, 0)
	for rows.Next() {
		var entry string
		if err := rows.Scan(&entry); err != nil {
			continue
		}
		list = append(list, map[string]interface{}{
			"entry": entry,
		})
	}
	return list, nil
}

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(), limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var list []map[string]interface{} = make([]map[string]interface{}, 0)
	for rows.Next() {
		var ip string
		var expiresAt time.Time
		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,
		})
	}
	return list, nil
}

func ListPermEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) {
	rows, err := db.Query(`
		SELECT entry
		FROM perm_whitelist
		ORDER BY created_at
		LIMIT ? OFFSET ?
	`, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var list []map[string]interface{} = make([]map[string]interface{}, 0)
	for rows.Next() {
		var entry string
		if err := rows.Scan(&entry); err != nil {
			continue
		}
		list = append(list, map[string]interface{}{
			"entry": entry,
		})
	}
	return list, nil
}

func ipMatchesEntry(ip, entry string) (bool, error) {
	if entry == ip {
		return true, nil
	}

	if len(entry) > 0 && entry[0] != '/' && containsSlash(entry) {
		_, ipnet, err := net.ParseCIDR(entry)
		if err != nil {
			return false, err
		}
		parsedIP := net.ParseIP(ip)
		if parsedIP == nil {
			return false, nil
		}
		return ipnet.Contains(parsedIP), nil
	}
	return false, nil
}

func containsSlash(s string) bool {
	for i := 0; i < len(s); i++ {
		if s[i] == '/' {
			return true
		}
	}
	return false
}

func isUniqueConstraintError(err error) bool {
	if err == nil {
		return false
	}
	var sqliteErr *sqlite.Error
	if errors.As(err, &sqliteErr) {
		return sqliteErr.Code() == 19
	}
	return false
}

func IsIP(s string) bool {
	return net.ParseIP(s) != nil
}

func IsCIDR(s string) bool {
	_, _, err := net.ParseCIDR(s)
	return err == nil
}

func RotateAuditLog(db *sql.DB, maxEntries int) error {
	if maxEntries <= 0 {
		maxEntries = 200000
	}
	var count int
	err := db.QueryRow(`SELECT COUNT(*) FROM audit_log`).Scan(&count)
	if err != nil {
		return err
	}
	if count <= maxEntries {
		return nil
	}
	rows, err := db.Query(`
		SELECT id FROM audit_log
		ORDER BY timestamp ASC
		LIMIT ?
	`, count-maxEntries)
	if err != nil {
		return err
	}
	var ids []int64
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err == nil {
			ids = append(ids, id)
		}
	}
	rows.Close()
	_, err = db.Exec(`
		DELETE FROM audit_log
		WHERE id IN (
			SELECT id FROM audit_log
			ORDER BY timestamp ASC
			LIMIT ?
		)
	`, count-maxEntries)
	return err
}