tile38/vendor/github.com/tidwall/resp/server.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
			}
		}
	}
}