// Package guard is the SSRF control plane for sift's outbound fetches.
//
// Threat model: the URL is attacker-influenced (prompt-injected pages tell
// agents to fetch things). Sift runs inside the tailnet, so without these
// checks it is an SSRF gadget into Honcho, Dixie, Mattermost, orion itself,
// and the LAN.
//
// Defence is two-layered and both layers must fire:
//
//  1. Pre-check: parse the URL, require http(s) / no userinfo / port 80 or 443,
//     resolve every A/AAAA, and reject the URL if *any* record is blocked.
//  2. Dial-time: net.Dialer.Control sees the literal IP about to be connect()ed
//     (including Happy Eyeballs races and a rebinding answer that arrived after
//     the pre-check) and applies the same address predicate.
//
// Redirects are followed manually (max 5); each hop re-runs the pre-check.
// There is no proxy, including from the environment.
package guard

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"time"
)

const maxRedirects = 5

// LookupIP resolves a hostname to addresses. Tests inject this.
type LookupIP func(ctx context.Context, host string) ([]netip.Addr, error)

// Guard holds the SSRF policy. A zero Guard is ready to use.
type Guard struct {
	LookupIP LookupIP
	// DenyPrefixes are extra CIDRs blocked at pre-check and at Control
	// (home WAN /32s that hairpin through the router into LAN services).
	// Set from config deny_prefixes / -deny-prefix; never from the URL.
	DenyPrefixes []netip.Prefix
	// AllowPrefixes bypasses the predicate. Production config never sets this;
	// only tests and the sifttest-tagged -test-allow-prefix flag do.
	AllowPrefixes []netip.Prefix

	// self is the discovered WAN deny set (v4 /32, v6 /64), swapped atomically.
	self atomic.Pointer[[]netip.Prefix]
}

func (g *Guard) allows(addr netip.Addr) bool {
	if g == nil {
		return false
	}
	addr = addr.Unmap()
	for _, p := range g.AllowPrefixes {
		if p.Contains(addr) {
			return true
		}
	}
	return false
}

func (g *Guard) lookup(ctx context.Context, host string) ([]netip.Addr, error) {
	if g.LookupIP != nil {
		return g.LookupIP(ctx, host)
	}
	return net.DefaultResolver.LookupNetIP(ctx, "ip", host)
}

// BlockedError is returned for every policy rejection. Error() is stable
// ("blocked: …") so the HTTP layer can put it in Firecrawl-shaped responses
// without agents retrying.
type BlockedError struct {
	Reason string
	Addr   netip.Addr
	Port   string
}

func (e *BlockedError) Error() string {
	if e.Addr.IsValid() {
		return fmt.Sprintf("blocked: %s %s", e.Reason, e.Addr)
	}
	if e.Port != "" {
		return fmt.Sprintf("blocked: %s %s", e.Reason, e.Port)
	}
	return "blocked: " + e.Reason
}

func blocked(reason string) error {
	return &BlockedError{Reason: reason}
}

func blockedAddr(reason string, addr netip.Addr) error {
	return &BlockedError{Reason: reason, Addr: addr}
}

func blockedPort(port string) error {
	return &BlockedError{Reason: "port", Port: port}
}

// ParseURL validates scheme, userinfo, host, and port. It does not resolve.
// Hostnames are returned as ASCII (non-ASCII IDN is rejected in P0: fail closed,
// punycode required). IP literals include inet_aton-style decimal/octal/hex
// forms so those never reach DNS.
func ParseURL(raw string) (*url.URL, netip.Addr, error) {
	if raw == "" {
		return nil, netip.Addr{}, blocked("empty url")
	}
	u, err := url.Parse(raw)
	if err != nil {
		return nil, netip.Addr{}, blocked("invalid url")
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return nil, netip.Addr{}, blocked("scheme " + u.Scheme)
	}
	if u.User != nil {
		return nil, netip.Addr{}, blocked("userinfo")
	}
	host := u.Hostname()
	if host == "" {
		return nil, netip.Addr{}, blocked("missing host")
	}
	port := u.Port()
	if port == "" {
		if u.Scheme == "https" {
			port = "443"
		} else {
			port = "80"
		}
	}
	if err := checkPort(port); err != nil {
		if addr, ok := parseIPLiteral(host); ok {
			return u, addr, err
		}
		return nil, netip.Addr{}, err
	}
	if addr, ok := parseIPLiteral(host); ok {
		return u, addr, nil
	}
	if err := checkASCIIHost(host); err != nil {
		return nil, netip.Addr{}, err
	}
	return u, netip.Addr{}, nil
}

func checkPort(port string) error {
	n, err := strconv.Atoi(port)
	if err != nil || n < 1 || n > 65535 {
		return blockedPort(port)
	}
	if n != 80 && n != 443 {
		return blockedPort(port)
	}
	return nil
}

