reforgerds-updater/internal/a2s/client.go

package a2s

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"time"
)

// Errors
var (
	ErrPlayerRead          = errors.New("failed to read player data")
	ErrMultiPacketInvalid  = errors.New("invalid multi-packet ID mismatch")
	ErrMultiPacketMismatch = errors.New("multi-packet assembly failed")
	errBzip2               = errors.New("bzip2 compressed response not supported")
)

// Flag type for request/response types
type Flag byte

const (
	PlayerRequest    Flag = 0x55
	ChallengeResponse Flag = 0x41

	DefaultBufferSize       = 1400
	DefaultDeadlineTimeout  = 5
	singlePacket     uint32 = 0xFFFFFFFF
)

// Client handles connection and options
type Client struct {
	Conn       *net.UDPConn
	Address    *net.UDPAddr
	Timeout    time.Duration
	BufferSize uint16
}

// New creates a new Client and dials the connection
func New(ip string, port int) (*Client, error) {
	return NewWithAddr(&net.UDPAddr{IP: net.ParseIP(ip), Port: port})
}

func NewWithString(addr string) (*Client, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	return NewWithAddr(udpAddr)
}

func NewWithAddr(addr *net.UDPAddr) (*Client, error) {
	client := &Client{
		Address:    addr,
		Timeout:    DefaultDeadlineTimeout * time.Second,
		BufferSize: DefaultBufferSize,
	}
	return client, client.Dial()
}

func (c *Client) Dial() error {
	conn, err := net.DialUDP("udp", nil, c.Address)
	if err != nil {
		return err
	}
	c.Conn = conn
	return nil
}

func (c *Client) Close() error {
	return c.Conn.Close()
}

func (c *Client) SetBufferSize(size uint16) {
	c.BufferSize = size
}

func (c *Client) SetDeadlineTimeout(seconds int) {
	c.Timeout = time.Duration(seconds) * time.Second
}

func (c *Client) Get(requestType Flag) ([]byte, Flag, time.Duration, error) {
	resp, duration, err := c.request(requestType, singlePacket)
	if err != nil {
		return nil, 0, 0, err
	}
	flag := Flag(resp[4])

	if flag == ChallengeResponse {
		challenge := binary.BigEndian.Uint32(resp[5:9])
		resp, _, err = c.request(requestType, challenge)
		if err != nil {
			return nil, 0, 0, err
		}
		flag = Flag(resp[4])
	}

	if err := validateResponseType(requestType, flag); err != nil {
		return resp[5:], flag, duration, err
	}

	return resp[5:], flag, duration, nil
}

func (c *Client) request(requestType Flag, challenge uint32) ([]byte, time.Duration, error) {
	req, err := createHeader(requestType, challenge)
	if err != nil {
		return nil, 0, err
	}

	start := time.Now()

	if _, err := c.Conn.Write(req); err != nil {
		return nil, 0, err
	}
	if err := c.Conn.SetReadDeadline(time.Now().Add(c.Timeout)); err != nil {
		return nil, 0, err
	}

	resp := make([]byte, c.BufferSize)
	n, err := c.Conn.Read(resp)
	if err != nil {
		return nil, 0, err
	}

	duration := time.Since(start)

	multi, err := isMultiPacket(resp)
	if err != nil {
		return resp, 0, err
	}

	if !multi {
		return resp[:n], duration, nil
	}

	packetID := binary.LittleEndian.Uint32(resp[4:8])
	packetCount := int(resp[8] & 0x0F)
	currentPacket := int(resp[9] & 0x0F)

	if (packetID & 0x80000000) != 0 {
		return nil, 0, errBzip2
	}

	packets := make(map[int][]byte)
	packets[currentPacket] = resp[12:n]

	for len(packets) < packetCount {
		buf := make([]byte, c.BufferSize)
		n, err := c.Conn.Read(buf)
		if err != nil {
			return nil, 0, err
		}
		if binary.LittleEndian.Uint32(buf[4:8]) != packetID {
			return nil, 0, ErrMultiPacketInvalid
		}

		currentPacket = int(buf[9] & 0x0F)
		if _, exists := packets[currentPacket]; !exists {
			packets[currentPacket] = buf[12:n]
		}
	}

	var assembledResp []byte
	for i := 0; i < packetCount; i++ {
		data, exists := packets[i]
		if !exists {
			return nil, 0, ErrMultiPacketMismatch
		}
		assembledResp = append(assembledResp, data...)
	}

	return assembledResp, duration, nil
}

// Helpers

func createHeader(requestType Flag, challenge uint32) ([]byte, error) {
	header := []byte{0xFF, 0xFF, 0xFF, 0xFF, byte(requestType)}
	if challenge != singlePacket {
		challengeBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(challengeBytes, challenge)
		header = append(header, challengeBytes...)
	}
	return header, nil
}

func validateResponseType(request, response Flag) error {
	if request == PlayerRequest && response != 0x44 {
		return fmt.Errorf("unexpected player response flag: 0x%X", response)
	}
	return nil
}

func isMultiPacket(buf []byte) (bool, error) {
	if len(buf) < 4 {
		return false, errors.New("packet too short")
	}
	header := binary.LittleEndian.Uint32(buf[:4])
	return header == 0xFFFFFFFE, nil
}