Wormhole/internal/cache/cache.go

package cache

import (
	"net"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/wormhole/wormhole/internal/store"
)

type OnlineAgent struct {
	ClientID    uint
	Version     string
	ConnectedAt time.Time
	LastPing    time.Time
}

type CompiledRule struct {
	Rule    store.IPRule
	CIDR    *net.IPNet
	Regex   *regexp.Regexp
	ExactIP net.IP
}

type FlowDelta struct {
	Inlet  int64
	Export int64
}

type Manager struct {
	mu sync.RWMutex

	clients      map[uint]*store.Client
	clientsByKey map[string]*store.Client
	tunnels      map[uint]*store.Tunnel
	tunnelsByPort map[string]*store.Tunnel
	hosts        map[uint]*store.Host
	hostIndex    []hostEntry

	onlineAgents map[uint]*OnlineAgent

	ipRules      []CompiledRule
	authPolicies map[string]*store.AuthPolicy
	visitorBindings map[string][]uint

	flowCounters map[string]*FlowDelta
}

type hostEntry struct {
	host *store.Host
}

func NewManager() *Manager {
	return &Manager{
		clients:       make(map[uint]*store.Client),
		clientsByKey:  make(map[string]*store.Client),
		tunnels:       make(map[uint]*store.Tunnel),
		tunnelsByPort: make(map[string]*store.Tunnel),
		hosts:         make(map[uint]*store.Host),
		onlineAgents:  make(map[uint]*OnlineAgent),
		authPolicies:  make(map[string]*store.AuthPolicy),
		visitorBindings: make(map[string][]uint),
		flowCounters:  make(map[string]*FlowDelta),
	}
}

func (m *Manager) ReloadAll(st *store.Store) error {
	clients, err := st.ListClients(0, true)
	if err != nil {
		return err
	}
	tunnels, err := st.ListTunnels(0, 0, true)
	if err != nil {
		return err
	}
	hosts, err := st.ListHosts(0, 0, true)
	if err != nil {
		return err
	}
	rules, err := st.ListIPRules("", 0)
	if err != nil {
		return err
	}
	policies, err := st.ListAuthPolicies("", 0)
	if err != nil {
		return err
	}
	visitors, err := st.ListAllResourceVisitors()
	if err != nil {
		return err
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	m.clients = make(map[uint]*store.Client)
	m.clientsByKey = make(map[string]*store.Client)
	for i := range clients {
		c := clients[i]
		m.clients[c.ID] = &c
		m.clientsByKey[c.VerifyKey] = &c
	}

	m.tunnels = make(map[uint]*store.Tunnel)
	m.tunnelsByPort = make(map[string]*store.Tunnel)
	for i := range tunnels {
		t := tunnels[i]
		m.tunnels[t.ID] = &t
		key := tunnelPortKey(t.ListenIP, t.ListenPort)
		m.tunnelsByPort[key] = &t
	}

	m.hosts = make(map[uint]*store.Host)
	m.hostIndex = nil
	for i := range hosts {
		h := hosts[i]
		m.hosts[h.ID] = &h
		m.hostIndex = append(m.hostIndex, hostEntry{host: &h})
	}

	m.ipRules = compileRules(rules)

	m.authPolicies = make(map[string]*store.AuthPolicy)
	for i := range policies {
		p := policies[i]
		key := authPolicyKey(p.ResourceType, p.ResourceID)
		m.authPolicies[key] = &p
	}

	m.visitorBindings = make(map[string][]uint)
	for i := range visitors {
		v := visitors[i]
		key := authPolicyKey(v.ResourceType, v.ResourceID)
		m.visitorBindings[key] = append(m.visitorBindings[key], v.UserID)
	}
	return nil
}

func tunnelPortKey(ip string, port int) string {
	if ip == "" {
		ip = "0.0.0.0"
	}
	return ip + ":" + strconv.Itoa(port)
}

func authPolicyKey(resourceType string, resourceID uint) string {
	return resourceType + ":" + itoa(resourceID)
}

func itoa(n uint) string {
	if n == 0 {
		return "0"
	}
	var b [20]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}

func compileRules(rules []store.IPRule) []CompiledRule {
	out := make([]CompiledRule, 0, len(rules))
	for _, r := range rules {
		cr := CompiledRule{Rule: r}
		switch r.PatternKind {
		case "cidr":
			_, network, err := net.ParseCIDR(r.Pattern)
			if err == nil {
				cr.CIDR = network
			}
		case "regex":
			re, err := regexp.Compile(r.Pattern)
			if err == nil {
				cr.Regex = re
			}
		default:
			cr.ExactIP = net.ParseIP(r.Pattern)
		}
		out = append(out, cr)
	}
	return out
}

func (m *Manager) SetClientOnline(clientID uint, version string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	now := time.Now()
	m.onlineAgents[clientID] = &OnlineAgent{
		ClientID:    clientID,
		Version:     version,
		ConnectedAt: now,
		LastPing:    now,
	}
}

func (m *Manager) SetClientOffline(clientID uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.onlineAgents, clientID)
}

