Request now uses reusable buffer

rename channel variables

README.md

@@ -52,16 +52,28 @@ func main() {
 }
 ```
 
-#### `WithDefaultTimeout(timeout time.Duration)`
+#### `WithLoginTimeout(timeout time.Duration)`
 
-You may want to change the default timeout if some of your responses are getting mixed together (stability can depend on connection to the server). For example, on LAN I can leave the default at 20ms, when connecting remotely I normally increase this to 50ms.
+If the server was just started or is currently performing a long operation like map rotating then it's possible to receive empty rcon responses. In which case you'll want to retry login. Use this functional option to set the max timeout for logins, it defaults to 5 seconds.
 
 For example:
 
 ```go
-	rcon, err := q3rcon.New(
-		host, port, password,
-		q3rcon.WithDefaultTimeout(50*time.Millisecond))
+rcon, err := q3rcon.New(
+	host, port, password,
+	q3rcon.WithLoginTimeout(2*time.Second))
+```
+
+#### `WithDefaultTimeout(timeout time.Duration)`
+
+You may want to change the default timeout if some of your responses are getting mixed together (stability can depend on connection to the server). For example, over localhost I can leave the default at 20ms, when connecting remotely I normally increase this to 50ms.
+
+For example:
+
+```go
+rcon, err := q3rcon.New(
+	host, port, password,
+	q3rcon.WithDefaultTimeout(50*time.Millisecond))
 ```
 
 #### `WithTimeouts(timeouts map[string]time.Duration)`
@@ -69,14 +81,14 @@ For example:
 Perhaps there are some requests that take a long time to receive a response but you want the whole response in one chunk. Pass a timeouts map, for example:
 
 ```go
-	timeouts := map[string]time.Duration{
-		"map_rotate":  1200 * time.Millisecond,
-		"map_restart": 1200 * time.Millisecond,
-	}
+timeouts := map[string]time.Duration{
+	"map_rotate":  1200 * time.Millisecond,
+	"map_restart": 1200 * time.Millisecond,
+}
 
-	rcon, err := q3rcon.New(
-		host, port, password,
-		q3rcon.WithTimeouts(timeouts))
+rcon, err := q3rcon.New(
+	host, port, password,
+	q3rcon.WithTimeouts(timeouts))
 ```
 
 ## Command line
@@ -103,8 +115,17 @@ q3rcon -h=localhost -p=30000 -P="rconpassword" -i
 
 If interactive mode is enabled, any arguments sent on the command line will be ignored.
 
-## `Logging`
+## Your own implementation
+
+The included CLI is a generic implementation, while it can be used out of the box you may find that some requests result in fragmented responses. The solution is to implement your own version, adjusting the timings with the functional options as detailed above. I could have increased the default timeouts but that would add unnecessary delay for most requests, so I decided to leave those details to the users of the package.
+
+Since you can include the q3rcon package into your own package you can easily make your own modifications, for example, I added [colour to the terminal][status] and [reformatted some of the responses][mapname].
+
+## Logging
 
 Log level may be set by passing the `-l` flag with a number from 0 up to 6 where
 
 0 = Panic, 1 = Fatal, 2 = Error, 3 = Warning, 4 = Info, 5 = Debug, 6 = Trace
+
+[status]: ./img/status.png
+[mapname]: ./img/mapname.png

cmd/q3rcon/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
@@ -16,6 +17,12 @@ import (
 
 var interactive bool
 
+func exit(err error) {
+	_, _ = fmt.Fprintf(os.Stderr, "Error: %s\n", err)
+	flag.Usage()
+	os.Exit(1)
+}
+
 func main() {
 	var (
 		host     string
@@ -28,8 +35,8 @@ func main() {
 	flag.StringVar(&host, "h", "localhost", "hostname of the server (shorthand)")
 	flag.IntVar(&port, "port", 28960, "port of the server")
 	flag.IntVar(&port, "p", 28960, "port of the server (shorthand)")
-	flag.StringVar(&password, "password", "", "hostname of the server")
-	flag.StringVar(&password, "P", "", "hostname of the server (shorthand)")
+	flag.StringVar(&password, "password", "", "rcon password")
+	flag.StringVar(&password, "P", "", "rcon password (shorthand)")
 
 	flag.BoolVar(&interactive, "interactive", false, "run in interactive mode")
 	flag.BoolVar(&interactive, "i", false, "run in interactive mode")
@@ -57,6 +64,11 @@ func main() {
 		return
 	}
 
+	if len(flag.Args()) == 0 {
+		err = errors.New("no rcon commands passed")
+		exit(err)
+	}
+
 	for _, arg := range flag.Args() {
 		resp, err := rcon.Send(arg)
 		if err != nil {
@@ -68,8 +80,7 @@ func main() {
 }
 
 func connectRcon(host string, port int, password string) (*q3rcon.Rcon, error) {
-	rcon, err := q3rcon.New(
-		host, port, password)
+	rcon, err := q3rcon.New(host, port, password)
 	if err != nil {
 		return nil, err
 	}

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 {

internal/packet/request.go

@@ -1,16 +1,23 @@
 package packet
 
-import "fmt"
+import (
+	"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)),
 	}
 }
 
@@ -19,7 +26,8 @@ func (r Request) Header() []byte {
 }
 
 func (r Request) Encode(cmd string) []byte {
-	datagram := r.Header()
-	datagram = append(datagram, fmt.Sprintf(" %s %s", r.password, cmd)...)
-	return datagram
+	r.buf.Reset()
+	r.buf.Write(r.Header())
+	r.buf.WriteString(fmt.Sprintf(" %s %s", r.password, cmd))
+	return r.buf.Bytes()
 }

option.go

@@ -0,0 +1,27 @@
+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,7 +1,9 @@
 package q3rcon
 
 import (
+	"bytes"
 	"errors"
+	"net"
 	"strings"
 	"time"
 
@@ -9,31 +11,16 @@ 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)
-
-func WithDefaultTimeout(timeout time.Duration) Option {
-	return func(rcon *Rcon) {
-		rcon.defaultTimeout = timeout
-	}
-}
-
-// WithTimeouts is a functional option to set the timeouts for responses
-func WithTimeouts(timeouts map[string]time.Duration) Option {
-	return func(rcon *Rcon) {
-		rcon.timeouts = timeouts
-	}
-}
+const respBufSiz = 2048
 
 type Rcon struct {
 	conn     conn.UDPConn
 	request  packet.Request
 	response packet.Response
 
+	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) {
@@ -47,9 +34,11 @@ 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),
 	}
@@ -58,8 +47,7 @@ func New(host string, port int, password string, options ...Option) (*Rcon, erro
 		o(r)
 	}
 
-	err = r.Login()
-	if err != nil {
+	if err = r.Login(); err != nil {
 		return nil, err
 	}
 
@@ -67,7 +55,7 @@ func New(host string, port int, password string, options ...Option) (*Rcon, erro
 }
 
 func (r Rcon) Login() error {
-	timeout := time.After(2 * time.Second)
+	timeout := time.After(r.loginTimeout)
 	for {
 		select {
 		case <-timeout:
@@ -96,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()
 }