gamestatus/source/source.go

package source

import (
	"encoding/json"
	"net"
	"strconv"
	"strings"
	"time"

	a2s "github.com/rumblefrog/go-a2s"

	"github.com/VexoraDevelopment/gamestatus/info"
)

type Options struct {
	Timeout time.Duration
}

type Client struct {
	opt Options
}

func New(opt Options) *Client {
	if opt.Timeout <= 0 {
		opt.Timeout = 2 * time.Second
	}
	return &Client{opt: opt}
}

func (c *Client) Connect(address string, port int) (*info.ServerInfo, error) {
	return c.query(address, port, false)
}

func (c *Client) QueryFull(address string, port int) (*info.ServerInfo, error) {
	return c.query(address, port, true)
}

func (c *Client) ConnectRust(address string, port int) (*info.ServerInfo, error) {
	out, err := c.query(address, port, false)
	if err != nil {
		return nil, err
	}
	out.Edition = info.EditionRust
	return out, nil
}

func (c *Client) QueryFullRust(address string, port int) (*info.ServerInfo, error) {
	out, err := c.query(address, port, true)
	if err != nil {
		return nil, err
	}
	out.Edition = info.EditionRust
	return out, nil
}

func (c *Client) ConnectCS2(address string, port int) (*info.ServerInfo, error) {
	out, err := c.query(address, port, false)
	if err != nil {
		return nil, err
	}
	out.Edition = info.EditionCS2
	return out, nil
}

func (c *Client) QueryFullCS2(address string, port int) (*info.ServerInfo, error) {
	out, err := c.query(address, port, true)
	if err != nil {
		return nil, err
	}
	out.Edition = info.EditionCS2
	return out, nil
}

func (c *Client) ConnectValve(address string, port int) (*info.ServerInfo, error) {
	out, err := c.query(address, port, false)
	if err != nil {
		return nil, err
	}
	out.Edition = info.EditionValve
	return out, nil
}

func (c *Client) QueryFullValve(address string, port int) (*info.ServerInfo, error) {
	out, err := c.query(address, port, true)
	if err != nil {
		return nil, err
	}
	out.Edition = info.EditionValve
	return out, nil
}

func (c *Client) query(address string, port int, full bool) (*info.ServerInfo, error) {
	addr := normalizeAddress(address, port)
	client, err := a2s.NewClient(addr, a2s.TimeoutOption(c.opt.Timeout))
	if err != nil {
		return nil, err
	}
	defer client.Close()

	start := time.Now()
	infoResp, err := client.QueryInfo()
	if err != nil {
		return nil, err
	}

	out := &info.ServerInfo{
		Edition:  info.EditionSource,
		Address:  address,
		IP:       resolveIP(address),
		Port:     effectivePort(address, port),
		MOTD:     infoResp.Name,
		Version:  infoResp.Version,
		Online:   int(infoResp.Players),
		Slots:    int(infoResp.MaxPlayers),
		Software: infoResp.Game,
		Game:     infoResp.Folder,
		Map:      infoResp.Map,
		Query:    full,
		Latency:  time.Since(start),
	}

	if full {
		playerResp, err := client.QueryPlayer()
		if err == nil && playerResp != nil {
			names := make([]string, 0, len(playerResp.Players))
			for _, p := range playerResp.Players {
				if p == nil {
					continue
				}
				name := strings.TrimSpace(p.Name)
				if name != "" {
					names = append(names, name)
				}
			}
			out.Players = names
		}
		rulesResp, err := client.QueryRules()
		if err == nil && rulesResp != nil {
			out.RawQuery = rulesResp.Rules
		}
	}

	raw, _ := json.Marshal(infoResp)
	out.Raw = raw
	return out, nil
}

func normalizeAddress(host string, port int) string {
	if port == 0 {
		port = a2s.DefaultPort
	}
	if strings.Contains(host, ":") {
		if _, _, err := net.SplitHostPort(host); err == nil {
			return host
		}
		if strings.Count(host, ":") > 1 && !strings.HasPrefix(host, "[") {
			return net.JoinHostPort(host, strconv.Itoa(port))
		}
	}
	return net.JoinHostPort(host, strconv.Itoa(port))
}

func effectivePort(host string, port int) int {
	if port != 0 {
		return port
	}
	if _, p, err := net.SplitHostPort(host); err == nil {
		if v, convErr := strconv.Atoi(p); convErr == nil {
			return v
		}
	}
	return a2s.DefaultPort
}

func resolveIP(host string) string {
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	host = strings.TrimPrefix(host, "[")
	host = strings.TrimSuffix(host, "]")
	if ip := net.ParseIP(host); ip != nil {
		return ip.String()
	}
	ips, err := net.LookupIP(host)
	if err != nil || len(ips) == 0 {
		return ""
	}
	return ips[0].String()
}