// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

package stun

import (
	"encoding/binary"
	"net"
	"net/netip"
	"testing"
	"time"
)

// bindingRequest builds a minimal RFC 5389 binding request.
func bindingRequest(txid [12]byte, change byte) []byte {
	var attrs []byte
	if change != 0 {
		attrs = append(attrs, attrHeader(attrChangeRequest, 4)...)
		attrs = append(attrs, 0, 0, 0, change)
	}
	msg := make([]byte, 20, 20+len(attrs))
	binary.BigEndian.PutUint16(msg[0:2], typeBindingRequest)
	binary.BigEndian.PutUint16(msg[2:4], uint16(len(attrs)))
	binary.BigEndian.PutUint32(msg[4:8], magicCookie)
	copy(msg[8:20], txid[:])
	return append(msg, attrs...)
}

// parseXorMapped extracts XOR-MAPPED-ADDRESS from a binding success.
func parseXorMapped(t *testing.T, resp []byte, txid [12]byte) netip.AddrPort {
	t.Helper()
	if binary.BigEndian.Uint16(resp[0:2]) != typeBindingSuccess {
		t.Fatalf("type = %#x, want binding success", resp[0:2])
	}
	msgLen := int(binary.BigEndian.Uint16(resp[2:4]))
	for off := 20; off+4 <= 20+msgLen; {
		at := binary.BigEndian.Uint16(resp[off : off+2])
		al := int(binary.BigEndian.Uint16(resp[off+2 : off+4]))
		if at == attrXorMapped {
			v := append([]byte(nil), resp[off+4:off+4+al]...)
			port := binary.BigEndian.Uint16(v[2:4]) ^ uint16(magicCookie>>16)
			var key [16]byte
			binary.BigEndian.PutUint32(key[0:4], magicCookie)
			copy(key[4:], txid[:])
			for i := 4; i < len(v); i++ {
				v[i] ^= key[i-4]
			}
			if v[1] == 0x01 {
				return netip.AddrPortFrom(netip.AddrFrom4([4]byte(v[4:8])), port)
			}
			return netip.AddrPortFrom(netip.AddrFrom16([16]byte(v[4:20])), port)
		}
		off += 4 + al + (4-al%4)%4
	}
	t.Fatal("no XOR-MAPPED-ADDRESS in response")
	return netip.AddrPort{}
}

func TestBindingAndChangePort(t *testing.T) {
	// Two loopback "addresses" is not possible portably, so exercise one
	// address (basic binding + change-port); the change-IP path needs the
	// two-address grid of a real deployment.
	srv, err := Listen([]string{"127.0.0.1:0"})
	if err != nil {
		t.Fatal(err)
	}
	// port 0 twice would collide at 1; rebind explicitly on free ports
	srv.Close()
	base := freePort(t)
	srv, err = Listen([]string{netip.AddrPortFrom(netip.MustParseAddr("127.0.0.1"), base).String()})
	if err != nil {
		t.Skipf("cannot bind %d/%d: %v", base, base+1, err)
	}
	defer srv.Close()
	go srv.Serve()

	client, err := net.DialUDP("udp", nil, srv.socks[0][0].conn.LocalAddr().(*net.UDPAddr))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	client.SetDeadline(time.Now().Add(2 * time.Second))

	txid := NewTxID()
	client.Write(bindingRequest(txid, 0))
	buf := make([]byte, 1500)
	n, err := client.Read(buf)
	if err != nil {
		t.Fatal(err)
	}
	mapped := parseXorMapped(t, buf[:n], txid)
	want := client.LocalAddr().(*net.UDPAddr).AddrPort()
	if mapped.Port() != want.Port() {
		t.Fatalf("mapped port %d, want %d", mapped.Port(), want.Port())
	}

	// CHANGE-REQUEST(port): response must come from the alternate port.
	// Dial-connected sockets drop packets from other sources, so use an
	// unconnected socket and inspect the reply's source.
	uc, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	if err != nil {
		t.Fatal(err)
	}
	defer uc.Close()
	uc.SetDeadline(time.Now().Add(2 * time.Second))
	txid2 := NewTxID()
	uc.WriteToUDPAddrPort(bindingRequest(txid2, changePort), srv.socks[0][0].addr)
	n, from, err := uc.ReadFromUDPAddrPort(buf)
	if err != nil {
		t.Fatal(err)
	}
	if from.Port() != srv.socks[0][1].addr.Port() {
		t.Fatalf("change-port reply came from %v, want alt port %d", from, srv.socks[0][1].addr.Port())
	}
	parseXorMapped(t, buf[:n], txid2)
}

func freePort(t *testing.T) uint16 {
	t.Helper()
	// Find two adjacent free ports for the primary/alternate pair.
	for tries := 0; tries < 20; tries++ {
		l, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
		if err != nil {
			continue
		}
		p := l.LocalAddr().(*net.UDPAddr).Port
		l.Close()
		l2, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: p + 1})
		if err != nil {
			continue
		}
		l2.Close()
		return uint16(p)
	}
	t.Skip("no adjacent free UDP ports found")
	return 0
}