fix(images): 异步生图转存的图片下载与同步回填策略对齐

异步生图任务把上游返回的图片 url 下载后转存对象存储，此前沿用默认 http.Client，与同步 Images 端点 url 回填（c227863d）的下载策略不一致。 - 下载目的地为 localhost、回环、私网、链路本地、CGNAT 等非公网地址时不下载，重定向的每一跳同样校验；直连时在 socket 层按实际连接 IP 判断 - 继续支持环境变量代理：经代理时先校验目标主机的解析结果 - 响应声明类型与字节嗅探都不是图片时任务失败，不再按 png 落盘 - 下载客户端进程内共享，复用连接池；超时与大小上限不变
2026-10-19 16:17:56 +00:00 · 2026-10-19 17:51:18 +08:00
parent b0a02a3304
commit ab71e12ec4
2 changed files with 386 additions and 11 deletions
@@ -8,10 +8,16 @@ import (
 	"fmt"
 	"io"
 	"mime"
+	"net"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
 )

 const defaultImageMaxDownloadBytes int64 = 32 << 20 // 32 MiB
@@ -53,8 +59,123 @@ func NewImageResultUploader(storage ImageStorage, prefix string, maxDownloadByte
 	}
 }

-func defaultImageDownloadHTTPClient() *http.Client {
-	return &http.Client{Timeout: 60 * time.Second}
+// defaultImageDownloadHTTPClient 是进程内共享的图片下载客户端，复用连接池。
+var defaultImageDownloadHTTPClient = sync.OnceValue(func() *http.Client {
+	return newImageDownloadHTTPClient(http.ProxyFromEnvironment)
+})
+
+// newImageDownloadHTTPClient 返回只允许下载公网目的地的客户端。
+// 图片 url 来自上游响应，是不可信输入：目的地为 localhost、回环、私网、链路本地、
+// 云元数据等地址时一律拒绝，重定向的每一跳都经同一 RoundTrip 校验。
+// 不经代理时在 socket 层校验真实连接 IP（防 DNS rebinding）；经环境代理时由代理解析目标，
+// 只能在发出前校验本地解析结果。
+func newImageDownloadHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client {
+	direct := newImageDownloadBaseTransport(imageDownloadDialControl)
+	proxied := newImageDownloadBaseTransport(nil)
+	proxied.Proxy = proxy
+	return &http.Client{
+		Timeout:   60 * time.Second,
+		Transport: &imageDownloadTransport{direct: direct, proxied: proxied, proxy: proxy, resolvesPrivate: isPrivateOrLoopbackHost},
+	}
+}
+
+// newImageDownloadBaseTransport 返回与 http.DefaultTransport 参数一致、但不带代理的 Transport。
+// control 非空时在每个 socket 连接前校验目标地址。
+func newImageDownloadBaseTransport(control func(network, address string, c syscall.RawConn) error) *http.Transport {
+	dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second, Control: control}
+	return &http.Transport{
+		DialContext:           dialer.DialContext,
+		ForceAttemptHTTP2:     true,
+		MaxIdleConns:          100,
+		IdleConnTimeout:       90 * time.Second,
+		TLSHandshakeTimeout:   10 * time.Second,
+		ExpectContinueTimeout: 1 * time.Second,
+	}
+}
+
+// imageDownloadDialControl 拒绝连接到非公网 IP。address 是 DNS 解析后的真实连接地址，
+// 多个解析结果时 Dialer 会跳过被拒地址继续尝试下一个。
+func imageDownloadDialControl(_, address string, _ syscall.RawConn) error {
+	host, _, err := net.SplitHostPort(address)
+	if err != nil {
+		return err
+	}
+	if ip := net.ParseIP(stripIPv6Zone(host)); ip == nil || isPrivateIP(ip) {
+		return &net.AddrError{Err: "address not allowed for image download", Addr: address}
+	}
+	return nil
+}
+
+// imageDownloadTransport 按请求是否走环境代理分派到对应 Transport，并在发出前拒绝非公网目的地。
+type imageDownloadTransport struct {
+	direct  http.RoundTripper
+	proxied http.RoundTripper
+	proxy   func(*http.Request) (*url.URL, error)
+	// resolvesPrivate 报告 host 的本地解析结果是否含非公网地址，仅用于经代理的请求。
+	resolvesPrivate func(ctx context.Context, host string) (bool, error)
+}
+
+func (t *imageDownloadTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	next, err := t.route(req)
+	if err != nil {
+		if req.Body != nil {
+			_ = req.Body.Close()
+		}
+		return nil, err
+	}
+	return next.RoundTrip(req)
+}
+
+func (t *imageDownloadTransport) route(req *http.Request) (http.RoundTripper, error) {
+	if req.URL == nil {
+		return nil, errors.New("image url is empty")
+	}
+	host := req.URL.Hostname()
+	if isBlockedImageDownloadHost(host) {
+		return nil, fmt.Errorf("image url host is not allowed: %s", host)
+	}
+	var proxyURL *url.URL
+	if t.proxy != nil {
+		var err error
+		if proxyURL, err = t.proxy(req); err != nil {
+			return nil, err
+		}
+	}
+	if proxyURL == nil {
+		return t.direct, nil
+	}
+	if t.resolvesPrivate == nil {
+		return nil, errors.New("image download proxy host check is not configured")
+	}
+	blocked, err := t.resolvesPrivate(req.Context(), host)
+	if err != nil {
+		return nil, fmt.Errorf("resolve image url host: %w", err)
+	}
+	if blocked {
+		return nil, fmt.Errorf("image url host is not allowed: %s", host)
+	}
+	return t.proxied, nil
+}
+
+// isBlockedImageDownloadHost 报告 host 是否为空、被拒绝的主机名，或非公网 IP 字面量。
+// 域名的解析结果由 socket 层或代理分支再校验。
+func isBlockedImageDownloadHost(host string) bool {
+	host = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(host)), ".")
+	if urlvalidator.IsBlockedHost(host) || isBlockedHostname(host) {
+		return true
+	}
+	if ip := net.ParseIP(stripIPv6Zone(host)); ip != nil {
+		return isPrivateIP(ip)
+	}
+	return false
+}
+
+// stripIPv6Zone 去掉 IPv6 字面量的 zone（如 fe80::1%eth0），使其可被 net.ParseIP 解析。
+func stripIPv6Zone(host string) string {
+	if i := strings.LastIndexByte(host, '%'); i >= 0 {
+		return host[:i]
+	}
+	return host
 }

 // Rewrite 将 result（上游生图响应 JSON）里的每张图片转存到对象存储，
