// @tag pubsub // @tag core package pubsub import ( "fmt" "regexp" "strings" "sync" ) // Hub maintains the mapping between channels and their subscribers. type Hub struct { mu sync.RWMutex subs map[string]map[*Client]struct{} patterns map[string]*patternEntry // Client registry for CLIENT LIST/INFO/KILL commands clients map[uint64]*Client // Rate limiting patternCount int // Total active pattern subscriptions channelCount int // Total active channel subscriptions } // MaxPatternSubscriptions limits pattern subscriptions to prevent DoS const MaxPatternSubscriptions = 10000 // MaxChannelSubscriptions limits channel subscriptions per client to prevent DoS const MaxChannelSubscriptions = 10000 // patternEntry holds compiled regex and subscribers for a pattern. type patternEntry struct { regex *regexp.Regexp clients map[*Client]struct{} } // patternCache caches compiled regex patterns to avoid recompilation. type patternCacheType struct { sync.Map } // Delete removes a pattern from the cache func (pc *patternCacheType) Delete(pattern string) { pc.Map.Delete(pattern) } var patternCache = patternCacheType{} // NewHub creates a new pubsub hub instance ready for use. func NewHub() *Hub { return &Hub{ subs: make(map[string]map[*Client]struct{}), patterns: make(map[string]*patternEntry), clients: make(map[uint64]*Client), } } // RegisterClient adds a client to the registry func (h *Hub) RegisterClient(client *Client) { h.mu.Lock() defer h.mu.Unlock() h.clients[client.ID()] = client } // UnregisterClient removes a client from the registry func (h *Hub) UnregisterClient(client *Client) { h.mu.Lock() defer h.mu.Unlock() delete(h.clients, client.ID()) } // GetClient returns a client by ID, or nil if not found func (h *Hub) GetClient(id uint64) *Client { h.mu.RLock() defer h.mu.RUnlock() return h.clients[id] } // GetAllClients returns a slice of all registered clients func (h *Hub) GetAllClients() []*Client { h.mu.RLock() defer h.mu.RUnlock() clients := make([]*Client, 0, len(h.clients)) for _, c := range h.clients { clients = append(clients, c) } return clients } // ClientCount returns the number of connected clients func (h *Hub) ClientCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.clients) } // Subscribe adds a client to a channel. func (h *Hub) Subscribe(client *Client, channel string) bool { if channel == "" { return false } // Must hold hub mutex when modifying the global subs map h.mu.Lock() defer h.mu.Unlock() // Check per-client subscription limit if client.SubsCount() >= MaxChannelSubscriptions { return false } clients, exists := h.subs[channel] if !exists { clients = make(map[*Client]struct{}) h.subs[channel] = clients } if _, already := clients[client]; already { return false } clients[client] = struct{}{} h.channelCount++ client.AddSubscription(channel) return true } // Unsubscribe removes a client from a channel. func (h *Hub) Unsubscribe(client *Client, channel string) bool { if channel == "" { return false } // Must hold hub mutex when modifying the global subs map h.mu.Lock() defer h.mu.Unlock() clients, exists := h.subs[channel] if !exists { return false } if _, ok := clients[client]; !ok { return false } delete(clients, client) if len(clients) == 0 { delete(h.subs, channel) } h.channelCount-- client.RemoveSubscription(channel) return true } // UnsubscribeAll removes a client from all subscribed channels and patterns. func (h *Hub) UnsubscribeAll(client *Client) { channels := client.GetChannels() for _, channel := range channels { h.Unsubscribe(client, channel) } patterns := client.GetPatterns() for _, pattern := range patterns { h.PatternUnsubscribe(client, pattern) } } // Publish sends a message to all subscribers of a channel. func (h *Hub) Publish(channel string, message []byte) int { // Copy clients under lock, then release lock before sending h.mu.RLock() clients, exists := h.subs[channel] if !exists || len(clients) == 0 { h.mu.RUnlock() return 0 } clientList := make([]*Client, 0, len(clients)) for client := range clients { clientList = append(clientList, client) } h.mu.RUnlock() count := 0 for _, client := range clientList { if err := client.sendMessage(channel, message); err != nil { h.removeClientFromChannel(client, channel) continue } count++ } return count } // PublishToPattern sends a message to all clients subscribed to patterns matching the channel. // nolint: cyclop func (h *Hub) PublishToPattern(channel string, message []byte) int { h.mu.RLock() // Collect unique patterns that match the channel using a map for O(1) deduplication matchedPatterns := make(map[string]bool) for pattern, entry := range h.patterns { if entry.regex.MatchString(channel) { matchedPatterns[pattern] = true } } if len(matchedPatterns) == 0 { h.mu.RUnlock() return 0 } type match struct { client *Client pattern string } var matches []match for pattern := range matchedPatterns { if entry, ok := h.patterns[pattern]; ok { for client := range entry.clients { matches = append(matches, match{client: client, pattern: pattern}) } } } h.mu.RUnlock() count := 0 for _, m := range matches { if err := m.client.sendPMessage(m.pattern, channel, message); err != nil { h.removeClientFromPattern(m.client, m.pattern) continue } count++ } return count } // PatternSubscribe adds a client to a pattern subscription. func (h *Hub) PatternSubscribe(client *Client, pattern string) (bool, error) { regex, err := getOrCompilePattern(pattern) if err != nil { return false, fmt.Errorf("invalid pattern: %v", err) } h.mu.Lock() defer h.mu.Unlock() entry, exists := h.patterns[pattern] if !exists { entry = &patternEntry{ regex: regex, clients: make(map[*Client]struct{}), } h.patterns[pattern] = entry } if _, already := entry.clients[client]; already { return false, nil } // Check pattern subscription limit to prevent DoS if h.patternCount >= MaxPatternSubscriptions { return false, fmt.Errorf("too many pattern subscriptions") } entry.clients[client] = struct{}{} h.patternCount++ client.AddPatternSubscription(pattern) return true, nil } // getOrCompilePattern returns a cached compiled regex or compiles and caches a new one. func getOrCompilePattern(pattern string) (*regexp.Regexp, error) { if cached, ok := patternCache.Load(pattern); ok { return cached.(*regexp.Regexp), nil } regex, err := compilePattern(pattern) if err != nil { return nil, err } patternCache.Store(pattern, regex) return regex, nil } func compilePattern(pattern string) (*regexp.Regexp, error) { var buf strings.Builder for i := 0; i < len(pattern); i++ { switch pattern[i] { case '*': buf.WriteString(".*") case '?': buf.WriteByte('.') case '\\': if i+1 < len(pattern) { buf.WriteByte('\\') buf.WriteByte(pattern[i+1]) i++ } case '.', '+', '(', ')', '[', ']', '{', '}', '^', '$', '|': buf.WriteByte('\\') buf.WriteByte(pattern[i]) default: buf.WriteByte(pattern[i]) } } return regexp.Compile("^" + buf.String() + "$") } // PatternUnsubscribe removes a client from a pattern subscription. func (h *Hub) PatternUnsubscribe(client *Client, pattern string) bool { h.mu.Lock() defer h.mu.Unlock() entry, exists := h.patterns[pattern] if !exists { return false } if _, ok := entry.clients[client]; !ok { return false } delete(entry.clients, client) if len(entry.clients) == 0 { delete(h.patterns, pattern) h.patternCount-- // Remove pattern from cache patternCache.Delete(pattern) } client.RemovePatternSubscription(pattern) return true } // PatternCount returns the total number of pattern subscriptions. func (h *Hub) PatternCount() int { h.mu.RLock() defer h.mu.RUnlock() count := 0 for _, entry := range h.patterns { count += len(entry.clients) } return count } // RemoveClientFromChannel handles cleanup when a client's connection fails. func (h *Hub) removeClientFromChannel(client *Client, channel string) { h.mu.Lock() defer h.mu.Unlock() if clients, ok := h.subs[channel]; ok { if _, exists := clients[client]; exists { delete(clients, client) if len(clients) == 0 { delete(h.subs, channel) } } } client.RemoveSubscription(channel) } // RemoveClientFromPattern handles cleanup when a client's connection fails for patterns. func (h *Hub) removeClientFromPattern(client *Client, pattern string) { h.mu.Lock() defer h.mu.Unlock() if entry, ok := h.patterns[pattern]; ok { if _, exists := entry.clients[client]; exists { delete(entry.clients, client) if len(entry.clients) == 0 { delete(h.patterns, pattern) } } } client.RemovePatternSubscription(pattern) } // ChannelCount returns the number of subscribers to a specific channel. func (h *Hub) ChannelCount(channel string) int { h.mu.RLock() defer h.mu.RUnlock() if clients, ok := h.subs[channel]; ok { return len(clients) } return 0 } // TotalChannels returns the total number of active channels. func (h *Hub) TotalChannels() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.subs) } // ChannelList returns a list of all active channel names, optionally filtered by pattern. func (h *Hub) ChannelList(pattern string) []string { h.mu.RLock() defer h.mu.RUnlock() if pattern == "" { channels := make([]string, 0, len(h.subs)) for ch := range h.subs { channels = append(channels, ch) } return channels } // Compile pattern once outside the loop (using cache) regex, err := getOrCompilePattern(pattern) if err != nil { return nil } channels := make([]string, 0, len(h.subs)) for ch := range h.subs { if regex.MatchString(ch) { channels = append(channels, ch) } } return channels }