From 8c1bc0f04cf8bdfb2fa56662eaf5398dfa1efdca Mon Sep 17 00:00:00 2001
From: onyx-and-iris <code@onyxandiris.online>
Date: Fri, 8 Nov 2024 15:30:47 +0000
Subject: [PATCH] move listen() into Rcon

rename channel variables
---
 internal/conn/conn.go | 47 +++++---------------------
 q3rcon.go             | 77 +++++++++++++++++++++++++------------------
 2 files changed, 53 insertions(+), 71 deletions(-)

diff --git a/internal/conn/conn.go b/internal/conn/conn.go
index 6dea113..9ae39e4 100644
--- a/internal/conn/conn.go
+++ b/internal/conn/conn.go
@@ -1,19 +1,15 @@
 package conn
 
 import (
-	"bytes"
 	"fmt"
 	"net"
-	"strings"
 	"time"
 
-	"github.com/onyx-and-iris/q3rcon/internal/packet"
 	log "github.com/sirupsen/logrus"
 )
 
 type UDPConn struct {
-	conn     *net.UDPConn
-	response packet.Response
+	conn *net.UDPConn
 }
 
 func New(host string, port int) (UDPConn, error) {
@@ -28,8 +24,7 @@ func New(host string, port int) (UDPConn, error) {
 	log.Infof("Outgoing address %s", conn.RemoteAddr())
 
 	return UDPConn{
-		conn:     conn,
-		response: packet.NewResponse(),
+		conn: conn,
 	}, nil
 }
 
@@ -42,39 +37,13 @@ func (c UDPConn) Write(buf []byte) (int, error) {
 	return n, nil
 }
 
-func (c UDPConn) Listen(timeout time.Duration, resp chan<- string, errChan chan<- error) {
-	c.conn.SetReadDeadline(time.Now().Add(timeout))
-	done := make(chan struct{})
-	var sb strings.Builder
-	buf := make([]byte, 2048)
-
-	for {
-		select {
-		case <-done:
-			resp <- sb.String()
-			return
-		default:
-			rlen, _, err := c.conn.ReadFromUDP(buf)
-			if err != nil {
-				e, ok := err.(net.Error)
-				if ok {
-					if e.Timeout() {
-						close(done)
-					} else {
-						errChan <- e
-						return
-					}
-				}
-			}
-			if rlen < len(c.response.Header()) {
-				continue
-			}
-
-			if bytes.HasPrefix(buf, c.response.Header()) {
-				sb.Write(buf[len(c.response.Header()):rlen])
-			}
-		}
+func (c UDPConn) ReadUntil(timeout time.Time, buf []byte) (int, error) {
+	c.conn.SetReadDeadline(timeout)
+	rlen, _, err := c.conn.ReadFromUDP(buf)
+	if err != nil {
+		return 0, err
 	}
+	return rlen, nil
 }
 
 func (c UDPConn) Close() error {
diff --git a/q3rcon.go b/q3rcon.go
index 13e68cc..2630338 100644
--- a/q3rcon.go
+++ b/q3rcon.go
@@ -1,7 +1,9 @@
 package q3rcon
 
 import (
+	"bytes"
 	"errors"
+	"net"
 	"strings"
 	"time"
 
@@ -9,29 +11,7 @@ import (
 	"github.com/onyx-and-iris/q3rcon/internal/packet"
 )
 
-// Option is a functional option type that allows us to configure the VbanTxt.
-type Option func(*Rcon)
-
-// WithLoginTimeout is a functional option to set the login timeout
-func WithLoginTimeout(timeout time.Duration) Option {
-	return func(rcon *Rcon) {
-		rcon.loginTimeout = timeout
-	}
-}
-
-// WithDefaultTimeout is a functional option to set the default response timeout
-func WithDefaultTimeout(timeout time.Duration) Option {
-	return func(rcon *Rcon) {
-		rcon.defaultTimeout = timeout
-	}
-}
-
-// WithTimeouts is a functional option to set the timeouts for responses per command
-func WithTimeouts(timeouts map[string]time.Duration) Option {
-	return func(rcon *Rcon) {
-		rcon.timeouts = timeouts
-	}
-}
+const respBufSiz = 2048
 
 type Rcon struct {
 	conn     conn.UDPConn
@@ -41,8 +21,6 @@ type Rcon struct {
 	loginTimeout   time.Duration
 	defaultTimeout time.Duration
 	timeouts       map[string]time.Duration
-
-	resp chan string
 }
 
 func New(host string, port int, password string, options ...Option) (*Rcon, error) {
@@ -56,9 +34,10 @@ func New(host string, port int, password string, options ...Option) (*Rcon, erro
 	}
 
 	r := &Rcon{
-		conn:           conn,
-		request:        packet.NewRequest(password),
-		resp:           make(chan string),
+		conn:     conn,
+		request:  packet.NewRequest(password),
+		response: packet.NewResponse(),
+
 		loginTimeout:   5 * time.Second,
 		defaultTimeout: 20 * time.Millisecond,
 		timeouts:       make(map[string]time.Duration),
@@ -105,21 +84,55 @@ func (r Rcon) Send(cmd string) (string, error) {
 		timeout = r.defaultTimeout
 	}
 
-	e := make(chan error)
-	go r.conn.Listen(timeout, r.resp, e)
+	respChan := make(chan string)
+	errChan := make(chan error)
+	go r.listen(timeout, respChan, errChan)
 	_, err := r.conn.Write(r.request.Encode(cmd))
 	if err != nil {
 		return "", err
 	}
 
 	select {
-	case err := <-e:
+	case err := <-errChan:
 		return "", err
-	case resp := <-r.resp:
+	case resp := <-respChan:
 		return strings.TrimPrefix(resp, string(r.response.Header())), nil
 	}
 }
 
+func (r Rcon) listen(timeout time.Duration, respChan chan<- string, errChan chan<- error) {
+	done := make(chan struct{})
+	respBuf := make([]byte, respBufSiz)
+	var sb strings.Builder
+
+	for {
+		select {
+		case <-done:
+			respChan <- sb.String()
+			return
+		default:
+			rlen, err := r.conn.ReadUntil(time.Now().Add(timeout), respBuf)
+			if err != nil {
+				e, ok := err.(net.Error)
+				if ok {
+					if e.Timeout() {
+						close(done)
+					} else {
+						errChan <- e
+						return
+					}
+				}
+			}
+
+			if rlen > len(r.response.Header()) {
+				if bytes.HasPrefix(respBuf, r.response.Header()) {
+					sb.Write(respBuf[len(r.response.Header()):rlen])
+				}
+			}
+		}
+	}
+}
+
 func (r Rcon) Close() {
 	r.conn.Close()
 }