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

// Package stun implements an unmodified RFC 5389 STUN binding responder with
// the RFC 5780 NAT-behavior-discovery attributes (OTHER-ADDRESS,
// RESPONSE-ORIGIN, CHANGE-REQUEST) when alternate addresses are available.
// No custom framing — interop with existing STUN tooling is a feature
// (spec §4). Each configured primary address gets two sockets: the given
// port and port+1 (the RFC 5780 alternate-port convention).
package stun

import (
	"crypto/rand"
	"encoding/binary"
	"log/slog"
	"net"
	"net/netip"
)

const (
	magicCookie = 0x2112A442

	typeBindingRequest = 0x0001
	typeBindingSuccess = 0x0101

	attrChangeRequest  = 0x0003
	attrXorMapped      = 0x0020
	attrSoftware       = 0x8022
	attrResponseOrigin = 0x802B
	attrOtherAddress   = 0x802C

	changeIP   = 0x04
	changePort = 0x02
)

// sock is one bound socket, addressable by (address index, port index).
type sock struct {
	conn *net.UDPConn
	addr netip.AddrPort
}

// Server holds the socket grid: addrs × {primary, alternate} ports.
type Server struct {
	// socks[i][0] = primary port, socks[i][1] = alt port for address i.
	socks [][2]*sock
}

// Listen binds primary+alternate sockets for every address. Addresses are
// "ip:port" specs; the alternate port is port+1.
func Listen(addrs []string) (*Server, error) {
	s := &Server{}
	for _, spec := range addrs {
		ap, err := netip.ParseAddrPort(spec)
		if err != nil {
			return nil, err
		}
		var pair [2]*sock
		for i, port := range []uint16{ap.Port(), ap.Port() + 1} {
			bind := netip.AddrPortFrom(ap.Addr(), port)
			conn, err := net.ListenUDP("udp", net.UDPAddrFromAddrPort(bind))
			if err != nil {
				s.Close()
				return nil, err
			}
			pair[i] = &sock{conn: conn, addr: bind}
		}
		s.socks = append(s.socks, pair)
	}
	return s, nil
}

func (s *Server) Close() {
	for _, pair := range s.socks {
		for _, sk := range pair {
			if sk != nil {
				sk.conn.Close()
			}
		}
	}
}

// Has5780 reports whether any address family has ≥2 addresses — the
// prerequisite for full NAT behavior discovery.
func (s *Server) Has5780() bool {
	var v4, v6 int
	for _, pair := range s.socks {
		if pair[0].addr.Addr().Is4() || pair[0].addr.Addr().Is4In6() {
			v4++
		} else {
			v6++
		}
	}
	return v4 >= 2 || v6 >= 2
}

// Serve starts one read loop per socket and blocks until the first error.
func (s *Server) Serve() error {
	errCh := make(chan error, len(s.socks)*2)
	for ai := range s.socks {
		for pi := range s.socks[ai] {
			go func(ai, pi int) { errCh <- s.loop(ai, pi) }(ai, pi)
		}
	}
	return <-errCh
}

func (s *Server) loop(ai, pi int) error {
	sk := s.socks[ai][pi]
	buf := make([]byte, 1500)
	for {
		n, raddr, err := sk.conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			return err
		}
		s.handle(ai, pi, buf[:n], raddr)
	}
}

// otherAddr finds the "diagonal" alternate for RFC 5780: different address
// (same family), different port. Returns nil when there is none.
func (s *Server) other(ai int, sameFamily bool, fam4 bool) int {
	for i, pair := range s.socks {
		if i == ai {
			continue
		}
		is4 := pair[0].addr.Addr().Is4() || pair[0].addr.Addr().Is4In6()
		if !sameFamily || is4 == fam4 {
			return i
		}
	}
	return -1
}

