package proxyproto

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"net"
)

// maxV2HeaderSize is the maximum acceptable size of a V2 header.
//
// A V2 header may be at most 16 bytes + 64KiB large. We enforce a lower limit
// to mitigate memory allocation DoS while allowing real-world legitimate
// headers. PP2_SUBTYPE_SSL_CLIENT_CERT is typically between 1 and 2KiB, so we
// use a 4KiB limit to leave some room for other TLVs.
const maxV2HeaderSize = 4096

var (
	lengthUnspec      = uint16(0)
	lengthV4          = uint16(12)
	lengthV6          = uint16(36)
	lengthUnix        = uint16(216)
	lengthUnspecBytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, lengthUnspec)
		return a
	}()
	lengthV4Bytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, lengthV4)
		return a
	}()
	lengthV6Bytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, lengthV6)
		return a
	}()
	lengthUnixBytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, lengthUnix)
		return a
	}()
	errUint16Overflow = errors.New("proxyproto: uint16 overflow")
)

type _ports struct {
	SrcPort uint16
	DstPort uint16
}

type _addr4 struct {
	Src     [4]byte
	Dst     [4]byte
	SrcPort uint16
	DstPort uint16
}

type _addr6 struct {
	Src [16]byte
	Dst [16]byte
	_ports
}

type _addrUnix struct {
	Src [108]byte
	Dst [108]byte
}

func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
	// Skip first 12 bytes (signature)
	for range 12 {
		if _, err = reader.ReadByte(); err != nil {
			return nil, fmt.Errorf("%w: %w", ErrCantReadProtocolVersionAndCommand, err)
		}
	}

	header = new(Header)
	header.Version = 2

	// Read the 13th byte, protocol version and command
	b13, err := reader.ReadByte()
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrCantReadProtocolVersionAndCommand, err)
	}
	header.Command = ProtocolVersionAndCommand(b13)
	if _, ok := supportedCommand[header.Command]; !ok {
		return nil, ErrUnsupportedProtocolVersionAndCommand
	}

	// Read the 14th byte, address family and protocol
	b14, err := reader.ReadByte()
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrCantReadAddressFamilyAndProtocol, err)
	}
	header.TransportProtocol = AddressFamilyAndProtocol(b14)
	// UNSPEC is only supported when LOCAL is set.
	if header.TransportProtocol == UNSPEC && header.Command != LOCAL {
		return nil, ErrUnsupportedAddressFamilyAndProtocol
	}

	// Make sure there are bytes available as specified in length
	var length uint16
	if err := binary.Read(reader, binary.BigEndian, &length); err != nil {
		return nil, fmt.Errorf("%w: %w", ErrCantReadLength, err)
	}
	if !header.validateLength(length) {
		return nil, ErrInvalidLength
	}

	// Return early if the length is zero, which means that
	// there's no address information and TLVs present for UNSPEC.
	if length == 0 {
		return header, nil
	}

	if length > maxV2HeaderSize {
		return nil, ErrInvalidLength
	}

	// Length-limited reader for payload section
	payloadReader := io.LimitReader(reader, int64(length)).(*io.LimitedReader)

	// Read addresses and ports for protocols other than UNSPEC.
	// Ignore address information for UNSPEC, and skip straight to read TLVs,
	// since the length is greater than zero.
	if header.TransportProtocol != UNSPEC {
		if header.TransportProtocol.IsIPv4() {
			var addr _addr4
			if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
				return nil, fmt.Errorf("%w: %w", ErrInvalidAddress, err)
			}
			header.SourceAddr = newIPAddr(header.TransportProtocol, addr.Src[:], addr.SrcPort)
			header.DestinationAddr = newIPAddr(header.TransportProtocol, addr.Dst[:], addr.DstPort)
		} else if header.TransportProtocol.IsIPv6() {
			var addr _addr6
			if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
				return nil, fmt.Errorf("%w: %w", ErrInvalidAddress, err)
			}
			header.SourceAddr = newIPAddr(header.TransportProtocol, addr.Src[:], addr.SrcPort)
			header.DestinationAddr = newIPAddr(header.TransportProtocol, addr.Dst[:], addr.DstPort)
		} else if header.TransportProtocol.IsUnix() {
			var addr _addrUnix
			if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
				return nil, fmt.Errorf("%w: %w", ErrInvalidAddress, err)
			}

			network := "unix"
			if header.TransportProtocol.IsDatagram() {
				network = "unixgram"
			}

			header.SourceAddr = &net.UnixAddr{
				Net:  network,
				Name: parseUnixName(addr.Src[:]),
			}
			header.DestinationAddr = &net.UnixAddr{
				Net:  network,
				Name: parseUnixName(addr.Dst[:]),
			}
		}
	}

	// Copy bytes for optional Type-Length-Value vector
	header.rawTLVs = make([]byte, payloadReader.N) // Allocate minimum size slice
	if _, err = io.ReadFull(payloadReader, header.rawTLVs); err != nil && err != io.EOF {
		return nil, err
	}

	if payloadReader.N != 0 {
		return nil, ErrInvalidLength
	}

	return header, nil
}

