优化缩略图

internal/service/image_proxy_stream_test.go

@@ -0,0 +1,188 @@
+package service
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"path/filepath"
+	"sync/atomic"
+	"testing"
+
+	"go.uber.org/zap"
+
+	"github.com/truewhile/MeBox/internal/config"
+)
+
+// newStreamImageProxy 构造一个允许访问指定测试上游的图片代理。
+func newStreamImageProxy(t *testing.T, upstreamURL string) *ImageProxy {
+	t.Helper()
+	parsed, err := url.Parse(upstreamURL)
+	if err != nil {
+		t.Fatalf("parse upstream url: %v", err)
+	}
+	proxy := NewImageProxy(&config.Config{Cache: config.CacheConfig{CacheDir: filepath.Join(t.TempDir(), "cache")}}, zap.NewNop())
+	proxy.SetAllowedRemoteHostsProvider(func() []string { return []string{parsed.Host} })
+	return proxy
+}
+
+// bigTestJPEG 在合法 JPEG 后面补一串数据，用来验证大图是流式落盘而不是
+// 整张读进内存（剧照/原图常有数兆字节）。
+func bigTestJPEG(t *testing.T, extraBytes int) []byte {
+	t.Helper()
+	out := make([]byte, 0, len(testJPEG)+extraBytes)
+	out = append(out, testJPEG...)
+	out = append(out, bytes.Repeat([]byte{0x5a}, extraBytes)...)
+	return out
+}
+
+// TestImageProxyStreamsRemoteImageIntoCacheFile 远程原图必须直接流式写到
+// 缓存文件，后续请求由缓存文件服务（上游只被请求一次）。
+func TestImageProxyStreamsRemoteImageIntoCacheFile(t *testing.T) {
+	payload := bigTestJPEG(t, 3<<20)
+
+	var calls int32
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&calls, 1)
+		w.Header().Set("Content-Type", "image/jpeg")
+		_, _ = w.Write(payload)
+	}))
+	defer upstream.Close()
+
+	proxy := newStreamImageProxy(t, upstream.URL)
+	raw := upstream.URL + "/poster.jpg"
+
+	rec := httptest.NewRecorder()
+	if err := proxy.Serve(t.Context(), rec, httptest.NewRequest(http.MethodGet, "/api/img", nil), raw); err != nil {
+		t.Fatalf("Serve failed: %v", err)
+	}
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200", rec.Code)
+	}
+	if !bytes.Equal(rec.Body.Bytes(), payload) {
+		t.Fatalf("body length = %d, want %d", rec.Body.Len(), len(payload))
+	}
+
+	_, cachePath, failPath := proxy.remoteImageCachePathsForValidated(raw)
+	stat, err := os.Stat(cachePath)
+	if err != nil {
+		t.Fatalf("expected the image to be streamed to %s: %v", cachePath, err)
+	}
+	if stat.Size() != int64(len(payload)) {
+		t.Fatalf("cached size = %d, want %d", stat.Size(), len(payload))
+	}
+	if _, err := os.Stat(failPath); err == nil {
+		t.Fatalf("unexpected failure marker at %s", failPath)
+	}
+	if entries, err := os.ReadDir(filepath.Dir(cachePath)); err == nil {
+		for _, entry := range entries {
+			if filepath.Ext(entry.Name()) == ".tmp" {
+				t.Fatalf("temporary file left behind: %s", entry.Name())
+			}
+		}
+	}
+
+	// 第二次请求命中缓存文件，不再回源。
+	second := httptest.NewRecorder()
+	if err := proxy.Serve(t.Context(), second, httptest.NewRequest(http.MethodGet, "/api/img", nil), raw); err != nil {
+		t.Fatalf("second Serve failed: %v", err)
+	}
+	if !bytes.Equal(second.Body.Bytes(), payload) {
+		t.Fatal("cached response body mismatch")
+	}
+	if got := atomic.LoadInt32(&calls); got != 1 {
+		t.Fatalf("upstream calls = %d, want 1 (second request must be served from cache)", got)
+	}
+}
+
+// TestImageProxyRejectsNonImageUpstreamWithoutCaching 流式写入必须在提交前
+// 校验内容，否则错误页会被永久缓存成图片。
+func TestImageProxyRejectsNonImageUpstreamWithoutCaching(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/html")
+		_, _ = w.Write([]byte("<html><body>not an image</body></html>"))
+	}))
+	defer upstream.Close()
+
+	proxy := newStreamImageProxy(t, upstream.URL)
+	raw := upstream.URL + "/poster.jpg"
+
+	rec := httptest.NewRecorder()
+	if err := proxy.Serve(t.Context(), rec, httptest.NewRequest(http.MethodGet, "/api/img", nil), raw); err != nil {
+		t.Fatalf("Serve failed: %v", err)
+	}
+	if rec.Body.Len() != len(transparent1x1PNG) {
+		t.Fatalf("body length = %d, want placeholder %d", rec.Body.Len(), len(transparent1x1PNG))
+	}
+
+	_, cachePath, failPath := proxy.remoteImageCachePathsForValidated(raw)
+	if _, err := os.Stat(cachePath); err == nil {
+		t.Fatal("non-image response must not be cached")
+	}
+	if _, err := os.Stat(failPath); err != nil {
+		t.Fatalf("expected a failure marker next to the cache entry: %v", err)
+	}
+}
+
+// TestImageProxyServesRemoteImageWhenCacheDirUnwritable 缓存目录不可写时
+// （磁盘只读/满）必须退回内存缓冲，不能让所有图片变成占位图。
+func TestImageProxyServesRemoteImageWhenCacheDirUnwritable(t *testing.T) {
+	payload := testJPEG
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "image/jpeg")
+		_, _ = w.Write(payload)
+	}))
+	defer upstream.Close()
+
+	// 把一个普通文件当作目录的父级，MkdirAll 必然失败。
+	root := t.TempDir()
+	blocker := filepath.Join(root, "blocker")
+	if err := os.WriteFile(blocker, []byte("x"), 0o644); err != nil {
+		t.Fatalf("write blocker: %v", err)
+	}
+	proxy := newStreamImageProxy(t, upstream.URL)
+	proxy.cacheDir = filepath.Join(blocker, "cache", "images")
+
+	raw := upstream.URL + "/poster.jpg"
+	rec := httptest.NewRecorder()
+	if err := proxy.Serve(t.Context(), rec, httptest.NewRequest(http.MethodGet, "/api/img", nil), raw); err != nil {
+		t.Fatalf("Serve failed: %v", err)
+	}
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200", rec.Code)
+	}
+	body, err := io.ReadAll(rec.Body)
+	if err != nil {
+		t.Fatalf("read body: %v", err)
+	}
+	if !bytes.Equal(body, payload) {
+		t.Fatalf("body length = %d, want %d", len(body), len(payload))
+	}
+}
+
+// TestImageProxyFetchReturnsCachedBytes Fetch 仍需返回字节（刮削写元数据用），
+// 内容来自刚写入的缓存文件。
+func TestImageProxyFetchReturnsCachedBytes(t *testing.T) {
+	payload := bigTestJPEG(t, 256<<10)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "image/jpeg")
+		_, _ = w.Write(payload)
+	}))
+	defer upstream.Close()
+
+	proxy := newStreamImageProxy(t, upstream.URL)
+	raw := upstream.URL + "/poster.jpg"
+
+	data, ctype, err := proxy.Fetch(t.Context(), raw)
+	if err != nil {
+		t.Fatalf("Fetch failed: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("fetched %d bytes, want %d", len(data), len(payload))
+	}
+	if ctype != "image/jpeg" {
+		t.Fatalf("content type = %q, want image/jpeg", ctype)
+	}
+}