go-ts/pkg/transport/socket.go

package transport

import (
	"fmt"
	"net"
	"sync"
	"time"

	"go-ts/pkg/protocol"
)

// TS3Conn handles the UDP connection to the TeamSpeak server.
type TS3Conn struct {
	conn      *net.UDPConn
	readQueue chan *protocol.Packet
	closeChan chan struct{}
	wg        sync.WaitGroup
	writeMu   sync.Mutex
}

// NewTS3Conn creates a new connection to the specified address.
func NewTS3Conn(address string) (*TS3Conn, error) {
	raddr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve address: %w", err)
	}

	conn, err := net.DialUDP("udp", nil, raddr)
	if err != nil {
		return nil, fmt.Errorf("failed to dial UDP: %w", err)
	}

	ts3c := &TS3Conn{
		conn:      conn,
		readQueue: make(chan *protocol.Packet, 100),
		closeChan: make(chan struct{}),
	}

	ts3c.startReader()

	return ts3c, nil
}

func (c *TS3Conn) startReader() {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		buf := make([]byte, 2048) // Max packet size should be around 500, but use larger buffer to be safe
		for {
			select {
			case <-c.closeChan:
				return
			default:
				c.conn.SetReadDeadline(time.Now().Add(1 * time.Second))
				n, _, err := c.conn.ReadFromUDP(buf)
				if err != nil {
					if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
						continue
					}
					// If closed, return
					select {
					case <-c.closeChan:
						return
					default:
						fmt.Printf("Error reading valid UDP packet: %v\n", err)
						continue
					}
				}

				// Parse packet
				pkt, err := protocol.Decode(buf[:n], false) // Server -> Client
				if err != nil {
					fmt.Printf("Failed to decode packet: %v\n", err)
					continue
				}

				select {
				case c.readQueue <- pkt:
				default:
					fmt.Println("Read queue full, dropping packet")
				}
			}
		}
	}()
}

// SendPacket sends a packet to the server.
func (c *TS3Conn) SendPacket(pkt *protocol.Packet) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	// Client -> Server
	bytes, err := pkt.Encode(true)
	if err != nil {
		return err
	}

	_, err = c.conn.Write(bytes)
	return err
}

// ReadPacket returns the next received packet channel.
func (c *TS3Conn) PacketChan() <-chan *protocol.Packet {
	return c.readQueue
}

// RemoteAddr returns the remote network address.
func (c *TS3Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// Close closes the connection.
func (c *TS3Conn) Close() error {
	close(c.closeChan)
	c.conn.Close()
	c.wg.Wait()
	return nil
}