func (m *Manager) TouchAgent(clientID uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if a, ok := m.onlineAgents[clientID]; ok {
		a.LastPing = time.Now()
	}
}

func (m *Manager) AgentLastPing(clientID uint) (time.Time, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	a, ok := m.onlineAgents[clientID]
	if !ok {
		return time.Time{}, false
	}
	return a.LastPing, true
}

func (m *Manager) IsOnline(clientID uint) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	_, ok := m.onlineAgents[clientID]
	return ok
}

func (m *Manager) OnlineCount() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.onlineAgents)
}

func (m *Manager) GetClientByKey(key string) (*store.Client, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	c, ok := m.clientsByKey[key]
	return c, ok
}

func (m *Manager) GetClient(id uint) (*store.Client, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	c, ok := m.clients[id]
	return c, ok
}

func (m *Manager) GetTunnel(id uint) (*store.Tunnel, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	t, ok := m.tunnels[id]
	return t, ok
}

func (m *Manager) GetHost(id uint) (*store.Host, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	h, ok := m.hosts[id]
	return h, ok
}

func (m *Manager) MatchHost(hostHeader, path, scheme string) *store.Host {
	m.mu.RLock()
	defer m.mu.RUnlock()

	hostHeader = strings.ToLower(strings.Split(hostHeader, ":")[0])
	var best *store.Host
	bestLen := -1

	for _, entry := range m.hostIndex {
		h := entry.host
		if !h.Status {
			continue
		}
		if h.Scheme != "all" && h.Scheme != scheme {
			continue
		}
		if !matchHostPattern(h.Host, hostHeader) {
			continue
		}
		loc := h.Location
		if loc == "" {
			loc = "/"
		}
		if !strings.HasPrefix(path, loc) {
			continue
		}
		if len(loc) > bestLen {
			best = h
			bestLen = len(loc)
		}
	}
	return best
}

func matchHostPattern(pattern, host string) bool {
	pattern = strings.ToLower(pattern)
	if pattern == host {
		return true
	}
	if strings.HasPrefix(pattern, "*.") {
		suffix := pattern[1:]
		return strings.HasSuffix(host, suffix) || host == pattern[2:]
	}
	return false
}

func (m *Manager) ReloadIPRules(st *store.Store) error {
	rules, err := st.ListIPRules("", 0)
	if err != nil {
		return err
	}
	m.mu.Lock()
	m.ipRules = compileRules(rules)
	m.mu.Unlock()
	return nil
}

func (m *Manager) AllowIP(ipStr string, scopes []Scope) bool {
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return false
	}
	m.mu.RLock()
	defer m.mu.RUnlock()

	var denyMatched, allowMatched bool
	hasAllowRules := false

	for _, cr := range m.ipRules {
		if !scopeMatch(cr.Rule, scopes) {
			continue
		}
		if !ruleMatch(cr, ip, ipStr) {
			continue
		}
		if cr.Rule.Type == "deny" {
			denyMatched = true
		}
		if cr.Rule.Type == "allow" {
			allowMatched = true
			hasAllowRules = true
		}
	}

	if denyMatched {
		return false
	}
	if hasAllowRules {
		return allowMatched
	}
	return true
}

type Scope struct {
	Kind       string
	ResourceID uint
}

func scopeMatch(rule store.IPRule, scopes []Scope) bool {
	if rule.Scope == "global" {
		return true
	}
	for _, s := range scopes {
		if rule.Scope == s.Kind && rule.ResourceID == s.ResourceID {
			return true
		}
	}
	return false
}

