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) { udpAddr, err := net.ResolveUDPAddr("udp4", fmt.Sprintf("%s:%d", host, port)) if err != nil { return UDPConn{}, err } conn, err := net.DialUDP("udp4", nil, udpAddr) if err != nil { return UDPConn{}, err } log.Infof("Outgoing address %s", conn.RemoteAddr()) return UDPConn{ conn: conn, response: packet.NewResponse(), }, nil } func (c UDPConn) Write(buf []byte) (int, error) { n, err := c.conn.Write(buf) if err != nil { return 0, err } return n, nil } func (c UDPConn) Listen(timeout time.Duration, resp chan<- string) { c.conn.SetReadDeadline(time.Now().Add(timeout)) ch := make(chan struct{}) var sb strings.Builder buf := make([]byte, 2048) for { select { case <-ch: 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(ch) } else { log.Error(e) } } } if rlen == 0 { continue } if bytes.HasPrefix(buf, c.response.Header()) { sb.Write(buf[len(c.response.Header()):rlen]) } } } } func (c UDPConn) Close() error { err := c.conn.Close() if err != nil { return err } return nil }