mirror of https://github.com/tidwall/tile38.git
136 lines
3.0 KiB
Go
136 lines
3.0 KiB
Go
package resp
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
// Server represents a RESP server which handles reading RESP Values.
|
|
type Server struct {
|
|
mu sync.RWMutex
|
|
handlers map[string]func(conn *Conn, args []Value) bool
|
|
accept func(conn *Conn) bool
|
|
}
|
|
|
|
// Conn represents a RESP network connection.
|
|
type Conn struct {
|
|
*Reader
|
|
*Writer
|
|
base net.Conn
|
|
RemoteAddr string
|
|
}
|
|
|
|
// NewConn returns a Conn.
|
|
func NewConn(conn net.Conn) *Conn {
|
|
return &Conn{
|
|
Reader: NewReader(conn),
|
|
Writer: NewWriter(conn),
|
|
base: conn,
|
|
RemoteAddr: conn.RemoteAddr().String(),
|
|
}
|
|
}
|
|
|
|
// NewServer returns a new Server.
|
|
func NewServer() *Server {
|
|
return &Server{
|
|
handlers: make(map[string]func(conn *Conn, args []Value) bool),
|
|
}
|
|
}
|
|
|
|
// HandleFunc registers the handler function for the given command.
|
|
// The conn parameter is a Conn type and it can be used to read and write further RESP messages from and to the connection.
|
|
// Returning false will close the connection.
|
|
func (s *Server) HandleFunc(command string, handler func(conn *Conn, args []Value) bool) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.handlers[strings.ToUpper(command)] = handler
|
|
}
|
|
|
|
// AcceptFunc registers a function for accepting connections.
|
|
// Calling this function is optional and it allows for total control over reading and writing RESP Values from and to the connections.
|
|
// Returning false will close the connection.
|
|
func (s *Server) AcceptFunc(accept func(conn *Conn) bool) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.accept = accept
|
|
}
|
|
|
|
// ListenAndServe listens on the TCP network address addr for incoming connections.
|
|
func (s *Server) ListenAndServe(addr string) error {
|
|
l, err := net.Listen("tcp", addr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer l.Close()
|
|
for {
|
|
conn, err := l.Accept()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
go func() {
|
|
err = s.handleConn(conn)
|
|
if err != nil {
|
|
if _, ok := err.(*errProtocol); ok {
|
|
io.WriteString(conn, "-ERR "+formSingleLine(err.Error())+"\r\n")
|
|
} else {
|
|
io.WriteString(conn, "-ERR unknown error\r\n")
|
|
}
|
|
}
|
|
conn.Close()
|
|
}()
|
|
}
|
|
}
|
|
|
|
func (s *Server) handleConn(nconn net.Conn) error {
|
|
conn := NewConn(nconn)
|
|
s.mu.RLock()
|
|
accept := s.accept
|
|
s.mu.RUnlock()
|
|
if accept != nil {
|
|
if !accept(conn) {
|
|
return nil
|
|
}
|
|
}
|
|
for {
|
|
v, _, _, err := conn.ReadMultiBulk()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
values := v.Array()
|
|
if len(values) == 0 {
|
|
continue
|
|
}
|
|
lccommandName := values[0].String()
|
|
commandName := strings.ToUpper(lccommandName)
|
|
s.mu.RLock()
|
|
h := s.handlers[commandName]
|
|
s.mu.RUnlock()
|
|
switch commandName {
|
|
case "QUIT":
|
|
if h == nil {
|
|
conn.WriteSimpleString("OK")
|
|
return nil
|
|
}
|
|
case "PING":
|
|
if h == nil {
|
|
if err := conn.WriteSimpleString("PONG"); err != nil {
|
|
return err
|
|
}
|
|
continue
|
|
}
|
|
}
|
|
if h == nil {
|
|
if err := conn.WriteError(errors.New("ERR unknown command '" + lccommandName + "'")); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
if !h(conn, values) {
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
}
|