// Package checker — STUN binding-request codec.
//
// Hand-rolled RFC 5389 binding request encoder + binding success response
// parser. Just enough to extract XOR-MAPPED-ADDRESS — no message integrity,
// no fingerprint, no ALTERNATE-SERVER, no STUN-USE-CANDIDATE. Used after
// socks5UDPAssociate succeeds to verify the relay actually forwards UDP
// to the public Internet.
//
// We deliberately avoid pulling in pion/stun: ~80 LOC of encoding/binary,
// one attribute type. Adding ~50 KB of compiled code + a transitive
// dependency for one round trip is overkill.
package checker

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
)

// Magic cookie defined by RFC 5389 §6.
const stunMagicCookie = 0x2112A442

// STUN message types and attribute types we care about.
const (
	stunBindingRequest         = 0x0001
	stunBindingSuccessResponse = 0x0101
	stunAttrXORMappedAddress   = 0x0020
	stunAddressFamilyIPv4      = 0x01
	stunAddressFamilyIPv6      = 0x02
)

// Sentinel errors so HintFor (and tests) can match specific failure modes.
var (
	ErrSTUNTooShort          = errors.New("stun: response shorter than 20-byte header")
	ErrSTUNBadMagicCookie    = errors.New("stun: magic cookie mismatch")
	ErrSTUNNotSuccess        = errors.New("stun: response is not a Binding Success Response")
	ErrSTUNTxIDMismatch      = errors.New("stun: transaction ID mismatch")
	ErrSTUNNoMappedAddress   = errors.New("stun: response has no XOR-MAPPED-ADDRESS attribute")
	ErrSTUNUnsupportedFamily = errors.New("stun: unsupported XOR-MAPPED-ADDRESS family")
)

// NewTransactionID returns 12 cryptographically-random bytes suitable for
// use as a STUN transaction ID. Errors only if rand.Reader fails — caller
// should propagate (the runtime is in trouble at that point anyway).
func NewTransactionID() ([12]byte, error) {
	var id [12]byte
	if _, err := rand.Read(id[:]); err != nil {
		return id, fmt.Errorf("stun: read random transaction id: %w", err)
	}
	return id, nil
}

// EncodeBindingRequest builds a 20-byte STUN Binding Request with the given
// transaction ID. RFC 5389 allows empty request bodies.
func EncodeBindingRequest(txID [12]byte) []byte {
	buf := make([]byte, 20)
	binary.BigEndian.PutUint16(buf[0:2], stunBindingRequest)
	binary.BigEndian.PutUint16(buf[2:4], 0) // attribute length
	binary.BigEndian.PutUint32(buf[4:8], stunMagicCookie)
	copy(buf[8:20], txID[:])
	return buf
}