@@ -215,9 +336,12 @@ func (u *ImageResultUploader) download(ctx context.Context, rawURL string) ([]by
 	if int64(len(data)) > limit {
 		return nil, "", fmt.Errorf("downloaded image exceeds %d bytes", limit)
 	}
-	contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
+	contentType := strings.ToLower(strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0]))
 	if !strings.HasPrefix(contentType, "image/") {
-		contentType = detectImageContentType(data)
+		// 声明与嗅探都不是图片时拒绝，避免把任意响应体（内部接口的 JSON/HTML 等）转存给租户。
+		if contentType = detectedImageContentType(data); contentType == "" {
+			return nil, "", errors.New("download image: content is not an image")
+		}
 	}
 	return data, contentType, nil
 }
@@ -1,12 +1,17 @@
 package service

 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
+	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
+	"sync/atomic"
 	"testing"
 	"time"

@@ -73,18 +78,24 @@ func TestImageResultUploaderRewritesB64JSON(t *testing.T) {
 }

 func TestImageResultUploaderRewritesURL(t *testing.T) {
-	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
-		w.Header().Set("Content-Type", "image/png")
-		_, _ = w.Write(pngBytes)
-	}))
-	defer upstream.Close()
+	var requested string
+	client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+		requested = req.URL.String()
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     http.Header{"Content-Type": []string{"image/png"}},
+			Body:       io.NopCloser(bytes.NewReader(pngBytes)),
+			Request:    req,
+		}, nil
+	})}

 	storage := &fakeImageStorage{}
