vibeproxy/internal/runtime/store.go

// Package runtime holds the active, in-memory configuration snapshot that
// the hot proxy/stream paths read. SQLite is only touched when the
// configuration changes; every request/connection reads this snapshot.
package runtime

import (
	"context"
	"crypto/tls"
	"database/sql"
	"net"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/owen/vibeproxy/internal/access"
	"github.com/owen/vibeproxy/internal/database"
	"github.com/owen/vibeproxy/internal/health"
	"github.com/owen/vibeproxy/internal/loadbalancer"
	"github.com/owen/vibeproxy/internal/models"
)

// ProxyHostRuntime is one HTTP(S) proxy or redirect host, resolved and
// ready to serve traffic.
type ProxyHostRuntime struct {
	Model      models.ProxyHost
	Pool       *loadbalancer.Pool
	AccessList *access.List // nil if none attached
}

// StreamRuntime is one TCP/UDP stream, ready to serve traffic.
type StreamRuntime struct {
	Model models.Stream
}

// Snapshot is an immutable view of the whole active configuration.
type Snapshot struct {
	HostsByDomain  map[string]*ProxyHostRuntime // lowercase domain -> host
	AccessLists    map[int64]*access.List
	Settings       models.Settings
	TrustedProxies []*net.IPNet
}

func (s *Snapshot) LookupHost(hostHeader string) *ProxyHostRuntime {
	name := hostHeader
	if h, _, err := net.SplitHostPort(hostHeader); err == nil {
		name = h
	}
	name = strings.ToLower(strings.TrimSuffix(name, "."))
	return s.HostsByDomain[name]
}

func (s *Snapshot) IsTrustedProxy(ip net.IP) bool {
	for _, n := range s.TrustedProxies {
		if n.Contains(ip) {
			return true
		}
	}
	return false
}

// Store owns the current Snapshot (read via atomic pointer for lock-free,
// allocation-free reads on the hot path) and knows how to rebuild it from
// the database.
type Store struct {
	db      *sql.DB
	current atomic.Pointer[Snapshot]

	checker   *health.Checker
	watchCtx  context.Context
	watchStop context.CancelFunc
	watchMu   sync.Mutex

	onHealthChange func(hostID int64, backendAddr string, healthy bool)
}

func NewStore(db *sql.DB) *Store {
	s := &Store{db: db}
	s.checker = health.NewChecker(s.handleHealthChange)
	s.current.Store(&Snapshot{HostsByDomain: map[string]*ProxyHostRuntime{}, AccessLists: map[int64]*access.List{}})
	return s
}

func (s *Store) Current() *Snapshot { return s.current.Load() }

func (s *Store) SetHealthChangeCallback(f func(hostID int64, backendAddr string, healthy bool)) {
	s.onHealthChange = f
}

func (s *Store) handleHealthChange(hostID int64, b *loadbalancer.Backend, healthy bool) {
	if s.onHealthChange != nil {
		s.onHealthChange(hostID, net.JoinHostPort(b.Model.Host, strconv.Itoa(b.Model.Port)), healthy)
	}
}

// Reload rebuilds the snapshot from the database and atomically swaps it
// in. Existing in-flight requests keep using the previous snapshot object
// until they finish; new requests immediately see the new one.
func (s *Store) Reload(ctx context.Context) error {
	hosts, err := database.ListProxyHosts(s.db)
	if err != nil {
		return err
	}
	accessLists, err := database.ListAccessLists(s.db)
	if err != nil {
		return err
	}
	settingsMap, err := database.AllSettings(s.db)
	if err != nil {
		return err
	}
	settings := settingsFromMap(settingsMap)

	compiledACLs := map[int64]*access.List{}
	for _, a := range accessLists {
		compiledACLs[a.ID] = access.Compile(a)
	}

	var trusted []*net.IPNet
	for _, c := range settings.TrustedProxyCIDRs {
		if _, n, err := net.ParseCIDR(c); err == nil {
			trusted = append(trusted, n)
		}
	}

	snap := &Snapshot{
		HostsByDomain:  map[string]*ProxyHostRuntime{},
		AccessLists:    compiledACLs,
		Settings:       settings,
		TrustedProxies: trusted,
	}

	// Stop the previous health-check watchers before starting new ones so
	// we don't leak goroutines across reloads.
	s.watchMu.Lock()
	if s.watchStop != nil {
		s.watchStop()
	}
	watchCtx, cancel := context.WithCancel(context.Background())
	s.watchCtx, s.watchStop = watchCtx, cancel
	s.watchMu.Unlock()

	for _, h := range hosts {
		if !h.Enabled {
			continue
		}
		var pool *loadbalancer.Pool
		if h.Type == models.ProxyHostTypeProxy {
			backends := make([]*loadbalancer.Backend, 0, len(h.Backends))
			for _, bm := range h.Backends {
				if !bm.Enabled {
					continue
				}
				b := loadbalancer.NewBackend(bm)
				backends = append(backends, b)
				s.checker.Watch(watchCtx, h.ID, b)
			}
			pool = loadbalancer.NewPool(h.LoadBalancingMethod, backends)
		}
		var acl *access.List
		if h.AccessListID != nil {
			acl = compiledACLs[*h.AccessListID]
		}
		rt := &ProxyHostRuntime{Model: h, Pool: pool, AccessList: acl}
		for _, d := range h.Domains {
			snap.HostsByDomain[strings.ToLower(d)] = rt
		}
	}

	s.current.Store(snap)
	return nil
}

func settingsFromMap(m map[string]string) models.Settings {
	s := models.Settings{
		AdminUsername:         get(m, "admin_username", "admin"),
		ACMEEmail:             get(m, "acme_email", ""),
		ACMEStaging:           get(m, "acme_staging", "false") == "true",
		ACMERenewalWindowDays: atoi(get(m, "acme_renewal_window_days", "30"), 30),
		CertStoragePath:       get(m, "cert_storage_path", "./data/certs"),
		LogRetentionDays:      atoi(get(m, "log_retention_days", "30"), 30),
		BackupEnabled:         get(m, "backup_enabled", "true") == "true",
		BackupIntervalHours:   atoi(get(m, "backup_interval_hours", "24"), 24),
		BackupRetentionCount:  atoi(get(m, "backup_retention_count", "7"), 7),
		BackupPath:            get(m, "backup_path", "./data/backups"),
	}
	if tp := get(m, "trusted_proxy_cidrs", ""); tp != "" {
		s.TrustedProxyCIDRs = strings.Split(tp, ",")
	}
	return s
}

func get(m map[string]string, k, def string) string {
	if v, ok := m[k]; ok && v != "" {
		return v
	}
	return def
}

func atoi(s string, def int) int {
	n, err := strconv.Atoi(s)
	if err != nil {
		return def
	}
	return n
}

// TLSCertLookup is implemented by the certificate manager and used for
// SNI-based certificate selection in the HTTPS listener.
type TLSCertLookup interface {
	GetCertificateForDomain(domain string) (*tls.Certificate, error)
}