Agent-System-Backend/internal/vision/fetch.go

package vision

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	agentllm "agent-native/runtime/llm"
)

var (
	errPrivateAddress = errors.New("地址不在允许范围")
	errNotImage       = errors.New("响应不是受支持的图片")
	errTooLarge       = errors.New("超过单图大小上限")
)

// FetchImage 下载一张公网图片。拒绝回环、私网、链路本地和云元数据地址，
// 并把连接钉在解析到的公网 IP 上，避免 DNS 重绑定。不走环境变量里的代理。
func FetchImage(ctx context.Context, rawURL string) (agentllm.Image, error) {
	if err := validateImageURL(ctx, rawURL); err != nil {
		return agentllm.Image{}, err
	}
	client := &http.Client{
		Timeout: 20 * time.Second,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= 3 {
				return errors.New("重定向过多")
			}
			if err := validateImageURL(req.Context(), req.URL.String()); err != nil {
				return err
			}
			return nil
		},
		Transport: &http.Transport{
			Proxy: nil,
			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
				host, port, err := net.SplitHostPort(addr)
				if err != nil {
					return nil, err
				}
				ip, err := pickPublicIP(ctx, host)
				if err != nil {
					return nil, err
				}
				var d net.Dialer
				d.Timeout = 10 * time.Second
				return d.DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
			},
			ResponseHeaderTimeout: 10 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
		},
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return agentllm.Image{}, err
	}
	req.Header.Set("User-Agent", "agent-system-vision/1.0")
	resp, err := client.Do(req)
	if err != nil {
		if errors.Is(err, errPrivateAddress) {
			return agentllm.Image{}, errPrivateAddress
		}
		return agentllm.Image{}, fmt.Errorf("下载失败")
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return agentllm.Image{}, fmt.Errorf("下载失败")
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, MaxImageBytes+1))
	if err != nil {
		return agentllm.Image{}, fmt.Errorf("下载失败")
	}
	if len(body) > MaxImageBytes {
		return agentllm.Image{}, errTooLarge
	}
	kind := MIMEOf(body)
	if kind == "" || !contentTypeOK(resp.Header.Get("Content-Type"), kind) {
		return agentllm.Image{}, errNotImage
	}
	// 只回传字节。URL 留给网关自己去拉会绕过这里的地址校验。
	return agentllm.Image{
		MIME: kind,
		Data: base64.StdEncoding.EncodeToString(body),
	}, nil
}

func publicReason(err error) string {
	switch {
	case err == nil:
		return ""
	case errors.Is(err, errPrivateAddress):
		return errPrivateAddress.Error()
	case errors.Is(err, errNotImage):
		return errNotImage.Error()
	case errors.Is(err, errTooLarge):
		return errTooLarge.Error()
	default:
		return "下载失败"
	}
}

func validateImageURL(ctx context.Context, raw string) error {
	u, err := url.Parse(raw)
	if err != nil || u.Host == "" {
		return errPrivateAddress
	}
	if u.User != nil {
		return errPrivateAddress
	}
	switch strings.ToLower(u.Scheme) {
	case "http", "https":
	default:
		return errPrivateAddress
	}
	if _, err := pickPublicIP(ctx, u.Hostname()); err != nil {
		return errPrivateAddress
	}
	return nil
}

func pickPublicIP(ctx context.Context, host string) (net.IP, error) {
	if ip := net.ParseIP(host); ip != nil {
		if err := checkIP(ip); err != nil {
			return nil, err
		}
		return ip, nil
	}
	if host == "" || strings.EqualFold(host, "localhost") {
		return nil, errPrivateAddress
	}
	addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host)
	if err != nil || len(addrs) == 0 {
		return nil, errPrivateAddress
	}
	var first net.IP
	for _, addr := range addrs {
		if err := checkIP(addr.IP); err != nil {
			return nil, err
		}
		if first == nil {
			first = addr.IP
		}
	}
	if first == nil {
		return nil, errPrivateAddress
	}
	return first, nil
}

var (
	cgnat = mustCIDR("100.64.0.0/10")
)

func mustCIDR(s string) *net.IPNet {
	_, n, err := net.ParseCIDR(s)
	if err != nil {
		panic(err)
	}
	return n
}

func checkIP(ip net.IP) error {
	if ip == nil {
		return errPrivateAddress
	}
	if v4 := ip.To4(); v4 != nil {
		ip = v4
	}
	switch {
	case ip.IsLoopback(), ip.IsPrivate(), ip.IsLinkLocalUnicast(), ip.IsLinkLocalMulticast(),
		ip.IsMulticast(), ip.IsUnspecified(), cgnat.Contains(ip):
		return errPrivateAddress
	default:
		return nil
	}
}

func contentTypeOK(header, sniffed string) bool {
	ct := strings.ToLower(strings.TrimSpace(strings.Split(header, ";")[0]))
	switch ct {
	case "", "application/octet-stream", "binary/octet-stream":
		return sniffed != ""
	default:
		return ct == sniffed
	}
}