480 lines
13 KiB
Go
480 lines
13 KiB
Go
// @tag server
|
|
// @tag handlers:pubsub
|
|
package server
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"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, len(channels))
|
|
for i, ch := range channels {
|
|
result[i] = protocol.Bulk(ch)
|
|
}
|
|
writer.WriteArray(result)
|
|
|
|
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 <subcommand> [<argument> [<argument> ...]]>",
|
|
" PUBSUB CHANNELS [pattern] - List active channels (optionally filtered by pattern)",
|
|
" PUBSUB NUMSUB [channel ...] - Get subscriber count for channels",
|
|
" 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
|
|
|
|
// Parse arguments
|
|
var killID uint64
|
|
var killAddr string
|
|
var skipMe = true
|
|
|
|
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
|
|
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 "SKIPME":
|
|
if i+1 < len(args) {
|
|
skipMe = strings.ToUpper(string(args[i+1])) == "YES"
|
|
i++
|
|
}
|
|
}
|
|
}
|
|
|
|
for _, c := range clients {
|
|
// Skip current client if requested
|
|
if skipMe && c.ID() == currentClient.ID() {
|
|
continue
|
|
}
|
|
|
|
shouldKill := false
|
|
|
|
if killID > 0 && c.ID() == killID {
|
|
shouldKill = true
|
|
}
|
|
if killAddr != "" && c.RemoteAddr() == killAddr {
|
|
shouldKill = true
|
|
}
|
|
|
|
if shouldKill {
|
|
// Unsubscribe and close
|
|
s.hub.UnsubscribeAll(c)
|
|
s.hub.UnregisterClient(c)
|
|
_ = c.Close()
|
|
killed++
|
|
}
|
|
}
|
|
|
|
return killed, 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
|
|
}
|