diff --git a/internal/lib/protocol.go b/internal/lib/protocol.go new file mode 100644 index 0000000..87acd76 --- /dev/null +++ b/internal/lib/protocol.go @@ -0,0 +1,66 @@ +/* +Copyright © 2024 tux <0xtux@pm.me> +*/ + +package lib + +import ( + "bufio" + "errors" + "fmt" + "net" + "strings" +) + +type Message struct { + CMD string + ARG string +} + +type ProtocolHandler struct { + reader *bufio.Reader + writer *bufio.Writer +} + +func InitProtocolHandler(conn net.Conn) *ProtocolHandler { + return &ProtocolHandler{ + reader: bufio.NewReader(conn), + writer: bufio.NewWriter(conn), + } +} + +func (p *ProtocolHandler) ReadMessage() (*Message, error) { + data, err := p.reader.ReadString('\n') + if err != nil { + return nil, err + } + + cmd, arg, err := p.parseMessage(data) + if err != nil { + return nil, errors.New("can't parse data") + } + + return &Message{ + CMD: cmd, + ARG: arg, + }, nil +} + +func (p *ProtocolHandler) WriteMessage(m *Message) error { + _, err := p.writer.WriteString(fmt.Sprintf("%s %s\n", m.CMD, m.ARG)) + if err != nil { + return err + } + return p.writer.Flush() +} + +func (p *ProtocolHandler) parseMessage(data string) (string, string, error) { + data = strings.TrimSpace(data) + d := strings.Fields(data) + + if len(d) != 2 { + return "", "", errors.New("invalid command") + } + + return d[0], d[1], nil +} diff --git a/internal/server/trok.go b/internal/server/trok.go index 24700f4..293d47b 100644 --- a/internal/server/trok.go +++ b/internal/server/trok.go @@ -5,9 +5,9 @@ Copyright © 2024 tux <0xtux@pm.me> package server import ( - "bufio" "net" + "github.com/0xtux/trok/internal/lib" "github.com/rs/zerolog/log" ) @@ -31,10 +31,10 @@ func (t *Trok) Stop() { } func (t *Trok) ControlConnHandler(conn net.Conn) { - reader := bufio.NewReader(conn) + p := lib.InitProtocolHandler(conn) for { - data, err := reader.ReadString('\n') + m, err := p.ReadMessage() if err != nil { if netErr, ok := err.(net.Error); ok && netErr.Timeout() { log.Warn().Msgf("connection timed out: %s", conn.RemoteAddr()) @@ -44,6 +44,26 @@ func (t *Trok) ControlConnHandler(conn net.Conn) { return } - log.Info().Msgf(data) + switch m.CMD { + + case "HELO": + t.handleCMDHELO(p, m) + + case "ACPT": + t.handleCMDACPT(p, m) + + default: + log.Info().Msgf("invalid command") + } } } + +func (t *Trok) handleCMDHELO(p *lib.ProtocolHandler, m *lib.Message) { + log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG) + p.WriteMessage(m) +} + +func (t *Trok) handleCMDACPT(p *lib.ProtocolHandler, m *lib.Message) { + log.Info().Msgf("[CMD] %s [ARG] %s", m.CMD, m.ARG) + p.WriteMessage(m) +}