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 } }