ai-gateway-go/internal/provider/url_validator.go

package provider

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"net/url"
	"strings"
	"time"
)

// ErrBlockedAddress 标记 base_url 解析到被禁止的网段(私网/特殊用途网段)。
// 该错误携带解析出的地址细节,只应记录在服务端日志,不得原样返回给客户端。
var ErrBlockedAddress = errors.New("base_url resolves to a blocked address")

func ValidateBaseURL(ctx context.Context, raw string, allowPrivate bool) (string, error) {
	parsed, err := url.Parse(strings.TrimSpace(raw))
	if err != nil {
		return "", err
	}
	if (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Hostname() == "" {
		return "", errors.New("base_url must be an absolute http(s) URL")
	}
	if parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
		return "", errors.New("base_url must not contain user info, query, or fragment")
	}
	if !allowPrivate {
		addresses, err := net.DefaultResolver.LookupIPAddr(ctx, parsed.Hostname())
		if err != nil {
			return "", fmt.Errorf("resolve base_url host: %w", err)
		}
		if len(addresses) == 0 {
			return "", errors.New("base_url host did not resolve")
		}
		for _, address := range addresses {
			if !IsPublicAddress(address.IP) {
				return "", fmt.Errorf("%w %s", ErrBlockedAddress, address.IP)
			}
		}
	}
	parsed.Path = strings.TrimRight(parsed.Path, "/")
	return parsed.String(), nil
}

// specialPurposePrefixes 是 Go netip 内建分类(loopback/private/link-local/
// multicast/unspecified)之外、但绝不应作为出站上游的 IANA 特殊用途网段。
// 内网服务常部署在 CGNAT(100.64/10)与 benchmark(198.18/15)段,而 6to4/NAT64
// 前缀可以把 IPv6 地址桥接回内网 IPv4,因此必须一并拦截。
var specialPurposePrefixes = []netip.Prefix{
	// IPv4 特殊用途网段(RFC 6890 及其更新)。
	netip.MustParsePrefix("100.64.0.0/10"),   // CGNAT 共享地址空间 RFC 6598
	netip.MustParsePrefix("192.0.0.0/24"),    // IETF 协议保留
	netip.MustParsePrefix("192.0.2.0/24"),    // TEST-NET-1 文档
	netip.MustParsePrefix("192.88.99.0/24"),  // 6to4 中继任播(已弃用)
	netip.MustParsePrefix("198.18.0.0/15"),   // 基准测试 RFC 2544
	netip.MustParsePrefix("198.51.100.0/24"), // TEST-NET-2 文档
	netip.MustParsePrefix("203.0.113.0/24"),  // TEST-NET-3 文档
	netip.MustParsePrefix("240.0.0.0/4"),     // 保留(含广播地址)
	// IPv6 特殊用途网段。
	netip.MustParsePrefix("2001:db8::/32"),  // 文档地址
	netip.MustParsePrefix("2001:10::/28"),   // ORCHID
	netip.MustParsePrefix("2002::/16"),      // 6to4:内嵌 IPv4,可桥接回内网
	netip.MustParsePrefix("64:ff9b::/96"),   // NAT64 知名前缀
	netip.MustParsePrefix("64:ff9b:1::/48"), // NAT64 本地使用前缀
}

// IsPublicAddress 报告 ip 是否为可安全访问的公网单播地址。IPv4-mapped
// IPv6(::ffff:a.b.c.d)先解映射为 IPv4 再判断,防止绕过。SSRF 防护统一使用
// 本函数,写入校验与拨号时校验共用同一份判定。
func IsPublicAddress(ip net.IP) bool {
	if ip == nil {
		return false
	}
	addr, ok := netip.AddrFromSlice(ip)
	if !ok {
		return false
	}
	addr = addr.Unmap()
	if !addr.IsValid() || addr.IsUnspecified() || addr.IsLoopback() || addr.IsMulticast() ||
		addr.IsLinkLocalUnicast() || addr.IsLinkLocalMulticast() || addr.IsInterfaceLocalMulticast() ||
		addr.IsPrivate() {
		return false
	}
	for _, prefix := range specialPurposePrefixes {
		if prefix.Contains(addr) {
			return false
		}
	}
	return true
}

// SafeDialContext 构造拨号函数:allowPrivate 为 false 时,在拨号前对解析出的
// 全部地址做 IsPublicAddress 校验,并按校验通过的地址直连(不再二次解析,
// 缓解 DNS rebinding)。gateway 数据平面与控制面客户端共用此实现。
func SafeDialContext(allowPrivate bool, timeout, keepAlive time.Duration) func(context.Context, string, string) (net.Conn, error) {
	dialer := &net.Dialer{Timeout: timeout, KeepAlive: keepAlive}
	if allowPrivate {
		return dialer.DialContext
	}
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		host, port, err := net.SplitHostPort(address)
		if err != nil {
			return nil, err
		}
		addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
		if err != nil {
			return nil, err
		}
		if len(addresses) == 0 {
			return nil, errors.New("upstream host did not resolve")
		}
		for _, address := range addresses {
			if !IsPublicAddress(address.IP) {
				return nil, fmt.Errorf("upstream resolved to blocked address %s", address.IP)
			}
		}
		return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
	}
}