PuppyProxy/messageserv.go

package puppy

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
)

/*
Message Server
*/

// A handler to handle a JSON message
type MessageHandler func(message []byte, conn net.Conn, logger *log.Logger, iproxy *InterceptingProxy)

// A listener that handles reading JSON messages and sending them to the correct handler
type MessageListener struct {
	handlers map[string]MessageHandler
	iproxy   *InterceptingProxy
	Logger   *log.Logger
}

type commandData struct {
	Command string
}

type errorMessage struct {
	Success bool
	Reason  string
}

// NewMessageListener creates a new message listener associated with the given intercepting proxy
func NewMessageListener(l *log.Logger, iproxy *InterceptingProxy) *MessageListener {
	m := &MessageListener{
		handlers: make(map[string]MessageHandler),
		iproxy:   iproxy,
		Logger:   l,
	}
	return m
}

// AddHandler will have the listener call the given handler when the "Command" parameter matches the given value
func (l *MessageListener) AddHandler(command string, handler MessageHandler) {
	l.handlers[strings.ToLower(command)] = handler
}

func (l *MessageListener) Handle(message []byte, conn net.Conn) error {
	var c commandData
	if err := json.Unmarshal(message, &c); err != nil {
		return fmt.Errorf("error parsing message: %s", err.Error())
	}

	handler, ok := l.handlers[strings.ToLower(c.Command)]
	if !ok {
		return fmt.Errorf("unknown command: %s", c.Command)
	}

	l.Logger.Printf("Calling handler for \"%s\"...", c.Command)
	handler(message, conn, l.Logger, l.iproxy)
	return nil
}

// Serve will have the listener serve messages on the given listener
func (l *MessageListener) Serve(nl net.Listener) {
	for {
		conn, err := nl.Accept()
		if err != nil {
			// Listener closed
			break
		}

		reader := bufio.NewReader(conn)
		go func() {
			for {
				m, err := ReadMessage(reader)
				l.Logger.Printf("> %s\n", m)
				if err != nil {
					if err != io.EOF {
						ErrorResponse(conn, "error reading message")
					}
					return
				}
				err = l.Handle(m, conn)
				if err != nil {
					ErrorResponse(conn, err.Error())
				}
			}
		}()
	}
}

// Error response writes an error message to the given writer
func ErrorResponse(w io.Writer, reason string) {
	var m errorMessage
	m.Success = false
	m.Reason = reason
	MessageResponse(w, m)
}

// MessageResponse writes a response to a given writer
func MessageResponse(w io.Writer, m interface{}) {
	b, err := json.Marshal(&m)
	if err != nil {
		panic(err)
	}
	w.Write(b)
	w.Write([]byte("\n"))
}

// ReadMessage reads a message from the given reader
func ReadMessage(r *bufio.Reader) ([]byte, error) {
	m, err := r.ReadBytes('\n')
	if err != nil {
		return nil, err
	}
	return m, nil
}