-	uploader := NewImageResultUploader(storage, "images/", 0, nil)
+	uploader := NewImageResultUploader(storage, "images/", 0, client)

-	result := json.RawMessage(`{"created":1,"data":[{"url":"` + upstream.URL + `/pic.png"}]}`)
+	result := json.RawMessage(`{"created":1,"data":[{"url":"https://images.example.test/pic.png"}]}`)
 	out, err := uploader.Rewrite(context.Background(), "imgtask_xyz", result)
 	require.NoError(t, err)
+	require.Equal(t, "https://images.example.test/pic.png", requested)

 	require.Len(t, storage.saved, 1)
 	require.Equal(t, pngBytes, storage.saved[0].data)
@@ -244,3 +255,243 @@ func TestImageTaskServiceCompleteOffloadFailureMarksFailed(t *testing.T) {
 	require.Contains(t, string(got.Error), "object storage")
 	require.NotContains(t, string(got.Result), "b64_json", "failed offload must not persist base64 to Redis")
 }
+
+func TestImageResultUploaderDefaultClientRejectsPrivateURL(t *testing.T) {
+	var hits atomic.Int32
+	internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		hits.Add(1)
+		_, _ = w.Write([]byte("GIF89a-payload"))
+	}))
+	defer internal.Close()
+
+	storage := &fakeImageStorage{}
+	uploader := NewImageResultUploader(storage, "images/", 0, nil)
+	result := json.RawMessage(`{"data":[{"url":"` + internal.URL + `/a.gif"}]}`)
+
+	_, err := uploader.Rewrite(context.Background(), "imgtask_private", result)
+	require.ErrorContains(t, err, "image url host is not allowed")
+	require.Zero(t, hits.Load(), "internal service must not be contacted")
+	require.Empty(t, storage.saved)
+}
+
+func TestImageTaskServiceCompletePrivateURLMarksFailed(t *testing.T) {
+	var hits atomic.Int32
+	internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		hits.Add(1)
+		_, _ = w.Write(pngBytes)
+	}))
+	defer internal.Close()
+
+	store := &imageTaskMemoryStore{}
+	storage := &fakeImageStorage{}
+	svc := NewImageTaskServiceWithUploader(store, NewImageResultUploader(storage, "images/", 0, nil), time.Hour, time.Minute)
+	owner := ImageTaskOwner{UserID: 1, APIKeyID: 2}
+	created, err := svc.Create(context.Background(), owner)
+	require.NoError(t, err)
+
+	result := json.RawMessage(`{"data":[{"url":"` + internal.URL + `/a.gif"}]}`)
+	require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result))
+
+	got, err := svc.Get(context.Background(), owner, created.ID)
+	require.NoError(t, err)
+	require.Equal(t, ImageTaskStatusFailed, got.Status)
+	require.Empty(t, got.ImageURL)
+	require.NotContains(t, string(got.Result), internal.URL)
+	require.Zero(t, hits.Load())
+	require.Empty(t, storage.saved)
+}
+
+func newCountingImageDownloadTransport(proxyURL *url.URL) (*imageDownloadTransport, *atomic.Int32, *atomic.Int32) {
+	var directCalls, proxiedCalls atomic.Int32
+	respond := func(counter *atomic.Int32) http.RoundTripper {
+		return roundTripFunc(func(req *http.Request) (*http.Response, error) {
+			counter.Add(1)
+			return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader(pngBytes)), Request: req}, nil
+		})
+	}
+	return &imageDownloadTransport{
+		direct:  respond(&directCalls),
+		proxied: respond(&proxiedCalls),
+		proxy:   func(*http.Request) (*url.URL, error) { return proxyURL, nil },
+		resolvesPrivate: func(ctx context.Context, host string) (bool, error) {
+			if host == "intranet.example.test" {
+				return true, nil
+			}
+			if host == "unresolvable.example.test" {
+				return false, errors.New("no such host")
+			}
+			return isPrivateOrLoopbackHost(ctx, host)
+		},
+	}, &directCalls, &proxiedCalls
+}
+
+func TestImageDownloadTransportRejectsNonPublicHosts(t *testing.T) {
+	hosts := []string{
+		"http://127.0.0.1:18082/a.gif",
+		"http://localhost/a.png",
+		"http://LOCALHOST./a.png",
+		"http://foo.localhost/a.png",
+		"http://10.0.0.8/a.png",
+		"http://172.16.3.4/a.png",
+		"http://192.168.1.1/a.png",
+		"http://100.64.0.1/a.png",
+		"http://169.254.169.254/latest/meta-data/",
+		"http://metadata.google.internal/computeMetadata/v1/",
+		"http://0.0.0.0/a.png",
+		"http://[::1]/a.png",
+		"http://[::ffff:127.0.0.1]/a.png",
+		"http://[fd00::1]/a.png",
+		"http://[fe80::1%25eth0]/a.png",
+		"http:///a.png",
+	}
+	for _, proxied := range []bool{false, true} {
+		for _, rawURL := range hosts {
+			var proxyURL *url.URL
+			if proxied {
+				proxyURL = &url.URL{Scheme: "http", Host: "127.0.0.1:7890"}
+			}
+			transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(proxyURL)
+			req, err := http.NewRequest(http.MethodGet, rawURL, nil)
+			require.NoError(t, err)
+
+			_, err = transport.RoundTrip(req)
+			require.Error(t, err, "proxied=%v url=%s", proxied, rawURL)
+			require.Zero(t, directCalls.Load(), rawURL)
+			require.Zero(t, proxiedCalls.Load(), rawURL)
+		}
+	}
+}
+
+func TestImageDownloadTransportRoutesPublicHosts(t *testing.T) {
+	transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(nil)
+	req, err := http.NewRequest(http.MethodGet, "https://203.0.113.10/a.png", nil)
+	require.NoError(t, err)
+	resp, err := transport.RoundTrip(req)
+	require.NoError(t, err)
+	_ = resp.Body.Close()
+	require.EqualValues(t, 1, directCalls.Load())
+	require.Zero(t, proxiedCalls.Load())
+
+	transport, directCalls, proxiedCalls = newCountingImageDownloadTransport(&url.URL{Scheme: "http", Host: "127.0.0.1:7890"})
+	req, err = http.NewRequest(http.MethodGet, "https://203.0.113.10/a.png", nil)
+	require.NoError(t, err)
+	resp, err = transport.RoundTrip(req)
+	require.NoError(t, err)
+	_ = resp.Body.Close()
+	require.Zero(t, directCalls.Load())
+	require.EqualValues(t, 1, proxiedCalls.Load(), "a loopback env proxy must still be usable for public targets")
+}
+
+func TestImageDownloadClientRejectsRedirectToPrivateHost(t *testing.T) {
+	var calls atomic.Int32
+	transport := &imageDownloadTransport{
+		direct: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+			calls.Add(1)
+			return &http.Response{
+				StatusCode: http.StatusFound,
+				Header:     http.Header{"Location": []string{"http://169.254.169.254/latest/meta-data/"}},
+				Body:       io.NopCloser(bytes.NewReader(nil)),
+				Request:    req,
+			}, nil
+		}),
+	}
+	uploader := NewImageResultUploader(&fakeImageStorage{}, "images/", 0, &http.Client{Transport: transport})
+
+	_, err := uploader.Rewrite(context.Background(), "imgtask_redirect", json.RawMessage(`{"data":[{"url":"https://images.example.test/a.png"}]}`))
+	require.ErrorContains(t, err, "image url host is not allowed: 169.254.169.254")
+	require.EqualValues(t, 1, calls.Load(), "the private redirect hop must not be sent")
+}
+
+func TestImageDownloadClientDialsOnlyPublicAddresses(t *testing.T) {
+	var hits atomic.Int32
+	internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		hits.Add(1)
+		_, _ = w.Write(pngBytes)
+	}))
+	defer internal.Close()
+
+	client := newImageDownloadHTTPClient(nil)
+	direct, ok := client.Transport.(*imageDownloadTransport).direct.(*http.Transport)
+	require.True(t, ok)
+	require.Nil(t, direct.Proxy, "the direct transport must not pick up an environment proxy")
+
+	// 绕过 RoundTrip 预检直接 dial：IP 字面量与经 DNS 解析到回环的域名都必须在 socket 层被拒。
+	conn, err := direct.DialContext(context.Background(), "tcp", internal.Listener.Addr().String())
+	if conn != nil {
+		_ = conn.Close()
+	}
+	require.ErrorContains(t, err, "address not allowed for image download")
+	_, port, err := net.SplitHostPort(internal.Listener.Addr().String())
+	require.NoError(t, err)
+	conn, err = direct.DialContext(context.Background(), "tcp", net.JoinHostPort("localhost", port))
+	if conn != nil {
+		_ = conn.Close()
+	}
+	require.ErrorContains(t, err, "address not allowed for image download")
+	require.Zero(t, hits.Load())
+}
+
+func newStaticImageDownloadClient(contentType string, body []byte) *http.Client {
+	return &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+		header := http.Header{}
+		if contentType != "" {
+			header.Set("Content-Type", contentType)
+		}
+		return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(bytes.NewReader(body)), Request: req}, nil
+	})}
+}
+
+func TestImageResultUploaderDownloadContentTypes(t *testing.T) {
+	avif := []byte("\x00\x00\x00\x1cftypavif-not-sniffable")
+	tests := []struct {
+		name        string
+		contentType string
+		body        []byte
+		wantType    string
+		wantErr     string
+	}{
+		{name: "declared png", contentType: "image/png", body: pngBytes, wantType: "image/png"},
+		{name: "octet-stream sniffed png", contentType: "application/octet-stream", body: pngBytes, wantType: "image/png"},
+		{name: "missing header sniffed png", body: pngBytes, wantType: "image/png"},
+		{name: "declared avif not sniffable", contentType: "Image/AVIF; charset=binary", body: avif, wantType: "image/avif"},
+		{name: "json body", contentType: "application/json", body: []byte(`{"ok":true}`), wantErr: "content is not an image"},
+		{name: "html body", contentType: "text/html", body: []byte("<html>admin</html>"), wantErr: "content is not an image"},
+		{name: "empty body", body: nil, wantErr: "content is not an image"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			storage := &fakeImageStorage{}
+			uploader := NewImageResultUploader(storage, "images/", 0, newStaticImageDownloadClient(tt.contentType, tt.body))
+			_, err := uploader.Rewrite(context.Background(), "imgtask_ct", json.RawMessage(`{"data":[{"url":"https://images.example.test/a"}]}`))
+			if tt.wantErr != "" {
+				require.ErrorContains(t, err, tt.wantErr)
+				require.Empty(t, storage.saved)
+				return
+			}
+			require.NoError(t, err)
+			require.Len(t, storage.saved, 1)
+			require.Equal(t, tt.wantType, storage.saved[0].contentType)
+		})
+	}
+}
+
+func TestImageDownloadDialControl(t *testing.T) {
+	for _, address := range []string{"127.0.0.1:443", "10.1.2.3:80", "100.100.100.200:80", "169.254.169.254:80", "[::1]:443", "[fe80::1%eth0]:443", "[::ffff:192.168.0.1]:443", "bad"} {
+		require.Error(t, imageDownloadDialControl("tcp", address, nil), address)
+	}
+	for _, address := range []string{"203.0.113.10:443", "[2001:db8::1]:443"} {
+		require.NoError(t, imageDownloadDialControl("tcp", address, nil), address)
+	}
+}
+
+func TestImageDownloadTransportProxiedChecksResolvedHost(t *testing.T) {
+	for _, host := range []string{"intranet.example.test", "unresolvable.example.test"} {
+		transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(&url.URL{Scheme: "http", Host: "127.0.0.1:7890"})
+		req, err := http.NewRequest(http.MethodGet, "https://"+host+"/a.png", nil)
+		require.NoError(t, err)
+		_, err = transport.RoundTrip(req)
+		require.Error(t, err, host)
+		require.Zero(t, directCalls.Load(), host)
+		require.Zero(t, proxiedCalls.Load(), "a proxied request must fail closed when the target is private or unresolvable: %s", host)
+	}
+}