chorus/internal/platform/storage/local_test.go

package storage

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"image"
	"image/color"
	"image/png"
	"io"
	"os"
	"path/filepath"
	"testing"

	corestorage "git.ilapage.cn/OPC/chorus/internal/core/storage"
)

func newTestStore(t *testing.T, maxBytes int64, maxPixels uint64) *Local {
	t.Helper()
	store, err := NewLocal(Config{
		Root: t.TempDir(), MaxObjectBytes: maxBytes, MaxImagePixels: maxPixels, ThumbnailMaxSide: 256,
		AllowedImageMIME: map[string]bool{"image/png": true, "image/jpeg": true, "image/webp": true},
	})
	if err != nil {
		t.Fatal(err)
	}
	return store
}

func TestPutWebPCreatesPNGThumbnail(t *testing.T) {
	data, err := base64.StdEncoding.DecodeString("UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEADsD+JaQAA3AAAAAA")
	if err != nil {
		t.Fatal(err)
	}
	store := newTestStore(t, 1024, 100)
	objects, err := store.PutImage(context.Background(), ImageRequest{Key: "webp/original", ThumbnailKey: "webp/thumbnail", OwnerID: 1, GenerationID: 1, ContentType: "image/webp", Source: bytes.NewReader(data)})
	if err != nil {
		t.Fatal(err)
	}
	if objects.Original.ContentType != "image/webp" || objects.Thumbnail.ContentType != "image/png" {
		t.Fatalf("objects=%#v", objects)
	}
}

func TestLocalPutOpenAtomicMetadataAndTraversal(t *testing.T) {
	store := newTestStore(t, 1024, 1_000_000)
	object, err := store.Put(context.Background(), corestorage.PutRequest{
		Key: "users/7/generations/9/result", OwnerID: 7, GenerationID: 9,
		ContentType: "text/plain", Source: bytes.NewBufferString("synthetic result"),
	})
	if err != nil {
		t.Fatal(err)
	}
	if object.OwnerID != 7 || object.GenerationID != 9 || object.Size != 16 {
		t.Fatalf("unexpected object: %#v", object)
	}
	reader, metadata, err := store.Open(context.Background(), object.Key)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
	data, _ := io.ReadAll(reader)
	if string(data) != "synthetic result" || metadata.OwnerID != 7 {
		t.Fatalf("open data=%q metadata=%#v", data, metadata)
	}
	if _, err := os.Stat(filepath.Join(store.root, "objects", "users", "7", "generations", "9", "result", "metadata.json")); err != nil {
		t.Fatalf("metadata not atomically placed: %v", err)
	}
	for _, key := range []string{"../outside", "/absolute", "a\\b", "C:/escape", "a//b"} {
		_, err := store.Put(context.Background(), corestorage.PutRequest{Key: key, OwnerID: 1, GenerationID: 1, ContentType: "x", Source: bytes.NewReader(nil)})
		if !errors.Is(err, ErrInvalidKey) {
			t.Errorf("Put(%q) error = %v", key, err)
		}
	}
}

func TestLocalFailureCleansTemporaryFiles(t *testing.T) {
	store := newTestStore(t, 4, 1_000_000)
	_, err := store.Put(context.Background(), corestorage.PutRequest{
		Key: "too-large", OwnerID: 1, GenerationID: 1, ContentType: "text/plain", Source: bytes.NewBufferString("12345"),
	})
	if !errors.Is(err, ErrObjectTooLarge) {
		t.Fatalf("Put() error = %v", err)
	}
	entries, err := os.ReadDir(store.temporaryRoot)
	if err != nil || len(entries) != 0 {
		t.Fatalf("temporary directory not cleaned: entries=%v error=%v", entries, err)
	}
	if _, err := os.Stat(filepath.Join(store.root, "objects", "too-large")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("failed object unexpectedly exists: %v", err)
	}
}

func TestLocalRejectsSymlinkComponents(t *testing.T) {
	store := newTestStore(t, 1024, 1_000_000)
	outside := t.TempDir()
	objectsRoot := filepath.Join(store.root, "objects")
	if err := os.Mkdir(objectsRoot, 0o700); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(outside, filepath.Join(objectsRoot, "linked")); err != nil {
		t.Skipf("symlinks unavailable on this platform: %v", err)
	}
	_, err := store.Put(context.Background(), corestorage.PutRequest{
		Key: "linked/escape", OwnerID: 1, GenerationID: 1, ContentType: "text/plain", Source: bytes.NewBufferString("x"),
	})
	if !errors.Is(err, ErrInvalidKey) {
		t.Fatalf("symlink component error = %v", err)
	}
	if _, err := os.Stat(filepath.Join(outside, "escape")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("storage wrote through symlink: %v", err)
	}
}

func TestPutImageValidatesPixelsMIMEAndCreatesThumbnail(t *testing.T) {
	store := newTestStore(t, 2<<20, 512*512)
	imageData := encodePNG(t, 512, 256)
	objects, err := store.PutImage(context.Background(), ImageRequest{
		Key: "images/original", ThumbnailKey: "images/thumb", OwnerID: 2, GenerationID: 3,
		ContentType: "image/png", Source: bytes.NewReader(imageData),
	})
	if err != nil {
		t.Fatal(err)
	}
	reader, _, err := store.Open(context.Background(), objects.Thumbnail.Key)
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
	config, format, err := image.DecodeConfig(reader)
	if err != nil || format != "png" || config.Width != 256 || config.Height != 128 {
		t.Fatalf("thumbnail config=%#v format=%s error=%v", config, format, err)
	}

	_, err = store.PutImage(context.Background(), ImageRequest{
		Key: "images/mismatch", ThumbnailKey: "images/mismatch-thumb", OwnerID: 2, GenerationID: 3,
		ContentType: "image/jpeg", Source: bytes.NewReader(imageData),
	})
	if !errors.Is(err, ErrUnsupportedImage) {
		t.Fatalf("MIME mismatch error = %v", err)
	}

	smallLimit := newTestStore(t, 2<<20, 10_000)
	_, err = smallLimit.PutImage(context.Background(), ImageRequest{
		Key: "images/bomb", ThumbnailKey: "images/bomb-thumb", OwnerID: 2, GenerationID: 3,
		ContentType: "image/png", Source: bytes.NewReader(imageData),
	})
	if !errors.Is(err, ErrImageTooLarge) {
		t.Fatalf("pixel limit error = %v", err)
	}
}

func encodePNG(t *testing.T, width, height int) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			img.Set(x, y, color.RGBA{R: uint8(x), G: uint8(y), B: 100, A: 255})
		}
	}
	var data bytes.Buffer
	if err := png.Encode(&data, img); err != nil {
		t.Fatal(err)
	}
	return data.Bytes()
}