// ParseBindingResponse decodes a STUN Binding Success Response and returns
// the public IP+port advertised in XOR-MAPPED-ADDRESS.
//
// Validates header (length, magic, message type, transaction ID), then
// walks the TLV attribute section and extracts the first XOR-MAPPED-ADDRESS
// attribute. Other attributes are skipped (per RFC 5389 §15 the
// "comprehension-optional" range is everything ≥ 0x8000; we skip every
// non-XOR-MAPPED-ADDRESS attribute regardless, since this is the only one
// we care about).
func ParseBindingResponse(buf []byte, expectedTxID [12]byte) (net.IP, uint16, error) {
	if len(buf) < 20 {
		return nil, 0, ErrSTUNTooShort
	}

	msgType := binary.BigEndian.Uint16(buf[0:2])
	attrLen := binary.BigEndian.Uint16(buf[2:4])
	cookie := binary.BigEndian.Uint32(buf[4:8])
	if cookie != stunMagicCookie {
		return nil, 0, ErrSTUNBadMagicCookie
	}
	if msgType != stunBindingSuccessResponse {
		return nil, 0, fmt.Errorf("%w: type=0x%04x", ErrSTUNNotSuccess, msgType)
	}

	var txID [12]byte
	copy(txID[:], buf[8:20])
	if txID != expectedTxID {
		return nil, 0, ErrSTUNTxIDMismatch
	}

	// Sanity check: attrLen must not exceed buffer.
	if int(attrLen) > len(buf)-20 {
		return nil, 0, fmt.Errorf("stun: attribute section length %d exceeds buffer (%d bytes after header)", attrLen, len(buf)-20)
	}

	attrs := buf[20 : 20+int(attrLen)]
	off := 0
	for off < len(attrs) {
		// Each attribute header is 4 bytes (type + length).
		if len(attrs)-off < 4 {
			return nil, 0, fmt.Errorf("stun: truncated attribute header at offset %d", off)
		}
		aType := binary.BigEndian.Uint16(attrs[off : off+2])
		aLen := binary.BigEndian.Uint16(attrs[off+2 : off+4])
		valStart := off + 4
		valEnd := valStart + int(aLen)
		if valEnd > len(attrs) {
			return nil, 0, fmt.Errorf("stun: attribute at offset %d claims length %d but only %d bytes remain", off, aLen, len(attrs)-valStart)
		}

		if aType == stunAttrXORMappedAddress {
			ip, port, err := parseXORMappedAddress(attrs[valStart:valEnd], txID)
			if err != nil {
				return nil, 0, err
			}
			return ip, port, nil
		}

		// Skip attribute including padding to next 4-byte boundary.
		paddedLen := (int(aLen) + 3) &^ 3
		next := valStart + paddedLen
		if next > len(attrs) {
			// Padding runs past end — treat as truncation.
			return nil, 0, fmt.Errorf("stun: attribute padding at offset %d runs past end", off)
		}
		off = next
	}

	return nil, 0, ErrSTUNNoMappedAddress
}

// parseXORMappedAddress decodes the value of an XOR-MAPPED-ADDRESS attribute.
//
// Layout (RFC 5389 §15.2):
//
//	0                   1                   2                   3
//	0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
//	+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//	|0 0 0 0 0 0 0 0|    Family     |         X-Port                |
//	+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
//	|                X-Address (Variable)
//	+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
func parseXORMappedAddress(val []byte, txID [12]byte) (net.IP, uint16, error) {
	if len(val) < 4 {
		return nil, 0, fmt.Errorf("stun: XOR-MAPPED-ADDRESS truncated (got %d bytes, need ≥4)", len(val))
	}
	family := val[1]
	xPort := binary.BigEndian.Uint16(val[2:4])
	port := xPort ^ uint16(stunMagicCookie>>16)

	switch family {
	case stunAddressFamilyIPv4:
		if len(val) < 8 {
			return nil, 0, fmt.Errorf("stun: XOR-MAPPED-ADDRESS IPv4 truncated (got %d bytes, need 8)", len(val))
		}
		xAddr := binary.BigEndian.Uint32(val[4:8])
		addr := xAddr ^ stunMagicCookie
		ip := make(net.IP, 4)
		binary.BigEndian.PutUint32(ip, addr)
		return ip, port, nil

	case stunAddressFamilyIPv6:
		if len(val) < 20 {
			return nil, 0, fmt.Errorf("stun: XOR-MAPPED-ADDRESS IPv6 truncated (got %d bytes, need 20)", len(val))
		}
		// XOR with magic_cookie || transaction_id (16 bytes total).
		var key [16]byte
		binary.BigEndian.PutUint32(key[0:4], stunMagicCookie)
		copy(key[4:16], txID[:])
		ip := make(net.IP, 16)
		for i := 0; i < 16; i++ {
			ip[i] = val[4+i] ^ key[i]
		}
		return ip, port, nil

	default:
		return nil, 0, fmt.Errorf("%w: family=0x%02x", ErrSTUNUnsupportedFamily, family)
	}
}