func (header *Header) formatVersion2() ([]byte, error) {
	var buf bytes.Buffer
	buf.Write(SIGV2)
	buf.WriteByte(header.Command.toByte())
	buf.WriteByte(header.TransportProtocol.toByte())
	if header.TransportProtocol.IsUnspec() {
		// For UNSPEC, write no addresses and ports but only TLVs if they are present
		hdrLen, err := addTLVLen(lengthUnspecBytes, len(header.rawTLVs))
		if err != nil {
			return nil, err
		}
		buf.Write(hdrLen)
	} else {
		var addrSrc, addrDst []byte
		if header.TransportProtocol.IsIPv4() {
			hdrLen, err := addTLVLen(lengthV4Bytes, len(header.rawTLVs))
			if err != nil {
				return nil, err
			}
			buf.Write(hdrLen)
			sourceIP, destIP, _ := header.IPs()
			addrSrc = sourceIP.To4()
			addrDst = destIP.To4()
		} else if header.TransportProtocol.IsIPv6() {
			hdrLen, err := addTLVLen(lengthV6Bytes, len(header.rawTLVs))
			if err != nil {
				return nil, err
			}
			buf.Write(hdrLen)
			sourceIP, destIP, _ := header.IPs()
			addrSrc = sourceIP.To16()
			addrDst = destIP.To16()
		} else if header.TransportProtocol.IsUnix() {
			buf.Write(lengthUnixBytes)
			sourceAddr, destAddr, ok := header.UnixAddrs()
			if !ok {
				return nil, ErrInvalidAddress
			}
			addrSrc = formatUnixName(sourceAddr.Name)
			addrDst = formatUnixName(destAddr.Name)
		}

		if addrSrc == nil || addrDst == nil {
			return nil, ErrInvalidAddress
		}
		buf.Write(addrSrc)
		buf.Write(addrDst)

		if sourcePort, destPort, ok := header.Ports(); ok {
			if sourcePort < 0 || sourcePort > math.MaxUint16 || destPort < 0 || destPort > math.MaxUint16 {
				return nil, ErrInvalidPortNumber
			}
			portBytes := make([]byte, 2)

			//nolint:gosec // Bounds are checked above.
			binary.BigEndian.PutUint16(portBytes, uint16(sourcePort))
			buf.Write(portBytes)

			//nolint:gosec // Bounds are checked above.
			binary.BigEndian.PutUint16(portBytes, uint16(destPort))
			buf.Write(portBytes)
		}
	}

	if len(header.rawTLVs) > 0 {
		buf.Write(header.rawTLVs)
	}

	return buf.Bytes(), nil
}

func (header *Header) validateLength(length uint16) bool {
	if header.TransportProtocol.IsIPv4() {
		return length >= lengthV4
	} else if header.TransportProtocol.IsIPv6() {
		return length >= lengthV6
	} else if header.TransportProtocol.IsUnix() {
		return length >= lengthUnix
	} else if header.TransportProtocol.IsUnspec() {
		return length >= lengthUnspec
	}
	return false
}

// addTLVLen adds the length of the TLV to the header length or errors on uint16 overflow.
func addTLVLen(cur []byte, tlvLen int) ([]byte, error) {
	if tlvLen == 0 {
		return cur, nil
	}
	curLen := binary.BigEndian.Uint16(cur)
	newLen := int(curLen) + tlvLen
	if newLen >= 1<<16 {
		return nil, errUint16Overflow
	}
	a := make([]byte, 2)
	//nolint:gosec // newLen bounds are validated above.
	binary.BigEndian.PutUint16(a, uint16(newLen))
	return a, nil
}

func newIPAddr(transport AddressFamilyAndProtocol, ip net.IP, port uint16) net.Addr {
	if transport.IsStream() {
		return &net.TCPAddr{IP: ip, Port: int(port)}
	}
	if transport.IsDatagram() {
		return &net.UDPAddr{IP: ip, Port: int(port)}
	}
	return nil
}

func parseUnixName(b []byte) string {
	before, _, ok := bytes.Cut(b, []byte{0})
	if !ok {
		return string(b)
	}
	return string(before)
}

func formatUnixName(name string) []byte {
	n := int(lengthUnix) / 2
	if len(name) >= n {
		return []byte(name[:n])
	}
	pad := make([]byte, n-len(name))
	return append([]byte(name), pad...)
}