diff --git a/internal/conn/conn.go b/internal/conn/conn.go index 0b622be..6dea113 100644 --- a/internal/conn/conn.go +++ b/internal/conn/conn.go @@ -42,15 +42,15 @@ func (c UDPConn) Write(buf []byte) (int, error) { return n, nil } -func (c UDPConn) Listen(timeout time.Duration, resp chan<- string) { +func (c UDPConn) Listen(timeout time.Duration, resp chan<- string, errChan chan<- error) { c.conn.SetReadDeadline(time.Now().Add(timeout)) - ch := make(chan struct{}) + done := make(chan struct{}) var sb strings.Builder buf := make([]byte, 2048) for { select { - case <-ch: + case <-done: resp <- sb.String() return default: @@ -59,13 +59,14 @@ func (c UDPConn) Listen(timeout time.Duration, resp chan<- string) { e, ok := err.(net.Error) if ok { if e.Timeout() { - close(ch) + close(done) } else { - log.Error(e) + errChan <- e + return } } } - if rlen == 0 { + if rlen < len(c.response.Header()) { continue } diff --git a/q3rcon.go b/q3rcon.go index 3b6ca4a..16cf960 100644 --- a/q3rcon.go +++ b/q3rcon.go @@ -7,7 +7,6 @@ import ( "github.com/onyx-and-iris/q3rcon/internal/conn" "github.com/onyx-and-iris/q3rcon/internal/packet" - log "github.com/sirupsen/logrus" ) // Option is a functional option type that allows us to configure the VbanTxt. @@ -50,7 +49,7 @@ func New(host string, port int, password string, options ...Option) (*Rcon, erro r := &Rcon{ conn: conn, request: packet.NewRequest(password), - resp: make(chan string, 1), + resp: make(chan string), defaultTimeout: 20 * time.Millisecond, timeouts: make(map[string]time.Duration), } @@ -97,14 +96,19 @@ func (r Rcon) Send(cmd string) (string, error) { timeout = r.defaultTimeout } - go r.conn.Listen(timeout, r.resp) + e := make(chan error) + go r.conn.Listen(timeout, r.resp, e) _, err := r.conn.Write(r.request.Encode(cmd)) if err != nil { return "", err } - log.Tracef("Sending '%s'", cmd) - return strings.TrimPrefix(<-r.resp, string(r.response.Header())), nil + select { + case err := <-e: + return "", err + case resp := <-r.resp: + return strings.TrimPrefix(resp, string(r.response.Header())), nil + } } func (r Rcon) Close() {