// 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"
	"syscall"
	"time"
)

var (
	cgnat      = netip.MustParsePrefix("100.64.0.0/10")     // tailnet IPv4 (CGNAT)
	tailnetULA = netip.MustParsePrefix("fd7a:115c:a1e0::/48") // tailnet IPv6
	nat64      = netip.MustParsePrefix("64:ff9b::/96")
	ipv4zero   = netip.MustParsePrefix("0.0.0.0/8")
)

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
}

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 {
		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 address predicate applied to every resolved record and to
// the literal IP in Dialer.Control.
func CheckAddr(addr netip.Addr) error {
	if !addr.IsValid() {
		return blocked("invalid address")
	}
	addr = addr.Unmap()

	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 nat64.Contains(addr) {
		var b [4]byte
		raw := addr.As16()
		copy(b[:], raw[12:])
		embedded := netip.AddrFrom4(b)
		if err := CheckAddr(embedded); err != nil {
			return blockedAddr("NAT64-embedded", embedded)
		}
		return blockedAddr("NAT64 prefix", addr)
	}

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

	if !addr.IsGlobalUnicast() {
		return blockedAddr("non-global address", addr)
	}
	if addr.Is4() && ipv4zero.Contains(addr) {
		return blockedAddr("unspecified address", addr)
	}
	if addr.Is4() && cgnat.Contains(addr) {
		return blockedAddr("tailnet address", addr)
	}
	if addr.Is6() && tailnetULA.Contains(addr) {
		return blockedAddr("tailnet 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 {
		return nil, nil, err
	}
	if literal.IsValid() {
		if err := 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 := 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 {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return blocked("dial address " + address)
	}
	if err := checkPort(port); err != nil {
		return err
	}
	addr, err := netip.ParseAddr(host)
	if err != nil {
		return blocked("dial address not an IP " + host)
	}
	return 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,
		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,
	}
}