package blacktip

import (
	"encoding/binary"
	"fmt"
	"net"
	"net/netip"
)

// PROXY protocol (HAProxy spec): a short header telling the target who the real
// client was. Without it a service behind the tunnel only ever sees the tunnel,
// so IP bans, rate limits, spam checks and logs all see the wrong address.
// Servers must be told to expect it and to trust only the edge/site that sends it.

var proxyV2Sig = []byte("\r\n\r\n\x00\r\nQUIT\n")

func addrPort(a net.Addr) (netip.Addr, uint16, bool) {
	ta, ok := a.(*net.TCPAddr)
	if !ok {
		ap, err := netip.ParseAddrPort(a.String())
		if err != nil {
			return netip.Addr{}, 0, false
		}
		return ap.Addr().Unmap(), ap.Port(), true
	}
	ip, ok := netip.AddrFromSlice(ta.IP)
	if !ok {
		return netip.Addr{}, 0, false
	}
	return ip.Unmap(), uint16(ta.Port), true
}

// proxyHeader builds a PROXY protocol header ("v1" or "v2") for a TCP connection from src to dst.
func proxyHeader(version string, src, dst net.Addr) []byte {
	sip, sport, ok1 := addrPort(src)
	dip, dport, ok2 := addrPort(dst)
	known := ok1 && ok2 && sip.Is4() == dip.Is4()
	if version == "v2" {
		h := append([]byte{}, proxyV2Sig...)
		if !known {
			return append(h, 0x20, 0x00, 0x00, 0x00) // LOCAL: no address information
		}
		if sip.Is4() {
			h = append(h, 0x21, 0x11, 0x00, 12)
			s4, d4 := sip.As4(), dip.As4()
			h = append(h, s4[:]...)
			h = append(h, d4[:]...)
		} else {
			h = append(h, 0x21, 0x21, 0x00, 36)
			s6, d6 := sip.As16(), dip.As16()
			h = append(h, s6[:]...)
			h = append(h, d6[:]...)
		}
		h = binary.BigEndian.AppendUint16(h, sport)
		return binary.BigEndian.AppendUint16(h, dport)
	}
	if !known {
		return []byte("PROXY UNKNOWN\r\n")
	}
	fam := "TCP4"
	if !sip.Is4() {
		fam = "TCP6"
	}
	return []byte(fmt.Sprintf("PROXY %s %s %s %d %d\r\n", fam, sip, dip, sport, dport))
}