// @tag server // @tag handlers:pubsub package server import ( "fmt" "strings" "time" "github.com/eleksir/pekoe/internal/glob" "github.com/eleksir/pekoe/internal/logger" "github.com/eleksir/pekoe/internal/protocol" "github.com/eleksir/pekoe/internal/pubsub" ) // handleSubscribe processes SUBSCRIBE command. // SUBSCRIBE channel [channel ...] - Subscribe to one or more channels. // Returns immediately after subscription (doesn't block like PSUBSCRIBE). func (s *Server) handleSubscribe(client *pubsub.Client, channels [][]byte, writer *ConnWriter) error { if len(channels) == 0 { writer.WriteError("ERR wrong number of arguments for 'subscribe' command") return nil } // Subscribe to each channel specified for _, ch := range channels { channel := string(ch) if s.hub.Subscribe(client, channel) { logger.Debug("client=%s event=subscribe channel=%s", client.RemoteAddr(), channel) // Send confirmation in format: ["subscribe", channel, count] writer.WriteArray([]protocol.Value{ protocol.Bulk("subscribe"), protocol.Bulk(channel), protocol.Int(int64(client.SubsCount())), }) } } return nil } // handleUnsubscribe processes UNSUBSCRIBE command. // UNSUBSCRIBE [channel [channel ...]] - Unsubscribe from channels. // If no channels specified, unsubscribe from all. func (s *Server) handleUnsubscribe(client *pubsub.Client, channels [][]byte, writer *ConnWriter) error { if len(channels) == 0 { // Unsubscribe from all channels chans := client.GetChannels() if len(chans) == 0 { // No subscriptions - send single unsubscribe message with -1 (Redis behavior) writer.WriteArray([]protocol.Value{ protocol.Bulk("unsubscribe"), protocol.Bulk("-1"), protocol.Int(0), }) return nil } for _, ch := range chans { s.hub.Unsubscribe(client, ch) writer.WriteArray([]protocol.Value{ protocol.Bulk("unsubscribe"), protocol.Bulk(ch), protocol.Int(int64(client.SubsCount())), }) } } else { // Unsubscribe from specified channels - only send response if actually subscribed for _, ch := range channels { channel := string(ch) removed := s.hub.Unsubscribe(client, channel) if removed { logger.Debug("client=%s event=unsubscribe channel=%s", client.RemoteAddr(), channel) writer.WriteArray([]protocol.Value{ protocol.Bulk("unsubscribe"), protocol.Bulk(channel), protocol.Int(int64(client.SubsCount())), }) } } } return nil } // handlePublish processes PUBLISH command. // PUBLISH channel message - Publish a message to a channel. // Returns the number of subscribers who received the message. func (s *Server) handlePublish(args [][]byte, writer *ConnWriter) error { if len(args) != 2 { writer.WriteError("ERR wrong number of arguments for 'publish' command") return nil } channel := string(args[0]) message := args[1] regularCount := s.hub.Publish(channel, message) patternCount := s.hub.PublishToPattern(channel, message) totalCount := regularCount + patternCount logger.Debug("event=publish channel=%s regular=%d pattern=%d", channel, regularCount, patternCount) writer.WriteInt(int64(totalCount)) return nil } // handlePubsub processes PUBSUB subcommands for introspection. // nolint: cyclop func (s *Server) handlePubsub(args [][]byte, writer *ConnWriter) error { if len(args) == 0 { writer.WriteError("ERR wrong number of arguments for 'pubsub' command") return nil } subcmd := strings.ToLower(string(args[0])) switch subcmd { case "numsub": // PUBSUB NUMSUB [channel ...] - Get subscriber count for channels if len(args) < 2 { writer.WriteError("ERR wrong number of arguments for 'pubsub numsub' command") return nil } // Build response: [[channel, count], ...] channels := args[1:] result := make([]protocol.Value, 0, len(channels)*2) for _, ch := range channels { channel := string(ch) count := s.hub.ChannelCount(channel) result = append(result, protocol.Bulk(ch), protocol.Int(int64(count))) } writer.WriteArray(result) case "channels": // PUBSUB CHANNELS [pattern] - List active channels, optionally filtered var pattern string if len(args) > 1 { pattern = string(args[1]) } channels := s.channelList(pattern) result := make([]protocol.Value, 0, len(channels)) for _, ch := range channels { if pattern == "" || glob.Match(pattern, ch) { result = append(result, protocol.Bulk(ch)) } } writer.WriteArray(result) case "shardchannels": // PUBSUB SHARDCHANNELS [pattern] - List active sharded channels // Pekoe doesn't support sharded pubsub, return empty list writer.WriteArray(nil) case "shardnumsub": // PUBSUB SHARDNUMSUB [shardchannel ...] - Get subscriber count for sharded channels // Pekoe doesn't support sharded pubsub, return empty array writer.WriteArray(nil) case "numpat": // PUBSUB NUMPAT - Count of unique pattern subscriptions writer.WriteInt(int64(s.hub.PatternCount())) case "help": // PUBSUB HELP - Returns helpful text about pubsub subcommands helpText := []string{ "PUBSUB [ [ ...]]>", " PUBSUB CHANNELS [pattern] - List active channels (optionally filtered by pattern)", " PUBSUB SHARDCHANNELS [pattern] - List active sharded channels (not supported)", " PUBSUB NUMSUB [channel ...] - Get subscriber count for channels", " PUBSUB SHARDNUMSUB [shardchannel ...] - Get subscriber count for sharded channels (not supported)", " PUBSUB NUMPAT - Get the number of pattern subscriptions", " PUBSUB HELP - Show this help message", } result := make([]protocol.Value, len(helpText)) for i, line := range helpText { result[i] = protocol.Bulk(line) } writer.WriteArray(result) default: writer.WriteError(fmt.Sprintf("ERR unknown pubsub subcommand '%s'", subcmd)) } return nil } // channelList returns all active channel names, optionally filtered by pattern. func (s *Server) channelList(pattern string) []string { channels := s.hub.ChannelList(pattern) if channels == nil { return []string{} } return channels } // handlePing processes PING command. // PING [message] - Returns PONG or the optional message. // Redis returns simple string for message, not bulk string. func (s *Server) handlePing(args [][]byte, writer *ConnWriter) error { if len(args) > 0 { writer.WriteStatus(string(args[0])) } else { writer.WriteStatus("PONG") } return nil } // handleQuit processes QUIT command. // QUIT - Close the connection gracefully. func (s *Server) handleQuit(client *pubsub.Client, writer *ConnWriter) error { writer.WriteStatus("OK") return ErrQuit } // handleClient processes CLIENT command. // CLIENT subcommand [argument [argument ...]] // nolint: cyclop func (s *Server) handleClient(client *pubsub.Client, args [][]byte, writer *ConnWriter) error { if len(args) == 0 { writer.WriteError("ERR wrong number of arguments for 'client' command") return nil } subcmd := strings.ToLower(string(args[0])) switch subcmd { case "id": // CLIENT ID - return the connection ID //nolint: gosec // client ID is sequential, overflow is theoretical writer.WriteInt(int64(client.ID())) case "setname": // CLIENT SETNAME connection-name if len(args) < 2 { writer.WriteError("ERR wrong number of arguments for 'client setname' command") return nil } name := string(args[1]) // Validate name - must be printable characters and no spaces if !isValidClientName(name) { writer.WriteError("ERR Client names cannot contain spaces, null characters or names longer than 255 characters") return nil } client.SetName(name) writer.WriteStatus("OK") case "getname": // CLIENT GETNAME - return the connection name name := client.GetName() if name == "" { writer.WriteBulk(nil) } else { writer.WriteBulk([]byte(name)) } case "list": // CLIENT LIST [TYPE normal|master|slave|pubsub] [ID min-id max-id] // For pekoe, we only support pubsub clients clientList := s.buildClientList(args[1:]) writer.WriteBulk([]byte(clientList)) case "info": // CLIENT INFO - return information about the client writer.WriteBulk([]byte(client.InfoString())) case "kill": // CLIENT KILL [ID min-id max-id] [ADDR ip:port] [LADDR ip:port] [TYPE normal|master|slave|pubsub] [USER username] [SKIPME yes|no] count, err := s.handleClientKill(client, args[1:]) if err != nil { writer.WriteError(err.Error()) return nil } writer.WriteInt(int64(count)) case "pause": // CLIENT PAUSE timeout [WRITE|ALL] // Acknowledge but don't actually implement - pubsub doesn't need this writer.WriteStatus("OK") case "unpause": // CLIENT UNPAUSE writer.WriteStatus("OK") case "reply": // CLIENT REPLY ON|OFF|SKIP if len(args) < 2 { writer.WriteError("ERR wrong number of arguments for 'client reply' command") return nil } mode := strings.ToUpper(string(args[1])) switch mode { case "ON": client.SetReplyMode(pubsub.ReplyOn) writer.WriteStatus("OK") case "OFF": client.SetReplyMode(pubsub.ReplyOff) // Don't send reply when setting to OFF _ = writer.Flush() case "SKIP": client.SetReplyMode(pubsub.ReplySkip) // Send one skip reply, then next will be skipped writer.WriteStatus("OK") default: writer.WriteError("ERR wrong argument for 'client reply' command") } case "tracking": // CLIENT TRACKING ON/OFF [REDIRECT client-id] [PREFIX prefix ...] [BCAST] [OPTIN] [OPTOUT] [NOLOOP] // Acknowledge but don't implement - pekoe is pub/sub only writer.WriteStatus("OK") case "no-evict": // CLIENT NO-EVICT ON|OFF writer.WriteStatus("OK") case "no-touch": // CLIENT NO-TOUCH ON|OFF writer.WriteStatus("OK") case "caching": // CLIENT CACHING YES|NO writer.WriteStatus("OK") default: // Unknown client subcommand writer.WriteError(fmt.Sprintf("ERR Unknown subcommand or wrong number of arguments for '%s'. Try CLIENT HELP.", subcmd)) } return nil } // isValidClientName checks if a client name is valid per Redis spec func isValidClientName(name string) bool { if len(name) == 0 || len(name) > 255 { return false } for _, c := range name { if c < 32 || c > 126 || c == ' ' { return false } } return true } // buildClientList builds CLIENT LIST response for all or filtered clients func (s *Server) buildClientList(args [][]byte) string { clients := s.hub.GetAllClients() // Parse optional filter arguments var filterType string for i := 0; i < len(args); i++ { arg := strings.ToLower(string(args[i])) if arg == "type" && i+1 < len(args) { filterType = strings.ToLower(string(args[i+1])) i++ } } var result string for _, c := range clients { // For pekoe, all clients are "pubsub" type if filterType != "" && filterType != "pubsub" { continue } result += c.ListString() + "\n" } return result } // handleClientKill handles CLIENT KILL subcommand // nolint: cyclop func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (int, error) { clients := s.hub.GetAllClients() killed := 0 var notFoundErr error // Parse arguments var killID uint64 var killIDMin uint64 var killIDMax uint64 var killAddr string var skipMe = true var maxAgeMs int64 for i := 0; i < len(args); i++ { arg := strings.ToUpper(string(args[i])) switch arg { case "ID": if i+2 < len(args) { // ID min max format - kill all IDs in range if _, err := fmt.Sscanf(string(args[i+1]), "%d", &killIDMin); err != nil { return 0, fmt.Errorf("invalid client ID: %s", args[i+1]) } if _, err := fmt.Sscanf(string(args[i+2]), "%d", &killIDMax); err != nil { return 0, fmt.Errorf("invalid client ID: %s", args[i+2]) } i += 2 } else if i+1 < len(args) { // Single ID if _, err := fmt.Sscanf(string(args[i+1]), "%d", &killID); err != nil { return 0, fmt.Errorf("invalid client ID: %s", args[i+1]) } i++ } case "ADDR": if i+1 < len(args) { killAddr = string(args[i+1]) i++ } case "LADDR": // LADDR not fully supported, just skip if i+1 < len(args) { i++ } case "SKIPME": if i+1 < len(args) { skipMe = strings.ToUpper(string(args[i+1])) == "YES" i++ } case "MAXAGE": if i+1 < len(args) { if _, err := fmt.Sscanf(string(args[i+1]), "%d", &maxAgeMs); err != nil { return 0, fmt.Errorf("invalid maxage: %s", args[i+1]) } if maxAgeMs <= 0 { return 0, fmt.Errorf("maxage must be greater than 0") } i++ } case "TYPE": // TYPE filter handled separately in handleClient if i+1 < len(args) { i++ } } } now := time.Now() for _, c := range clients { // Skip current client if requested if skipMe && c.ID() == currentClient.ID() { continue } shouldKill := false // Check ID (single or range) if killID > 0 { shouldKill = c.ID() == killID } else if killIDMin > 0 || killIDMax > 0 { id := c.ID() if killIDMin > 0 && id < killIDMin { shouldKill = false } else if killIDMax > 0 && id > killIDMax { shouldKill = false } else { shouldKill = true } } // Check ADDR if !shouldKill && killAddr != "" { shouldKill = c.RemoteAddr() == killAddr } // Check MAXAGE if shouldKill && maxAgeMs > 0 { age := now.Sub(c.CreatedAt()).Milliseconds() shouldKill = age >= maxAgeMs } if shouldKill { // Unsubscribe and close s.hub.UnsubscribeAll(c) s.hub.UnregisterClient(c) _ = c.Close() killed++ } } // For single ID/ADDR that wasn't found, Redis returns error if killed == 0 && (killID > 0 || killAddr != "") { return 0, fmt.Errorf("No such client") } return killed, notFoundErr } // handleSelect processes SELECT command. // SELECT db - Switch to a different database (0-15). // For pekoe (pubsub-only), this only tracks the selected DB for client compatibility. // Redis default max DB is 16 (0-15), configurable via max-databases in config. func (s *Server) handleSelect(client *pubsub.Client, args [][]byte, writer *ConnWriter) error { if len(args) != 1 { writer.WriteError("ERR wrong number of arguments for 'select' command") return nil } dbIndexStr := string(args[0]) var dbIndex int if _, err := fmt.Sscanf(dbIndexStr, "%d", &dbIndex); err != nil { writer.WriteError("ERR invalid DB index") return nil } if dbIndex < 0 || dbIndex >= MaxDatabases { writer.WriteError("ERR DB index is out of range") return nil } client.SetSelectedDB(dbIndex) writer.WriteStatus("OK") return nil } // handlePsubscribe processes PSUBSCRIBE command. // PSUBSCRIBE pattern [pattern ...] - Subscribe to channels matching patterns. func (s *Server) handlePsubscribe(client *pubsub.Client, patterns [][]byte, writer *ConnWriter) error { if len(patterns) == 0 { writer.WriteError("ERR wrong number of arguments for 'psubscribe' command") return nil } for _, p := range patterns { pattern := string(p) ok, err := s.hub.PatternSubscribe(client, pattern) if err != nil { writer.WriteError(fmt.Sprintf("ERR invalid pattern: %v", err)) return nil } if ok { logger.Debug("client=%s event=psubscribe pattern=%s", client.RemoteAddr(), pattern) // Send confirmation in format: ["psubscribe", pattern, count] writer.WriteArray([]protocol.Value{ protocol.Bulk("psubscribe"), protocol.Bulk(pattern), protocol.Int(int64(client.PatternSubsCount())), }) } } return nil } // handlePunsubscribe processes PUNSUBSCRIBE command. // PUNSUBSCRIBE [pattern [pattern ...]] - Unsubscribe from patterns. func (s *Server) handlePunsubscribe(client *pubsub.Client, patterns [][]byte, writer *ConnWriter) error { if len(patterns) == 0 { // Unsubscribe from all patterns pats := client.GetPatterns() if len(pats) == 0 { // No pattern subscriptions - send single punsubscribe message with -1 (Redis behavior) writer.WriteArray([]protocol.Value{ protocol.Bulk("punsubscribe"), protocol.Bulk("-1"), protocol.Int(0), }) return nil } for _, pattern := range pats { s.hub.PatternUnsubscribe(client, pattern) writer.WriteArray([]protocol.Value{ protocol.Bulk("punsubscribe"), protocol.Bulk(pattern), protocol.Int(int64(client.PatternSubsCount())), }) } } else { // Unsubscribe from specified patterns - only send response if actually subscribed for _, p := range patterns { pattern := string(p) removed := s.hub.PatternUnsubscribe(client, pattern) if removed { logger.Debug("client=%s event=punsubscribe pattern=%s", client.RemoteAddr(), pattern) writer.WriteArray([]protocol.Value{ protocol.Bulk("punsubscribe"), protocol.Bulk(pattern), protocol.Int(int64(client.PatternSubsCount())), }) } } } return nil }