func checkASCIIHost(host string) error {
	for _, r := range host {
		if r > 127 {
			return blocked("non-ascii host (punycode required)")
		}
		if r < 0x20 || strings.ContainsRune(" /\\", r) {
			return blocked("invalid host")
		}
	}
	if strings.Contains(host, "..") {
		return blocked("invalid host")
	}
	return nil
}

// parseIPLiteral accepts canonical IPs plus inet_aton-style IPv4 (decimal,
// octal, hex, 1–4 dotted parts) and IPv4-mapped / zone-id IPv6.
func parseIPLiteral(host string) (netip.Addr, bool) {
	if addr, err := netip.ParseAddr(host); err == nil {
		return addr, true
	}
	if addr, ok := parseIPv4AtOn(host); ok {
		return addr, true
	}
	return netip.Addr{}, false
}

func parseIPv4AtOn(s string) (netip.Addr, bool) {
	parts := strings.Split(s, ".")
	if len(parts) < 1 || len(parts) > 4 {
		return netip.Addr{}, false
	}
	nums := make([]uint64, len(parts))
	for i, p := range parts {
		n, ok := parseIPv4Part(p)
		if !ok {
			return netip.Addr{}, false
		}
		nums[i] = n
	}
	var b [4]byte
	switch len(nums) {
	case 1:
		if nums[0] > 0xffffffff {
			return netip.Addr{}, false
		}
		b[0] = byte(nums[0] >> 24)
		b[1] = byte(nums[0] >> 16)
		b[2] = byte(nums[0] >> 8)
		b[3] = byte(nums[0])
	case 2:
		if nums[0] > 0xff || nums[1] > 0xffffff {
			return netip.Addr{}, false
		}
		b[0] = byte(nums[0])
		b[1] = byte(nums[1] >> 16)
		b[2] = byte(nums[1] >> 8)
		b[3] = byte(nums[1])
	case 3:
		if nums[0] > 0xff || nums[1] > 0xff || nums[2] > 0xffff {
			return netip.Addr{}, false
		}
		b[0] = byte(nums[0])
		b[1] = byte(nums[1])
		b[2] = byte(nums[2] >> 8)
		b[3] = byte(nums[2])
	case 4:
		for i := 0; i < 4; i++ {
			if nums[i] > 0xff {
				return netip.Addr{}, false
			}
			b[i] = byte(nums[i])
		}
	}
	return netip.AddrFrom4(b), true
}

func parseIPv4Part(s string) (uint64, bool) {
	if s == "" {
		return 0, false
	}
	switch {
	case len(s) > 2 && (s[0] == '0' && (s[1] == 'x' || s[1] == 'X')):
		n, err := strconv.ParseUint(s[2:], 16, 32)
		return n, err == nil
	case len(s) > 1 && s[0] == '0':
		n, err := strconv.ParseUint(s, 8, 32)
		return n, err == nil
	default:
		n, err := strconv.ParseUint(s, 10, 32)
		return n, err == nil
	}
}

// CheckAddr is the default (no extra prefixes) address predicate.
func CheckAddr(addr netip.Addr) error {
	var g Guard
	return g.CheckAddr(addr)
}

// CheckAddr is the address predicate applied to every resolved record and to
// the literal IP in Dialer.Control. Extra DenyPrefixes are checked first.
func (g *Guard) CheckAddr(addr netip.Addr) error {
	if !addr.IsValid() {
		return blocked("invalid address")
	}
	addr = addr.Unmap()

	if g != nil {
		for _, p := range g.AllowPrefixes {
			if p.Contains(addr) {
				return nil
			}
		}
		for _, p := range g.DenyPrefixes {
			if p.Contains(addr) {
				return blockedAddr("denied prefix", addr)
			}
		}
		if extras := g.self.Load(); extras != nil {
			for _, p := range *extras {
				if p.Contains(addr) {
					return blockedAddr("denied prefix", addr)
				}
			}
		}
	}

	switch {
	case addr.IsUnspecified():
		return blockedAddr("unspecified address", addr)
	case addr.IsLoopback():
		return blockedAddr("loopback address", addr)
	case addr.IsPrivate():
		return blockedAddr("private address", addr)
	case addr.IsMulticast(), addr.IsInterfaceLocalMulticast(), addr.IsLinkLocalMulticast():
		return blockedAddr("multicast address", addr)
	case addr.IsLinkLocalUnicast():
		return blockedAddr("link-local address", addr)
	}

	if err := checkIANA(addr); err != nil {
		return err
	}

	if v4, ok := ipv4Compatible(addr); ok {
		return g.CheckAddr(v4)
	}

	if !addr.IsGlobalUnicast() {
		return blockedAddr("non-global address", addr)
	}
	return nil
}

