This is a fork of:
https://github.com/roglew/puppy
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
115 lines
2.1 KiB
115 lines
2.1 KiB
7 years ago
|
package main
|
||
|
|
||
|
import (
|
||
|
"bufio"
|
||
|
"encoding/json"
|
||
|
"fmt"
|
||
|
"strings"
|
||
|
"io"
|
||
|
"log"
|
||
|
"net"
|
||
|
)
|
||
|
|
||
|
/*
|
||
|
Message Server
|
||
|
*/
|
||
|
|
||
|
type MessageHandler func([]byte, net.Conn, *log.Logger, *InterceptingProxy)
|
||
|
|
||
|
type MessageListener struct {
|
||
|
handlers map[string]MessageHandler
|
||
|
iproxy *InterceptingProxy
|
||
|
Logger *log.Logger
|
||
|
}
|
||
|
|
||
|
type commandData struct {
|
||
|
Command string
|
||
|
}
|
||
|
|
||
|
type errorMessage struct {
|
||
|
Success bool
|
||
|
Reason string
|
||
|
}
|
||
|
|
||
|
func NewMessageListener(l *log.Logger, iproxy *InterceptingProxy) *MessageListener {
|
||
|
m := &MessageListener{
|
||
|
handlers: make(map[string]MessageHandler),
|
||
|
iproxy: iproxy,
|
||
|
Logger: l,
|
||
|
}
|
||
|
return m
|
||
|
}
|
||
|
|
||
|
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
|
||
|
}
|
||
|
|
||
|
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())
|
||
|
}
|
||
|
}
|
||
|
}()
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func ErrorResponse(w io.Writer, reason string) {
|
||
|
var m errorMessage
|
||
|
m.Success = false
|
||
|
m.Reason = reason
|
||
|
MessageResponse(w, m)
|
||
|
}
|
||
|
|
||
|
func MessageResponse(w io.Writer, m interface{}) {
|
||
|
b, err := json.Marshal(&m)
|
||
|
if err != nil {
|
||
|
panic(err)
|
||
|
}
|
||
|
MainLogger.Printf("< %s\n", string(b))
|
||
|
w.Write(b)
|
||
|
w.Write([]byte("\n"))
|
||
|
}
|
||
|
|
||
|
func ReadMessage(r *bufio.Reader) ([]byte, error) {
|
||
|
m, err := r.ReadBytes('\n')
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
return m, nil
|
||
|
}
|