redcon/redcon.go

425 lines
8.3 KiB
Go

package redcon
import (
"bytes"
"errors"
"io"
"net"
"strconv"
"sync"
)
type Conn interface {
RemoteAddr() string
Close() error
WriteError(msg string)
WriteString(str string)
WriteBulk(bulk string)
WriteInt(num int)
WriteArray(count int)
WriteNull()
}
var (
errUnbalancedQuotes = &errProtocol{"unbalanced quotes in request"}
errInvalidBulkLength = &errProtocol{"invalid bulk length"}
errInvalidMultiBulkLength = &errProtocol{"invalid multibulk length"}
)
type errProtocol struct {
msg string
}
func (err *errProtocol) Error() string {
return "Protocol error: " + err.msg
}
func ListenAndServe(
addr string, handler func(conn Conn, cmds [][]string),
accept func(conn Conn) bool, closed func(conn Conn, err error),
) error {
ln, err := net.Listen("tcp", addr)
if err != nil {
return err
}
defer ln.Close()
var mu sync.Mutex
for {
conn, err := ln.Accept()
if err != nil {
return err
}
wr := newWriter(conn)
wrc := &connWriter{wr, conn.RemoteAddr().String()}
if accept != nil && !accept(wrc) {
conn.Close()
continue
}
go func() {
var err error
defer func() {
conn.Close()
if closed != nil {
mu.Lock()
defer mu.Unlock()
if err == io.EOF {
err = nil
}
closed(wrc, err)
}
}()
rd := newReader(conn)
err = func() error {
for {
cmds, err := rd.ReadCommands()
if err != nil {
if err, ok := err.(*errProtocol); ok {
// All protocol errors should attempt a response to
// the client. Ignore errors.
wr.WriteError("ERR " + err.Error())
wr.Flush()
}
return err
}
if len(cmds) > 0 {
handler(wrc, cmds)
}
if wr.err != nil {
if wr.err == errClosed {
return nil
}
return wr.err
}
if err := wr.Flush(); err != nil {
return err
}
}
}()
}()
}
}
type connWriter struct {
wr *respWriter
addr string
}
func (wrc *connWriter) Close() error {
return wrc.wr.Close()
}
func (wrc *connWriter) WriteString(str string) {
wrc.wr.WriteString(str)
}
func (wrc *connWriter) WriteBulk(bulk string) {
wrc.wr.WriteBulk(bulk)
}
func (wrc *connWriter) WriteInt(num int) {
wrc.wr.WriteInt(num)
}
func (wrc *connWriter) WriteError(msg string) {
wrc.wr.WriteError(msg)
}
func (wrc *connWriter) WriteArray(count int) {
wrc.wr.WriteMultiBulkStart(count)
}
func (wrc *connWriter) WriteNull() {
wrc.wr.WriteNull()
}
func (wrc *connWriter) RemoteAddr() string {
return wrc.addr
}
// Reader represents a RESP command reader.
type respReader struct {
r io.Reader // base reader
b []byte // unprocessed bytes
a []byte // static read buffer
}
// NewReader returns a RESP command reader.
func newReader(r io.Reader) *respReader {
return &respReader{
r: r,
a: make([]byte, 8192),
}
}
// ReadCommands reads one or more commands from the reader.
func (r *respReader) ReadCommands() ([][]string, error) {
if len(r.b) > 0 {
// we have some potential commands.
var cmds [][]string
next:
switch r.b[0] {
default:
// just a plain text command
for i := 0; i < len(r.b); i++ {
if r.b[i] == '\n' {
var line []byte
if i > 0 && r.b[i-1] == '\r' {
line = r.b[:i-1]
} else {
line = r.b[:i]
}
var args []string
var quote bool
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, string(nline))
}
line = line[i+1:]
continue outer
}
if c == '"' {
if i != 0 {
return nil, errUnbalancedQuotes
}
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 == '"' {
quote = false
args = append(args, string(nline))
line = line[i+1:]
if len(line) > 0 && line[0] != ' ' {
return nil, errUnbalancedQuotes
}
continue outer
} else if c == '\\' {
escape = true
continue
}
}
nline = append(nline, c)
}
if quote {
return nil, errUnbalancedQuotes
}
if len(line) > 0 {
args = append(args, string(line))
}
break
}
if len(args) > 0 {
cmds = append(cmds, args)
}
r.b = r.b[i+1:]
if len(r.b) > 0 {
goto next
} else {
goto done
}
}
}
case '*':
// resp formatted command
var si int
outer2:
for i := 0; i < len(r.b); i++ {
var args []string
if r.b[i] == '\n' {
if r.b[i-1] != '\r' {
return nil, errInvalidMultiBulkLength
}
ni, err := strconv.ParseInt(string(r.b[si+1:i-1]), 10, 64)
if err != nil || ni <= 0 {
return nil, errInvalidMultiBulkLength
}
args = make([]string, 0, int(ni))
for j := 0; j < int(ni); j++ {
// read bulk length
i++
if i < len(r.b) {
if r.b[i] != '$' {
return nil, &errProtocol{"expected '$', got '" +
string(r.b[i]) + "'"}
}
si = i
for ; i < len(r.b); i++ {
if r.b[i] == '\n' {
if r.b[i-1] != '\r' {
return nil, errInvalidBulkLength
}
s := string(r.b[si+1 : i-1])
ni2, err := strconv.ParseInt(s, 10, 64)
if err != nil || ni2 < 0 {
return nil, errInvalidBulkLength
}
if i+int(ni2)+2 >= len(r.b) {
// not ready
break outer2
}
if r.b[i+int(ni2)+2] != '\n' ||
r.b[i+int(ni2)+1] != '\r' {
return nil, errInvalidBulkLength
}
arg := string(r.b[i+1 : i+1+int(ni2)])
i += int(ni2) + 2
args = append(args, arg)
break
}
}
}
}
if len(args) == cap(args) {
cmds = append(cmds, args)
r.b = r.b[i+1:]
if len(r.b) > 0 {
goto next
} else {
goto done
}
}
}
}
}
done:
if len(r.b) == 0 {
r.b = nil
}
if len(cmds) > 0 {
return cmds, nil
}
}
n, err := r.r.Read(r.a[:])
if err != nil {
if err == io.EOF {
if len(r.b) > 0 {
return nil, io.ErrUnexpectedEOF
}
}
return nil, err
}
r.b = append(r.b, r.a[:n]...)
return r.ReadCommands()
}
var errClosed = errors.New("closed")
type respWriter struct {
w io.Writer
b *bytes.Buffer
err error
}
func newWriter(w io.Writer) *respWriter {
return &respWriter{w: w, b: &bytes.Buffer{}}
}
func (w *respWriter) WriteNull() error {
if w.err != nil {
return w.err
}
w.b.WriteString("$-1\r\n")
return nil
}
func (w *respWriter) WriteMultiBulkStart(count int) error {
if w.err != nil {
return w.err
}
w.b.WriteByte('*')
w.b.WriteString(strconv.FormatInt(int64(count), 10))
w.b.WriteString("\r\n")
return nil
}
func (w *respWriter) WriteBulk(bulk string) error {
if w.err != nil {
return w.err
}
w.b.WriteByte('$')
w.b.WriteString(strconv.FormatInt(int64(len(bulk)), 10))
w.b.WriteString("\r\n")
w.b.WriteString(bulk)
w.b.WriteString("\r\n")
return nil
}
func (w *respWriter) Flush() error {
if w.err != nil {
return w.err
}
if w.b.Len() == 0 {
return nil
}
if _, err := w.b.WriteTo(w.w); err != nil {
w.err = err
return err
}
w.b.Reset()
return nil
}
func (w *respWriter) WriteMultiBulk(bulks []string) error {
if err := w.WriteMultiBulkStart(len(bulks)); err != nil {
return err
}
for _, bulk := range bulks {
if err := w.WriteBulk(bulk); err != nil {
return err
}
}
return nil
}
func (w *respWriter) WriteError(msg string) error {
if w.err != nil {
return w.err
}
w.b.WriteByte('-')
w.b.WriteString(msg)
w.b.WriteString("\r\n")
return nil
}
func (w *respWriter) WriteString(msg string) error {
if w.err != nil {
return w.err
}
w.b.WriteByte('+')
w.b.WriteString(msg)
w.b.WriteString("\r\n")
return nil
}
func (w *respWriter) WriteInt(num int) error {
if w.err != nil {
return w.err
}
w.b.WriteByte(':')
w.b.WriteString(strconv.FormatInt(int64(num), 10))
w.b.WriteString("\r\n")
return nil
}
func (w *respWriter) Close() error {
if w.err != nil {
return w.err
}
if err := w.Flush(); err != nil {
w.err = err
return err
}
w.err = errClosed
return nil
}