spark-server/internal/module/upload/logic/credential_test.go

package logic

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"spark-server/internal/module/upload/entity"
	"spark-server/internal/pkg/errs"
	"spark-server/internal/pkg/reqctx"
)

type fakeStore struct {
	put, get string
	putErr   error
}

func (f *fakeStore) PresignPut(_ context.Context, key, _ string, _ int64, _ time.Duration) (string, error) {
	if f.putErr != nil {
		return "", f.putErr
	}
	if f.put != "" {
		return f.put, nil
	}
	return "http://rustfs.test/" + key + "?sig=1", nil
}

func (f *fakeStore) PresignGet(_ context.Context, key string, _ time.Duration) (string, error) {
	if f.get != "" {
		return f.get, nil
	}
	return "http://rustfs.test/" + key + "?get=1", nil
}

func testUpload(t *testing.T) *Service {
	t.Helper()
	return NewService(&fakeStore{}, Options{
		TTL: 5 * time.Minute,
		Now: func() time.Time { return time.Date(2026, 9, 17, 8, 0, 0, 0, time.UTC) },
	})
}

func TestCreateCredentialChatVsAvatar(t *testing.T) {
	svc := testUpload(t)
	ctx := reqctx.With(t.Context(), &reqctx.RequestContext{UserID: 42, Language: "zh-CN"})
	chat, err := svc.CreateCredential(ctx, 42, &CreateInput{Purpose: entity.PurposeChatImage, Mime: "image/jpeg", Size: 1024})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(chat.ObjectKey, "chat/image/20260917/") || chat.Method != "PUT" ||
		chat.Headers["Content-Type"] != "image/jpeg" || chat.Headers["Content-Length"] != "1024" ||
		chat.Headers["If-None-Match"] != "*" || strings.Contains(chat.ObjectKey, "/42/") {
		t.Fatalf("%+v", chat)
	}
	avatar, err := svc.CreateCredential(ctx, 42, &CreateInput{Purpose: entity.PurposeAvatar, Mime: "image/png", Size: 2048})
	if err != nil || !strings.HasPrefix(avatar.ObjectKey, "perm/avatar/20260917/") {
		t.Fatalf("%+v %v", avatar, err)
	}
	group, err := svc.CreateCredential(ctx, 42, &CreateInput{Purpose: entity.PurposeGroupAvatar, Mime: "image/jpeg", Size: 1024, GroupID: 7})
	if err != nil || !strings.HasPrefix(group.ObjectKey, "perm/group/7/") {
		t.Fatalf("%+v %v", group, err)
	}
	if _, err := svc.CreateCredential(ctx, 42, &CreateInput{Purpose: entity.PurposeMomentImage, Mime: "image/gif", Size: 2048}); err != nil {
		t.Fatalf("moment gif %v", err)
	}
	if _, err := svc.CreateCredential(ctx, 42, &CreateInput{Purpose: entity.PurposeGroupAvatar, Mime: "image/jpeg", Size: 1024}); err == nil {
		t.Fatal("group avatar without group_id")
	}
}

func TestCreateCredentialRejectsBadMIMEAndSize(t *testing.T) {
	svc := testUpload(t)
	ctx := reqctx.With(t.Context(), &reqctx.RequestContext{UserID: 1, Language: "zh-CN"})
	if _, err := svc.CreateCredential(ctx, 1, &CreateInput{Purpose: entity.PurposeAvatar, Mime: "application/pdf", Size: 100}); err == nil {
		t.Fatal("pdf avatar")
	}
	if _, err := svc.CreateCredential(ctx, 1, &CreateInput{Purpose: entity.PurposeChatImage, Mime: "image/jpeg", Size: 50 << 20}); err == nil {
		t.Fatal("huge image")
	}
	for _, mime := range []string{"image/heic", "image/heif", "image/webp"} {
		if _, err := svc.CreateCredential(ctx, 1, &CreateInput{Purpose: entity.PurposeChatImage, Mime: mime, Size: 1024}); err == nil {
			t.Fatalf("accepted %s", mime)
		}
	}
	if _, err := svc.CreateCredential(ctx, 0, &CreateInput{Purpose: entity.PurposeChatImage, Mime: "image/jpeg", Size: 1}); !errors.Is(err, errs.ErrUnauthorized) {
		t.Fatalf("uid %v", err)
	}
}