func ruleMatch(cr CompiledRule, ip net.IP, ipStr string) bool {
	switch cr.Rule.PatternKind {
	case "cidr":
		return cr.CIDR != nil && cr.CIDR.Contains(ip)
	case "regex":
		return cr.Regex != nil && cr.Regex.MatchString(ipStr)
	default:
		return cr.ExactIP != nil && cr.ExactIP.Equal(ip)
	}
}

func (m *Manager) GetAuthPolicy(resourceType string, resourceID uint) (*store.AuthPolicy, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	p, ok := m.authPolicies[authPolicyKey(resourceType, resourceID)]
	return p, ok
}

func (m *Manager) GetResourceVisitorIDs(resourceType string, resourceID uint) ([]uint, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	ids, ok := m.visitorBindings[authPolicyKey(resourceType, resourceID)]
	if !ok {
		return nil, false
	}
	out := make([]uint, len(ids))
	copy(out, ids)
	return out, true
}

func (m *Manager) SetResourceVisitors(resourceType string, resourceID uint, userIDs []uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	key := authPolicyKey(resourceType, resourceID)
	if len(userIDs) == 0 {
		delete(m.visitorBindings, key)
		return
	}
	copied := make([]uint, len(userIDs))
	copy(copied, userIDs)
	m.visitorBindings[key] = copied
}

func (m *Manager) DeleteResourceVisitors(resourceType string, resourceID uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.visitorBindings, authPolicyKey(resourceType, resourceID))
}

func (m *Manager) AddFlow(resourceType string, resourceID uint, inlet, export int64) {
	key := resourceType + ":" + itoa(resourceID)
	m.mu.Lock()
	defer m.mu.Unlock()
	delta, ok := m.flowCounters[key]
	if !ok {
		delta = &FlowDelta{}
		m.flowCounters[key] = delta
	}
	atomic.AddInt64(&delta.Inlet, inlet)
	atomic.AddInt64(&delta.Export, export)
}

func (m *Manager) DrainFlow() map[string]FlowDelta {
	m.mu.Lock()
	defer m.mu.Unlock()
	out := make(map[string]FlowDelta, len(m.flowCounters))
	for k, v := range m.flowCounters {
		inlet := atomic.SwapInt64(&v.Inlet, 0)
		export := atomic.SwapInt64(&v.Export, 0)
		if inlet > 0 || export > 0 {
			out[k] = FlowDelta{Inlet: inlet, Export: export}
		}
	}
	return out
}

func (m *Manager) UpsertClient(c *store.Client) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.clients[c.ID] = c
	m.clientsByKey[c.VerifyKey] = c
}

func (m *Manager) DeleteClient(id uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if c, ok := m.clients[id]; ok {
		delete(m.clientsByKey, c.VerifyKey)
	}
	delete(m.clients, id)
	delete(m.onlineAgents, id)
}

func (m *Manager) UpsertTunnel(t *store.Tunnel) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.tunnels[t.ID] = t
	m.tunnelsByPort[tunnelPortKey(t.ListenIP, t.ListenPort)] = t
}

func (m *Manager) DeleteTunnel(id uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if t, ok := m.tunnels[id]; ok {
		delete(m.tunnelsByPort, tunnelPortKey(t.ListenIP, t.ListenPort))
	}
	delete(m.tunnels, id)
}

func (m *Manager) UpsertHost(h *store.Host) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.hosts[h.ID] = h
	m.hostIndex = nil
	for _, host := range m.hosts {
		m.hostIndex = append(m.hostIndex, hostEntry{host: host})
	}
}

func (m *Manager) DeleteHost(id uint) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.hosts, id)
	m.hostIndex = nil
	for _, host := range m.hosts {
		m.hostIndex = append(m.hostIndex, hostEntry{host: host})
	}
}

func (m *Manager) ListEnabledTunnels() []*store.Tunnel {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var out []*store.Tunnel
	for _, t := range m.tunnels {
		if t.Status {
			out = append(out, t)
		}
	}
	return out
}

func (m *Manager) SetTunnelRunStatus(id uint, running bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if t, ok := m.tunnels[id]; ok {
		t.RunStatus = running
	}
}