// ipv4Compatible reports deprecated IPv4-compatible IPv6 (::/96). Unspecified
// and loopback are handled before this is called. IPv4-mapped addresses are
// Unmap'd first.
func ipv4Compatible(addr netip.Addr) (netip.Addr, bool) {
	if !addr.Is6() || addr.Is4In6() {
		return netip.Addr{}, false
	}
	a := addr.As16()
	for i := 0; i < 12; i++ {
		if a[i] != 0 {
			return netip.Addr{}, false
		}
	}
	v4 := netip.AddrFrom4([4]byte{a[12], a[13], a[14], a[15]})
	if !v4.IsValid() || v4.IsUnspecified() {
		return netip.Addr{}, false
	}
	return v4, true
}

// CheckURL parses, and if the host is a name, resolves every record. Any
// blocked address fails the whole URL (mixed-record trick).
func (g *Guard) CheckURL(ctx context.Context, raw string) (*url.URL, []netip.Addr, error) {
	u, literal, err := ParseURL(raw)
	if err != nil {
		if literal.IsValid() && g.allows(literal) {
			err = nil
		} else {
			return nil, nil, err
		}
	}
	if literal.IsValid() {
		if err := g.CheckAddr(literal); err != nil {
			return nil, nil, err
		}
		return u, []netip.Addr{literal}, nil
	}
	addrs, err := g.lookup(ctx, u.Hostname())
	if err != nil {
		return nil, nil, fmt.Errorf("blocked: resolve %s: %w", u.Hostname(), err)
	}
	if len(addrs) == 0 {
		return nil, nil, blocked("no addresses for " + u.Hostname())
	}
	for _, a := range addrs {
		if err := g.CheckAddr(a); err != nil {
			return nil, nil, err
		}
	}
	return u, addrs, nil
}

// Control is a net.Dialer Control / ControlContext hook. address is the
// remote IP:port about to be connect()ed.
func (g *Guard) Control(network, address string, _ syscall.RawConn) error {
	switch network {
	case "tcp", "tcp4", "tcp6":
	default:
		return blocked("network " + network)
	}
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return blocked("dial address " + address)
	}
	addr, err := netip.ParseAddr(host)
	if err != nil {
		return blocked("dial address not an IP " + host)
	}
	if g.allows(addr) {
		return nil
	}
	if err := checkPort(port); err != nil {
		return err
	}
	return g.CheckAddr(addr)
}

// ControlContext is the context-aware form; policy does not use ctx today.
func (g *Guard) ControlContext(_ context.Context, network, address string, c syscall.RawConn) error {
	return g.Control(network, address, c)
}

// CheckRedirect is an http.Client CheckRedirect hook. Max 5 hops; every hop
// is re-validated (parse + resolve + address predicate).
func (g *Guard) CheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= maxRedirects {
		return blocked("too many redirects")
	}
	if req.URL == nil {
		return blocked("redirect missing url")
	}
	_, _, err := g.CheckURL(req.Context(), req.URL.String())
	return err
}

// Dialer returns a net.Dialer whose Control hook is the address predicate.
// Connect timeout is 5s as specified in the plan.
func (g *Guard) Dialer() *net.Dialer {
	return &net.Dialer{
		Timeout:        5 * time.Second,
		Control:        g.Control,
		ControlContext: g.ControlContext,
	}
}

// Transport is an HTTP transport with no proxy, the guarded dialer, and a
// 20s overall request timeout left to the caller via context. Body size caps
// land in P1.
func (g *Guard) Transport() *http.Transport {
	d := g.Dialer()
	return &http.Transport{
		Proxy:                 nil, // never ProxyFromEnvironment
		DialContext:           d.DialContext,
		ForceAttemptHTTP2:     true,
		MaxIdleConns:          8,
		IdleConnTimeout:       30 * time.Second,
		TLSHandshakeTimeout:   5 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		// Go gunzips transparently when this is false. P1 size caps MUST
		// apply to the decompressed resp.Body (io.LimitReader on the body
		// the client returns); that also covers zip-bombs.
		DisableCompression: false,
	}
}

// Client returns an http.Client that refuses blocked URLs on redirect.
// Callers must still CheckURL the initial URL; Client.CheckRedirect covers hops.
func (g *Guard) Client() *http.Client {
	return &http.Client{
		Transport:     g.Transport(),
		CheckRedirect: g.CheckRedirect,
		Timeout:       20 * time.Second,
	}
}

// SetSelfDeny atomically replaces the discovered WAN deny prefixes.
// A nil or empty list clears the extra set; static DenyPrefixes are unchanged.
func (g *Guard) SetSelfDeny(prefixes []netip.Prefix) {
	if g == nil {
		return
	}
	if len(prefixes) == 0 {
		g.self.Store(nil)
		return
	}
	cp := append([]netip.Prefix(nil), prefixes...)
	g.self.Store(&cp)
}

// SelfDeny returns a copy of the current discovered deny prefixes.
func (g *Guard) SelfDeny() []netip.Prefix {
	if g == nil {
		return nil
	}
	p := g.self.Load()
	if p == nil {
		return nil
	}
	return append([]netip.Prefix(nil), *p...)
}