// Package netutil holds small address helpers shared by the DNS server, the
// rate limiter and the HTTP layer.
package netutil

import (
	"net"
	"net/netip"
	"strings"
)

// PrefixSet answers "is this address in one of these networks?".
//
// It backs rate-limit exemptions, query-log ignore lists and trusted-proxy
// configuration. Invalid entries are skipped rather than rejected: these lists
// are conveniences, and the security-critical recursion ACL uses a stricter
// constructor that reports errors.
type PrefixSet struct {
	prefixes []netip.Prefix
}

// NewPrefixSet compiles a list of CIDR blocks or bare addresses.
func NewPrefixSet(in []string) *PrefixSet {
	s := &PrefixSet{}
	for _, v := range in {
		if p, ok := ParsePrefix(v); ok {
			s.prefixes = append(s.prefixes, p)
		}
	}
	return s
}

// ParsePrefix accepts a CIDR block or a bare address, returning a masked
// prefix. A bare address becomes a host route.
func ParsePrefix(v string) (netip.Prefix, bool) {
	v = strings.TrimSpace(v)
	if v == "" {
		return netip.Prefix{}, false
	}
	if strings.Contains(v, "/") {
		p, err := netip.ParsePrefix(v)
		if err != nil {
			return netip.Prefix{}, false
		}
		return p.Masked(), true
	}
	addr, err := netip.ParseAddr(v)
	if err != nil {
		return netip.Prefix{}, false
	}
	a := addr.Unmap()
	return netip.PrefixFrom(a, a.BitLen()), true
}

// Contains reports whether addr falls inside the set.
func (s *PrefixSet) Contains(addr netip.Addr) bool {
	if s == nil || len(s.prefixes) == 0 {
		return false
	}
	ip := addr.Unmap()
	for _, p := range s.prefixes {
		if p.Addr().Is4() == ip.Is4() && p.Contains(ip) {
			return true
		}
	}
	return false
}

// Empty reports whether the set has no usable entries.
func (s *PrefixSet) Empty() bool { return s == nil || len(s.prefixes) == 0 }

// Len reports how many prefixes the set holds.
func (s *PrefixSet) Len() int {
	if s == nil {
		return 0
	}
	return len(s.prefixes)
}

// AddrFromNetAddr extracts the IP from a net.Addr, which is how both the DNS
// listeners and net/http hand us the client address.
func AddrFromNetAddr(a net.Addr) netip.Addr {
	switch v := a.(type) {
	case *net.UDPAddr:
		if addr, ok := netip.AddrFromSlice(v.IP); ok {
			return addr.Unmap()
		}
	case *net.TCPAddr:
		if addr, ok := netip.AddrFromSlice(v.IP); ok {
			return addr.Unmap()
		}
	}
	if a == nil {
		return netip.Addr{}
	}
	host, _, err := net.SplitHostPort(a.String())
	if err != nil {
		host = a.String()
	}
	addr, err := netip.ParseAddr(host)
	if err != nil {
		return netip.Addr{}
	}
	return addr.Unmap()
}

// AddrFromHostPort parses a "host:port" or bare-host string into an address.
func AddrFromHostPort(s string) (netip.Addr, bool) {
	s = strings.TrimSpace(s)
	if s == "" {
		return netip.Addr{}, false
	}
	if host, _, err := net.SplitHostPort(s); err == nil {
		s = host
	}
	addr, err := netip.ParseAddr(strings.Trim(s, "[]"))
	if err != nil {
		return netip.Addr{}, false
	}
	return addr.Unmap(), true
}