func TestCreateDownloadValidatesKey(t *testing.T) {
	svc := testUpload(t)
	ctx := reqctx.With(t.Context(), &reqctx.RequestContext{UserID: 1, Language: "en"})
	got, err := svc.CreateDownload(ctx, 1, &DownloadInput{ObjectKey: "chat/20260917/1/2.jpg"})
	if err != nil || got.Method != "GET" || !strings.Contains(got.DownloadURL, "chat/20260917/1/2.jpg") {
		t.Fatalf("%+v %v", got, err)
	}
	if _, err := svc.CreateDownload(ctx, 1, &DownloadInput{ObjectKey: "../secret"}); err == nil {
		t.Fatal("path escape")
	}
}

func TestCreateDownloadsBatch(t *testing.T) {
	svc := testUpload(t)
	ctx := reqctx.With(t.Context(), &reqctx.RequestContext{UserID: 1, Language: "zh-CN"})
	keys := []string{
		"perm/moment/image/20260918/a.jpg",
		"perm/moment/image/20260918/b.jpg",
		"chat/image/20260918/c.jpg",
	}
	got, err := svc.CreateDownloads(ctx, 1, &BatchDownloadInput{ObjectKeys: keys})
	if err != nil || got == nil || len(got.Items) != 3 {
		t.Fatalf("%+v %v", got, err)
	}
	if got.Items[0].ObjectKey != keys[0] || got.Items[1].Method != "GET" || !strings.Contains(got.Items[2].DownloadURL, keys[2]) {
		t.Fatalf("%+v", got.Items)
	}
	if _, err := svc.CreateDownloads(ctx, 1, &BatchDownloadInput{}); err == nil {
		t.Fatal("empty")
	}
	tooMany := make([]string, 10)
	for i := range tooMany {
		tooMany[i] = "chat/image/20260918/x.jpg"
	}
	if _, err := svc.CreateDownloads(ctx, 1, &BatchDownloadInput{ObjectKeys: tooMany}); err == nil {
		t.Fatal("too many")
	}
	if _, err := svc.CreateDownloads(ctx, 1, &BatchDownloadInput{ObjectKeys: []string{"../secret"}}); err == nil {
		t.Fatal("path escape")
	}
}

func TestCreateCredentialsBatch(t *testing.T) {
	svc := testUpload(t)
	ctx := reqctx.With(t.Context(), &reqctx.RequestContext{UserID: 42, Language: "zh-CN"})
	part := func(size int64) CreateInput {
		return CreateInput{Purpose: entity.PurposeChatImage, Mime: "image/jpeg", Size: size}
	}
	got, err := svc.CreateCredentials(ctx, 42, &BatchCreateInput{Items: []CreateInput{part(1024), part(2048), part(4096)}})
	if err != nil || got == nil || len(got.Items) != 3 {
		t.Fatalf("%+v %v", got, err)
	}
	if got.Items[0].ObjectKey == got.Items[1].ObjectKey || got.Items[1].Headers["Content-Length"] != "2048" {
		t.Fatalf("%+v", got.Items)
	}
	if _, err := svc.CreateCredentials(ctx, 42, &BatchCreateInput{}); err == nil {
		t.Fatal("empty")
	}
	if _, err := svc.CreateCredentials(ctx, 42, &BatchCreateInput{Items: []CreateInput{part(1), part(1), part(1), part(1)}}); err == nil {
		t.Fatal("too many")
	}
}