redcon/append.go

451 lines
11 KiB
Go
Raw Normal View History

2017-06-26 07:53:51 +03:00
package redcon
import (
2020-03-04 00:16:44 +03:00
"fmt"
"reflect"
"sort"
2017-06-26 07:53:51 +03:00
"strconv"
"strings"
)
// Kind is the kind of command
type Kind int
const (
// Redis is returned for Redis protocol commands
Redis Kind = iota
// Tile38 is returnd for Tile38 native protocol commands
Tile38
// Telnet is returnd for plain telnet commands
Telnet
)
var errInvalidMessage = &errProtocol{"invalid message"}
// ReadNextCommand reads the next command from the provided packet. It's
2017-07-03 15:29:46 +03:00
// possible that the packet contains multiple commands, or zero commands
2017-06-26 07:53:51 +03:00
// when the packet is incomplete.
2017-07-03 15:45:30 +03:00
// 'argsbuf' is an optional reusable buffer and it can be nil.
// 'complete' indicates that a command was read. false means no more commands.
// 'args' are the output arguments for the command.
// 'kind' is the type of command that was read.
2017-06-26 07:53:51 +03:00
// 'leftover' is any remaining unused bytes which belong to the next command.
2017-07-03 15:45:30 +03:00
// 'err' is returned when a protocol error was encountered.
func ReadNextCommand(packet []byte, argsbuf [][]byte) (
complete bool, args [][]byte, kind Kind, leftover []byte, err error,
2017-06-26 07:53:51 +03:00
) {
2017-07-03 15:45:30 +03:00
args = argsbuf[:0]
2017-06-26 07:53:51 +03:00
if len(packet) > 0 {
if packet[0] != '*' {
if packet[0] == '$' {
return readTile38Command(packet, args)
}
return readTelnetCommand(packet, args)
}
2017-07-03 15:29:46 +03:00
// standard redis command
for s, i := 1, 1; i < len(packet); i++ {
2017-06-26 07:53:51 +03:00
if packet[i] == '\n' {
if packet[i-1] != '\r' {
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet, errInvalidMultiBulkLength
2017-06-26 07:53:51 +03:00
}
2017-07-03 15:29:46 +03:00
count, ok := parseInt(packet[s : i-1])
if !ok || count < 0 {
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet, errInvalidMultiBulkLength
2017-06-26 07:53:51 +03:00
}
i++
2017-07-03 15:29:46 +03:00
if count == 0 {
2017-07-03 15:45:30 +03:00
return true, args[:0], Redis, packet[i:], nil
2017-07-03 15:29:46 +03:00
}
2017-06-26 07:53:51 +03:00
nextArg:
for j := 0; j < count; j++ {
if i == len(packet) {
break
}
if packet[i] != '$' {
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet,
2017-06-26 07:53:51 +03:00
&errProtocol{"expected '$', got '" +
string(packet[i]) + "'"}
}
for s := i + 1; i < len(packet); i++ {
if packet[i] == '\n' {
if packet[i-1] != '\r' {
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet, errInvalidBulkLength
2017-06-26 07:53:51 +03:00
}
n, ok := parseInt(packet[s : i-1])
if !ok || count <= 0 {
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet, errInvalidBulkLength
2017-06-26 07:53:51 +03:00
}
i++
if len(packet)-i >= n+2 {
if packet[i+n] != '\r' || packet[i+n+1] != '\n' {
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet, errInvalidBulkLength
2017-06-26 07:53:51 +03:00
}
args = append(args, packet[i:i+n])
i += n + 2
if j == count-1 {
2017-07-03 15:29:46 +03:00
// done reading
2017-07-03 15:45:30 +03:00
return true, args, Redis, packet[i:], nil
2017-06-26 07:53:51 +03:00
}
continue nextArg
}
break
}
}
break
}
break
}
}
}
2017-07-03 15:45:30 +03:00
return false, args[:0], Redis, packet, nil
2017-06-26 07:53:51 +03:00
}
2017-07-03 15:45:30 +03:00
func readTile38Command(packet []byte, argsbuf [][]byte) (
complete bool, args [][]byte, kind Kind, leftover []byte, err error,
2017-06-26 07:53:51 +03:00
) {
2017-07-03 15:45:30 +03:00
for i := 1; i < len(packet); i++ {
if packet[i] == ' ' {
n, ok := parseInt(packet[1:i])
2017-06-26 07:53:51 +03:00
if !ok || n < 0 {
2017-07-03 15:45:30 +03:00
return false, args[:0], Tile38, packet, errInvalidMessage
2017-06-26 07:53:51 +03:00
}
i++
2017-07-03 15:45:30 +03:00
if len(packet) >= i+n+2 {
if packet[i+n] != '\r' || packet[i+n+1] != '\n' {
return false, args[:0], Tile38, packet, errInvalidMessage
2017-06-26 07:53:51 +03:00
}
2017-07-03 15:45:30 +03:00
line := packet[i : i+n]
2017-06-26 07:53:51 +03:00
reading:
for len(line) != 0 {
if line[0] == '{' {
// The native protocol cannot understand json boundaries so it assumes that
// a json element must be at the end of the line.
args = append(args, line)
break
}
if line[0] == '"' && line[len(line)-1] == '"' {
if len(args) > 0 &&
strings.ToLower(string(args[0])) == "set" &&
strings.ToLower(string(args[len(args)-1])) == "string" {
// Setting a string value that is contained inside double quotes.
// This is only because of the boundary issues of the native protocol.
args = append(args, line[1:len(line)-1])
break
}
}
i := 0
for ; i < len(line); i++ {
if line[i] == ' ' {
value := line[:i]
if len(value) > 0 {
args = append(args, value)
}
line = line[i+1:]
continue reading
}
}
args = append(args, line)
break
}
2017-07-03 15:45:30 +03:00
return true, args, Tile38, packet[i+n+2:], nil
2017-06-26 07:53:51 +03:00
}
break
}
}
2017-07-03 15:45:30 +03:00
return false, args[:0], Tile38, packet, nil
2017-06-26 07:53:51 +03:00
}
2017-07-03 15:45:30 +03:00
func readTelnetCommand(packet []byte, argsbuf [][]byte) (
complete bool, args [][]byte, kind Kind, leftover []byte, err error,
2017-06-26 07:53:51 +03:00
) {
// just a plain text command
2017-07-03 15:45:30 +03:00
for i := 0; i < len(packet); i++ {
if packet[i] == '\n' {
2017-06-26 07:53:51 +03:00
var line []byte
2017-07-03 15:45:30 +03:00
if i > 0 && packet[i-1] == '\r' {
line = packet[:i-1]
2017-06-26 07:53:51 +03:00
} else {
2017-07-03 15:45:30 +03:00
line = packet[:i]
2017-06-26 07:53:51 +03:00
}
var quote bool
var quotech byte
var escape bool
outer:
for {
nline := make([]byte, 0, len(line))
for i := 0; i < len(line); i++ {
c := line[i]
if !quote {
if c == ' ' {
if len(nline) > 0 {
args = append(args, nline)
}
line = line[i+1:]
continue outer
}
if c == '"' || c == '\'' {
if i != 0 {
2017-07-03 15:45:30 +03:00
return false, args[:0], Telnet, packet, errUnbalancedQuotes
2017-06-26 07:53:51 +03:00
}
quotech = c
quote = true
line = line[i+1:]
continue outer
}
} else {
if escape {
escape = false
switch c {
case 'n':
c = '\n'
case 'r':
c = '\r'
case 't':
c = '\t'
}
} else if c == quotech {
quote = false
quotech = 0
args = append(args, nline)
line = line[i+1:]
if len(line) > 0 && line[0] != ' ' {
2017-07-03 15:45:30 +03:00
return false, args[:0], Telnet, packet, errUnbalancedQuotes
2017-06-26 07:53:51 +03:00
}
continue outer
} else if c == '\\' {
escape = true
continue
}
}
nline = append(nline, c)
}
if quote {
2017-07-03 15:45:30 +03:00
return false, args[:0], Telnet, packet, errUnbalancedQuotes
2017-06-26 07:53:51 +03:00
}
if len(line) > 0 {
args = append(args, line)
}
break
}
2017-07-03 15:45:30 +03:00
return true, args, Telnet, packet[i+1:], nil
2017-06-26 07:53:51 +03:00
}
}
2017-07-03 15:45:30 +03:00
return false, args[:0], Telnet, packet, nil
2017-06-26 07:53:51 +03:00
}
2017-11-09 15:50:32 +03:00
// appendPrefix will append a "$3\r\n" style redis prefix for a message.
func appendPrefix(b []byte, c byte, n int64) []byte {
if n >= 0 && n <= 9 {
return append(b, c, byte('0'+n), '\r', '\n')
}
b = append(b, c)
b = strconv.AppendInt(b, n, 10)
return append(b, '\r', '\n')
}
2017-06-26 07:53:51 +03:00
// AppendUint appends a Redis protocol uint64 to the input bytes.
func AppendUint(b []byte, n uint64) []byte {
b = append(b, ':')
b = strconv.AppendUint(b, n, 10)
return append(b, '\r', '\n')
}
// AppendInt appends a Redis protocol int64 to the input bytes.
func AppendInt(b []byte, n int64) []byte {
2017-11-09 15:50:32 +03:00
return appendPrefix(b, ':', n)
2017-06-26 07:53:51 +03:00
}
// AppendArray appends a Redis protocol array to the input bytes.
func AppendArray(b []byte, n int) []byte {
2017-11-09 15:50:32 +03:00
return appendPrefix(b, '*', int64(n))
2017-06-26 07:53:51 +03:00
}
// AppendBulk appends a Redis protocol bulk byte slice to the input bytes.
func AppendBulk(b []byte, bulk []byte) []byte {
2017-11-09 15:50:32 +03:00
b = appendPrefix(b, '$', int64(len(bulk)))
2017-06-26 07:53:51 +03:00
b = append(b, bulk...)
return append(b, '\r', '\n')
}
// AppendBulkString appends a Redis protocol bulk string to the input bytes.
func AppendBulkString(b []byte, bulk string) []byte {
2017-11-09 15:50:32 +03:00
b = appendPrefix(b, '$', int64(len(bulk)))
2017-06-26 07:53:51 +03:00
b = append(b, bulk...)
return append(b, '\r', '\n')
}
// AppendString appends a Redis protocol string to the input bytes.
func AppendString(b []byte, s string) []byte {
b = append(b, '+')
b = append(b, stripNewlines(s)...)
return append(b, '\r', '\n')
}
// AppendError appends a Redis protocol error to the input bytes.
func AppendError(b []byte, s string) []byte {
b = append(b, '-')
b = append(b, stripNewlines(s)...)
return append(b, '\r', '\n')
}
// AppendOK appends a Redis protocol OK to the input bytes.
func AppendOK(b []byte) []byte {
return append(b, '+', 'O', 'K', '\r', '\n')
}
func stripNewlines(s string) string {
for i := 0; i < len(s); i++ {
if s[i] == '\r' || s[i] == '\n' {
s = strings.Replace(s, "\r", " ", -1)
s = strings.Replace(s, "\n", " ", -1)
break
}
}
return s
}
// AppendTile38 appends a Tile38 message to the input bytes.
func AppendTile38(b []byte, data []byte) []byte {
b = append(b, '$')
b = strconv.AppendInt(b, int64(len(data)), 10)
b = append(b, ' ')
b = append(b, data...)
return append(b, '\r', '\n')
}
// AppendNull appends a Redis protocol null to the input bytes.
func AppendNull(b []byte) []byte {
return append(b, '$', '-', '1', '\r', '\n')
}
2020-01-09 15:08:23 +03:00
// AppendBulkFloat appends a float64, as bulk bytes.
func AppendBulkFloat(b []byte, f float64) []byte {
mark1 := len(b)
b = strconv.AppendFloat(b, f, 'f', -1, 64)
mark2 := len(b)
b = AppendBulk(b, b[mark1:mark2])
mark3 := len(b)
copy(b[mark1:], b[mark2:])
b = b[:mark1+(mark3-mark2)]
return b
}
// AppendBulkInt appends an int64, as bulk bytes.
func AppendBulkInt(b []byte, x int64) []byte {
mark1 := len(b)
b = strconv.AppendInt(b, x, 10)
mark2 := len(b)
b = AppendBulk(b, b[mark1:mark2])
mark3 := len(b)
copy(b[mark1:], b[mark2:])
b = b[:mark1+(mark3-mark2)]
return b
}
// AppendBulkUint appends an uint64, as bulk bytes.
func AppendBulkUint(b []byte, x uint64) []byte {
mark1 := len(b)
b = strconv.AppendUint(b, x, 10)
mark2 := len(b)
b = AppendBulk(b, b[mark1:mark2])
mark3 := len(b)
copy(b[mark1:], b[mark2:])
b = b[:mark1+(mark3-mark2)]
return b
}
2020-03-04 00:16:44 +03:00
// AppendAny appends any type to valid Redis type
func AppendAny(b []byte, v interface{}) []byte {
switch v := v.(type) {
case nil:
b = AppendNull(b)
case error:
b = AppendError(b, "ERR "+v.Error())
case string:
b = AppendBulkString(b, v)
case []byte:
b = AppendBulk(b, v)
case bool:
if v {
b = AppendInt(b, 1)
} else {
b = AppendInt(b, 0)
}
case int:
b = AppendInt(b, int64(v))
case int8:
b = AppendInt(b, int64(v))
case int16:
b = AppendInt(b, int64(v))
case int32:
b = AppendInt(b, int64(v))
case int64:
b = AppendInt(b, int64(v))
case uint:
b = AppendUint(b, uint64(v))
case uint8:
b = AppendUint(b, uint64(v))
case uint16:
b = AppendUint(b, uint64(v))
case uint32:
b = AppendUint(b, uint64(v))
case uint64:
b = AppendUint(b, uint64(v))
case float32:
b = AppendBulkFloat(b, float64(v))
case float64:
b = AppendBulkFloat(b, float64(v))
default:
vv := reflect.ValueOf(v)
switch vv.Kind() {
case reflect.Slice:
n := vv.Len()
b = AppendArray(b, n)
for i := 0; i < n; i++ {
b = AppendAny(b, vv.Index(i).Interface())
}
case reflect.Map:
n := vv.Len()
b = AppendArray(b, n*2)
var i int
var strKey bool
var strsKeyItems []strKeyItem
iter := vv.MapRange()
for iter.Next() {
key := iter.Key().Interface()
if i == 0 {
if _, ok := key.(string); ok {
strKey = true
strsKeyItems = make([]strKeyItem, n)
}
}
if strKey {
strsKeyItems[i] = strKeyItem{
key.(string), iter.Value().Interface(),
}
} else {
b = AppendAny(b, key)
b = AppendAny(b, iter.Value().Interface())
}
i++
}
if strKey {
sort.Slice(strsKeyItems, func(i, j int) bool {
return strsKeyItems[i].key < strsKeyItems[j].key
})
for _, item := range strsKeyItems {
b = AppendBulkString(b, item.key)
b = AppendAny(b, item.value)
}
}
default:
b = AppendBulkString(b, fmt.Sprint(v))
}
}
return b
}
type strKeyItem struct {
key string
value interface{}
}