internal/conn/conn.go

@@ -1,15 +1,19 @@
 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
 }
 
 func New(host string, port int) (UDPConn, error) {
@@ -25,6 +29,7 @@ func New(host string, port int) (UDPConn, error) {
 
 	return UDPConn{
 		conn:     conn,
+		response: packet.NewResponse(),
 	}, nil
 }
 
@@ -37,13 +42,39 @@ func (c UDPConn) Write(buf []byte) (int, error) {
 	return n, nil
 }
 
-func (c UDPConn) ReadUntil(timeout time.Time, buf []byte) (int, error) {
+func (c UDPConn) Listen(timeout time.Duration, resp chan<- string, errChan chan<- error) {
-	c.conn.SetReadDeadline(timeout)
+	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 {
-		return 0, err
+				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])
+			}
+		}
 	}
-	return rlen, nil
 }
 
 func (c UDPConn) Close() error {

internal/packet/request.go

@@ -1,23 +1,16 @@
 package packet
 
-import (
+import "fmt"
-	"bytes"
-	"fmt"
-)
-
-const bufSz = 512
 
 type Request struct {
 	magic    []byte
 	password string
-	buf      *bytes.Buffer
 }
 
 func NewRequest(password string) Request {
 	return Request{
 		magic:    []byte{'\xff', '\xff', '\xff', '\xff'},
 		password: password,
-		buf:      bytes.NewBuffer(make([]byte, bufSz)),
 	}
 }
 
@@ -26,8 +19,7 @@ func (r Request) Header() []byte {
 }
 
 func (r Request) Encode(cmd string) []byte {
-	r.buf.Reset()
+	datagram := r.Header()
-	r.buf.Write(r.Header())
+	datagram = append(datagram, fmt.Sprintf(" %s %s", r.password, cmd)...)
-	r.buf.WriteString(fmt.Sprintf(" %s %s", r.password, cmd))
+	return datagram
-	return r.buf.Bytes()
 }

option.go

@@ -1,27 +0,0 @@
-package q3rcon
-
-import "time"
-
-// 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
-	}
-}

q3rcon.go

@@ -1,9 +1,7 @@
 package q3rcon
 
 import (
-	"bytes"
 	"errors"
-	"net"
 	"strings"
 	"time"
 
@@ -11,7 +9,29 @@ import (
 	"github.com/onyx-and-iris/q3rcon/internal/packet"
 )
 
-const respBufSiz = 2048
+// 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
+	}
+}
 
 type Rcon struct {
 	conn     conn.UDPConn
@@ -21,6 +41,8 @@ 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) {
@@ -36,8 +58,7 @@ func New(host string, port int, password string, options ...Option) (*Rcon, erro
 	r := &Rcon{
 		conn:           conn,
 		request:        packet.NewRequest(password),
-		response: packet.NewResponse(),
+		resp:           make(chan string),
-
 		loginTimeout:   5 * time.Second,
 		defaultTimeout: 20 * time.Millisecond,
 		timeouts:       make(map[string]time.Duration),
@@ -84,55 +105,21 @@ func (r Rcon) Send(cmd string) (string, error) {
 		timeout = r.defaultTimeout
 	}
 
-	respChan := make(chan string)
+	e := make(chan error)
-	errChan := make(chan error)
+	go r.conn.Listen(timeout, r.resp, e)
-	go r.listen(timeout, respChan, errChan)
 	_, err := r.conn.Write(r.request.Encode(cmd))
 	if err != nil {
 		return "", err
 	}
 
 	select {
-	case err := <-errChan:
+	case err := <-e:
 		return "", err
-	case resp := <-respChan:
+	case resp := <-r.resp:
 		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()
 }