zlh-agent/internal/minecraft/status.go

package minecraft

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"strings"
	"time"
)

type StatusResponse struct {
	Version struct {
		Name     string `json:"name"`
		Protocol int    `json:"protocol"`
	} `json:"version"`
	Players struct {
		Max    int `json:"max"`
		Online int `json:"online"`
		Sample []struct {
			Name string `json:"name"`
			ID   string `json:"id"`
		} `json:"sample"`
	} `json:"players"`
}

func ProtocolForVersion(version string) int {
	v := strings.TrimSpace(strings.TrimPrefix(version, "v"))
	switch v {
	case "1.21.1":
		return 767
	case "1.21":
		return 767
	case "1.20.4":
		return 765
	case "1.20.1":
		return 763
	case "1.19.4":
		return 762
	default:
		return 754
	}
}

func QueryStatus(host string, port int, protocol int) (StatusResponse, error) {
	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.DialTimeout("tcp", addr, 3*time.Second)
	if err != nil {
		return StatusResponse{}, err
	}
	defer conn.Close()

	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))

	if err := writeHandshake(conn, host, port, protocol); err != nil {
		return StatusResponse{}, err
	}
	if err := writeStatusRequest(conn); err != nil {
		return StatusResponse{}, err
	}

	payload, err := readPacket(conn)
	if err != nil {
		return StatusResponse{}, err
	}

	r := bytes.NewReader(payload)
	packetID, err := readVarInt(r)
	if err != nil {
		return StatusResponse{}, err
	}
	if packetID != 0x00 {
		return StatusResponse{}, fmt.Errorf("unexpected packet id: %d", packetID)
	}

	respStr, err := readString(r)
	if err != nil {
		return StatusResponse{}, err
	}

	var status StatusResponse
	if err := json.Unmarshal([]byte(respStr), &status); err != nil {
		return StatusResponse{}, err
	}
	return status, nil
}

func writeHandshake(w io.Writer, host string, port int, protocol int) error {
	var payload bytes.Buffer
	if err := writeVarInt(&payload, 0x00); err != nil {
		return err
	}
	if err := writeVarInt(&payload, protocol); err != nil {
		return err
	}
	if err := writeString(&payload, host); err != nil {
		return err
	}
	if err := binary.Write(&payload, binary.BigEndian, uint16(port)); err != nil {
		return err
	}
	if err := writeVarInt(&payload, 0x01); err != nil {
		return err
	}

	return writePacket(w, payload.Bytes())
}

func writeStatusRequest(w io.Writer) error {
	var payload bytes.Buffer
	if err := writeVarInt(&payload, 0x00); err != nil {
		return err
	}
	return writePacket(w, payload.Bytes())
}

func writePacket(w io.Writer, payload []byte) error {
	var buf bytes.Buffer
	if err := writeVarInt(&buf, len(payload)); err != nil {
		return err
	}
	if _, err := buf.Write(payload); err != nil {
		return err
	}
	_, err := w.Write(buf.Bytes())
	return err
}

func readPacket(r io.Reader) ([]byte, error) {
	length, err := readVarInt(r)
	if err != nil {
		return nil, err
	}
	if length <= 0 || length > 1<<20 {
		return nil, fmt.Errorf("invalid packet length: %d", length)
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}
	return buf, nil
}

func writeVarInt(w io.Writer, value int) error {
	for {
		temp := byte(value & 0x7F)
		value >>= 7
		if value != 0 {
			temp |= 0x80
		}
		if _, err := w.Write([]byte{temp}); err != nil {
			return err
		}
		if value == 0 {
			break
		}
	}
	return nil
}

func readVarInt(r io.Reader) (int, error) {
	numRead := 0
	result := 0
	for {
		if numRead > 5 {
			return 0, fmt.Errorf("varint too long")
		}
		b := make([]byte, 1)
		if _, err := r.Read(b); err != nil {
			return 0, err
		}
		value := int(b[0] & 0x7F)
		result |= value << (7 * numRead)
		numRead++
		if b[0]&0x80 == 0 {
			break
		}
	}
	return result, nil
}

func writeString(w io.Writer, s string) error {
	if err := writeVarInt(w, len(s)); err != nil {
		return err
	}
	_, err := w.Write([]byte(s))
	return err
}

func readString(r io.Reader) (string, error) {
	length, err := readVarInt(r)
	if err != nil {
		return "", err
	}
	if length < 0 || length > 1<<20 {
		return "", fmt.Errorf("invalid string length: %d", length)
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}