func (s *Server) handle(ai, pi int, pkt []byte, raddr netip.AddrPort) {
	if len(pkt) < 20 || binary.BigEndian.Uint16(pkt[0:2]) != typeBindingRequest {
		return
	}
	if binary.BigEndian.Uint32(pkt[4:8]) != magicCookie {
		return
	}
	msgLen := int(binary.BigEndian.Uint16(pkt[2:4]))
	if 20+msgLen > len(pkt) {
		return
	}
	var txid [12]byte
	copy(txid[:], pkt[8:20])

	// Parse CHANGE-REQUEST if present (RFC 5780 §7.2).
	var change byte
	for off := 20; off+4 <= 20+msgLen; {
		at := binary.BigEndian.Uint16(pkt[off : off+2])
		al := int(binary.BigEndian.Uint16(pkt[off+2 : off+4]))
		if off+4+al > len(pkt) {
			break
		}
		if at == attrChangeRequest && al >= 4 {
			change = pkt[off+7]
		}
		off += 4 + al + (4-al%4)%4 // attributes are 32-bit aligned
	}

	// Pick the responding socket per CHANGE-REQUEST.
	fam4 := raddr.Addr().Is4() || raddr.Addr().Is4In6()
	rai, rpi := ai, pi
	if change&changeIP != 0 {
		if o := s.other(ai, true, fam4); o >= 0 {
			rai = o
		} else {
			return // cannot honor — RFC says error response; silence is safer for a probe target
		}
	}
	if change&changePort != 0 {
		rpi = 1 - pi
	}
	responder := s.socks[rai][rpi]

	resp := buildResponse(txid, raddr, responder.addr, s.otherAddress(ai, fam4))
	if _, err := responder.conn.WriteToUDPAddrPort(resp, raddr); err != nil {
		slog.Debug("stun write failed", "to", raddr, "err", err)
	}
}

// otherAddress computes the OTHER-ADDRESS attribute value (alt IP, alt port)
// for the client's family, or an invalid AddrPort when unavailable.
func (s *Server) otherAddress(ai int, fam4 bool) netip.AddrPort {
	if o := s.other(ai, true, fam4); o >= 0 {
		return s.socks[o][1].addr
	}
	return netip.AddrPort{}
}

func buildResponse(txid [12]byte, mapped, origin, other netip.AddrPort) []byte {
	attrs := xorMappedAttr(attrXorMapped, mapped, txid)
	attrs = append(attrs, addrAttr(attrResponseOrigin, origin)...)
	if other.IsValid() {
		attrs = append(attrs, addrAttr(attrOtherAddress, other)...)
	}
	sw := []byte("echolot")
	attrs = append(attrs, attrHeader(attrSoftware, len(sw))...)
	attrs = append(attrs, pad4(sw)...)

	msg := make([]byte, 20, 20+len(attrs))
	binary.BigEndian.PutUint16(msg[0:2], typeBindingSuccess)
	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...)
}

func attrHeader(typ uint16, valLen int) []byte {
	h := make([]byte, 4)
	binary.BigEndian.PutUint16(h[0:2], typ)
	binary.BigEndian.PutUint16(h[2:4], uint16(valLen))
	return h
}

func pad4(b []byte) []byte {
	for len(b)%4 != 0 {
		b = append(b, 0)
	}
	return b
}

// addrValue encodes the RFC 5389 address structure (family, port, address).
func addrValue(ap netip.AddrPort) []byte {
	addr := ap.Addr().Unmap()
	if addr.Is4() {
		v := make([]byte, 8)
		v[1] = 0x01
		binary.BigEndian.PutUint16(v[2:4], ap.Port())
		a4 := addr.As4()
		copy(v[4:], a4[:])
		return v
	}
	v := make([]byte, 20)
	v[1] = 0x02
	binary.BigEndian.PutUint16(v[2:4], ap.Port())
	a16 := addr.As16()
	copy(v[4:], a16[:])
	return v
}

func addrAttr(typ uint16, ap netip.AddrPort) []byte {
	v := addrValue(ap)
	return append(attrHeader(typ, len(v)), v...)
}

// xorMappedAttr encodes XOR-MAPPED-ADDRESS (RFC 5389 §15.2).
func xorMappedAttr(typ uint16, ap netip.AddrPort, txid [12]byte) []byte {
	v := addrValue(ap)
	binary.BigEndian.PutUint16(v[2:4], ap.Port()^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]
	}
	return append(attrHeader(typ, len(v)), v...)
}

// NewTxID is exported for tests and client code.
func NewTxID() [12]byte {
	var t [12]byte
	_, _ = rand.Read(t[:])
	return t
}