pekoe/internal/pubsub/hub.go

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