rony-llm-agent/pkg/tools/sandbox/network.go

package sandbox

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// NetworkPolicy controls which URLs network-facing tools (e.g. webfetch) may
// reach — Phase 2 §8.1 egress control. The zero value is a usable default:
// http/https only, all domains, private ranges allowed (Rony is local-first,
// so talking to localhost is normal), but cloud-metadata endpoints always
// blocked — no configuration can open those, since leaking instance
// credentials is never what a fetch tool is for.
type NetworkPolicy struct {
	// AllowSchemes lists permitted URL schemes; empty means http and https.
	AllowSchemes []string
	// AllowDomains, when non-empty, is an allowlist: only these hosts (or
	// their subdomains) may be fetched.
	AllowDomains []string
	// DenyDomains lists hosts (and their subdomains) that may never be
	// fetched, evaluated before AllowDomains.
	DenyDomains []string
	// BlockPrivateIPs, when true, refuses loopback, RFC1918/4193 and
	// link-local addresses — both literal IPs in the URL and, via
	// HTTPClient's dial-time check, whatever a hostname actually resolves
	// to (defeating DNS-rebinding tricks that pass a hostname check but
	// resolve to an internal address).
	BlockPrivateIPs bool
}

// metadataIPs are cloud instance-metadata endpoints (AWS/GCP/Azure IMDS and
// the AWS ECS/EKS variant). Fetching them exfiltrates instance credentials,
// so they're refused unconditionally.
var metadataIPs = []string{"169.254.169.254", "169.254.170.2", "fd00:ec2::254"}

// Validate reports whether rawURL is allowed by the policy. It checks the
// scheme, the host against deny/allow lists, and — for literal IP hosts —
// the IP itself. Hostnames that resolve to blocked IPs are caught later at
// dial time by HTTPClient; call that too for full coverage.
func (p *NetworkPolicy) Validate(rawURL string) error {
	u, err := url.Parse(rawURL)
	if err != nil {
		return fmt.Errorf("egress policy: invalid url: %w", err)
	}

	scheme := strings.ToLower(u.Scheme)
	schemes := p.AllowSchemes
	if len(schemes) == 0 {
		schemes = []string{"http", "https"}
	}
	schemeOK := false
	for _, s := range schemes {
		if scheme == strings.ToLower(s) {
			schemeOK = true
			break
		}
	}
	if !schemeOK {
		return fmt.Errorf("egress policy: scheme %q not allowed", u.Scheme)
	}

	host := strings.ToLower(u.Hostname())
	if host == "" {
		return fmt.Errorf("egress policy: url has no host")
	}

	for _, d := range p.DenyDomains {
		if hostMatches(host, d) {
			return fmt.Errorf("egress policy: host %q is denied", host)
		}
	}
	if len(p.AllowDomains) > 0 {
		allowed := false
		for _, d := range p.AllowDomains {
			if hostMatches(host, d) {
				allowed = true
				break
			}
		}
		if !allowed {
			return fmt.Errorf("egress policy: host %q is not in the allowlist", host)
		}
	}

	if ip := net.ParseIP(host); ip != nil {
		if err := p.checkIP(ip); err != nil {
			return err
		}
	}
	return nil
}

// HTTPClient returns an *http.Client that re-checks every connection's
// resolved IP at dial time, so a hostname that passed Validate can't smuggle
// a request to a blocked address (DNS rebinding, or a benign-looking name
// resolving to a metadata endpoint). Redirects are re-validated too — a
// permitted URL redirecting to a blocked one is refused.
func (p *NetworkPolicy) HTTPClient(timeout time.Duration) *http.Client {
	dialer := &net.Dialer{Timeout: 15 * time.Second}
	transport := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			host, port, err := net.SplitHostPort(addr)
			if err != nil {
				return nil, err
			}
			ips, err := net.DefaultResolver.LookupIP(ctx, "ip", host)
			if err != nil {
				return nil, err
			}
			for _, ip := range ips {
				if err := p.checkIP(ip); err != nil {
					return nil, err
				}
			}
			// Dial one of the vetted IPs directly (rather than the
			// hostname) so the connection can't re-resolve to something
			// that was never checked.
			return dialer.DialContext(ctx, network, net.JoinHostPort(ips[0].String(), port))
		},
	}
	return &http.Client{
		Timeout:   timeout,
		Transport: transport,
		CheckRedirect: func(req *http.Request, _ []*http.Request) error {
			return p.Validate(req.URL.String())
		},
	}
}

// checkIP enforces the always-on metadata block and, when BlockPrivateIPs is
// set, the private/loopback/link-local ranges.
func (p *NetworkPolicy) checkIP(ip net.IP) error {
	for _, m := range metadataIPs {
		if ip.Equal(net.ParseIP(m)) {
			return fmt.Errorf("egress policy: cloud metadata endpoint %s is always blocked", ip)
		}
	}
	if !p.BlockPrivateIPs {
		return nil
	}
	if ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() {
		return fmt.Errorf("egress policy: private/internal address %s is blocked", ip)
	}
	return nil
}

// hostMatches reports whether host equals domain or is a subdomain of it.
func hostMatches(host, domain string) bool {
	domain = strings.ToLower(strings.TrimPrefix(domain, "."))
	return host == domain || strings.HasSuffix(host, "."+domain)
}