move listen() into Rcon

rename channel variables
This commit is contained in:
onyx-and-iris 2024-11-08 15:30:47 +00:00
parent dd1d530d44
commit 8c1bc0f04c
2 changed files with 53 additions and 71 deletions

View File

@ -1,19 +1,15 @@
package conn package conn
import ( import (
"bytes"
"fmt" "fmt"
"net" "net"
"strings"
"time" "time"
"github.com/onyx-and-iris/q3rcon/internal/packet"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
type UDPConn struct { type UDPConn struct {
conn *net.UDPConn conn *net.UDPConn
response packet.Response
} }
func New(host string, port int) (UDPConn, error) { func New(host string, port int) (UDPConn, error) {
@ -29,7 +25,6 @@ func New(host string, port int) (UDPConn, error) {
return UDPConn{ return UDPConn{
conn: conn, conn: conn,
response: packet.NewResponse(),
}, nil }, nil
} }
@ -42,39 +37,13 @@ func (c UDPConn) Write(buf []byte) (int, error) {
return n, nil return n, nil
} }
func (c UDPConn) Listen(timeout time.Duration, resp chan<- string, errChan chan<- error) { func (c UDPConn) ReadUntil(timeout time.Time, buf []byte) (int, error) {
c.conn.SetReadDeadline(time.Now().Add(timeout)) c.conn.SetReadDeadline(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) rlen, _, err := c.conn.ReadFromUDP(buf)
if err != nil { if err != nil {
e, ok := err.(net.Error) return 0, err
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 { func (c UDPConn) Close() error {

View File

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