memos/server/runner/s3presign/runner_test.go

package s3presign

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"google.golang.org/protobuf/types/known/timestamppb"

	"github.com/usememos/memos/internal/storage/s3"
	"github.com/usememos/memos/internal/testutil/fakes3"
	storepb "github.com/usememos/memos/proto/gen/store"
	"github.com/usememos/memos/store"
	teststore "github.com/usememos/memos/store/test"
)

func TestCheckAndPresignRefreshesAndPinsAttachments(t *testing.T) {
	ctx := context.Background()
	fake := fakes3.New(t, "presign-attachments")
	config := fake.Config("presign-attachments")
	ts := teststore.NewTestingStore(ctx, t)
	defer ts.Close()

	storageID := "s3-current"
	_, err := ts.UpsertInstanceSetting(ctx, &storepb.InstanceSetting{
		Key: storepb.InstanceSettingKey_STORAGE,
		Value: &storepb.InstanceSetting_StorageSetting{StorageSetting: &storepb.InstanceStorageSetting{
			Storages: []*storepb.Storage{{
				Id:     storageID,
				Name:   "Current S3",
				Type:   storepb.StorageType_STORAGE_TYPE_S3,
				Config: &storepb.Storage_S3Config{S3Config: config},
			}},
			DefaultStorageId: storageID,
		}},
	})
	require.NoError(t, err)
	user, err := ts.CreateUser(ctx, &store.User{
		Username: "presign-runner-user",
		Role:     store.RoleUser,
		Email:    "presign-runner-user@example.com",
	})
	require.NoError(t, err)

	driver, err := s3.NewDriver(ctx, config)
	require.NoError(t, err)
	content := []byte("legacy attachment refreshed by the runner")
	legacyKey, err := driver.UploadObject(ctx, "legacy/presign.txt", "text/plain", bytes.NewReader(content))
	require.NoError(t, err)

	freshTime := timestamppb.Now()
	createS3Attachment(ctx, t, ts, user.ID, "presign-legacy", legacyKey, "", config, nil, "")
	createS3Attachment(ctx, t, ts, user.ID, "presign-expired", "expired.txt", storageID, nil, timestamppb.New(time.Now().Add(-5*24*time.Hour)), "expired-reference")
	createS3Attachment(ctx, t, ts, user.ID, "presign-fresh", "fresh.txt", storageID, nil, freshTime, "still-valid-reference")
	createS3Attachment(ctx, t, ts, user.ID, "presign-dangling", "dangling.txt", "missing-storage", nil, nil, "dangling-reference")

	NewRunner(ts).RunOnce(ctx)

	legacy := getAttachmentByUID(ctx, t, ts, "presign-legacy")
	require.Equal(t, storageID, legacy.Payload.GetS3Object().GetStorageId())
	require.NotNil(t, legacy.Payload.GetS3Object().GetLastPresignedTime())
	require.Contains(t, legacy.Reference, "X-Amz-Signature=")
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, legacy.Reference, nil)
	require.NoError(t, err)
	response, err := http.DefaultClient.Do(request)
	require.NoError(t, err)
	defer response.Body.Close()
	require.Equal(t, http.StatusOK, response.StatusCode)
	downloaded, err := io.ReadAll(response.Body)
	require.NoError(t, err)
	require.Equal(t, content, downloaded)

	expired := getAttachmentByUID(ctx, t, ts, "presign-expired")
	require.NotEqual(t, "expired-reference", expired.Reference)
	require.Contains(t, expired.Reference, "X-Amz-Signature=")
	require.True(t, expired.Payload.GetS3Object().GetLastPresignedTime().AsTime().After(time.Now().Add(-time.Minute)))

	fresh := getAttachmentByUID(ctx, t, ts, "presign-fresh")
	require.Equal(t, "still-valid-reference", fresh.Reference)
	require.Equal(t, freshTime.AsTime(), fresh.Payload.GetS3Object().GetLastPresignedTime().AsTime())

	fallback := getAttachmentByUID(ctx, t, ts, "presign-dangling")
	require.Contains(t, fallback.Reference, "X-Amz-Signature=")
	require.NotNil(t, fallback.Payload.GetS3Object().GetLastPresignedTime())
	require.Equal(t, "missing-storage", fallback.Payload.GetS3Object().GetStorageId())
}

func TestCheckAndPresignProcessesMultipleBatches(t *testing.T) {
	ctx := context.Background()
	fake := fakes3.New(t, "presign-batches")
	config := fake.Config("presign-batches")
	ts := teststore.NewTestingStore(ctx, t)
	defer ts.Close()

	const storageID = "s3-batches"
	_, err := ts.UpsertInstanceSetting(ctx, &storepb.InstanceSetting{
		Key: storepb.InstanceSettingKey_STORAGE,
		Value: &storepb.InstanceSetting_StorageSetting{StorageSetting: &storepb.InstanceStorageSetting{
			Storages: []*storepb.Storage{{
				Id:     storageID,
				Name:   "Batch S3",
				Type:   storepb.StorageType_STORAGE_TYPE_S3,
				Config: &storepb.Storage_S3Config{S3Config: config},
			}},
			DefaultStorageId: storageID,
		}},
	})
	require.NoError(t, err)
	user, err := ts.CreateUser(ctx, &store.User{
		Username: "presign-batch-user",
		Role:     store.RoleUser,
		Email:    "presign-batch-user@example.com",
	})
	require.NoError(t, err)

	for i := 0; i < 101; i++ {
		uid := fmt.Sprintf("presign-batch-%03d", i)
		createS3Attachment(ctx, t, ts, user.ID, uid, uid+".txt", storageID, nil, nil, "")
	}

	NewRunner(ts).RunOnce(ctx)

	s3StorageType := storepb.AttachmentStorageType_S3
	attachments, err := ts.ListAttachments(ctx, &store.FindAttachment{StorageType: &s3StorageType, SkipDefaultLimit: true})
	require.NoError(t, err)
	require.Len(t, attachments, 101)
	for _, attachment := range attachments {
		require.Contains(t, attachment.Reference, "X-Amz-Signature=", attachment.UID)
		require.NotNil(t, attachment.Payload.GetS3Object().GetLastPresignedTime(), attachment.UID)
	}
}

func createS3Attachment(
	ctx context.Context,
	t *testing.T,
	ts *store.Store,
	creatorID int32,
	uid string,
	key string,
	storageID string,
	embeddedConfig *storepb.StorageS3Config,
	lastPresignedTime *timestamppb.Timestamp,
	reference string,
) {
	t.Helper()
	_, err := ts.CreateAttachment(ctx, &store.Attachment{
		UID:         uid,
		CreatorID:   creatorID,
		Filename:    uid + ".txt",
		Type:        "text/plain",
		Size:        1,
		StorageType: storepb.AttachmentStorageType_S3,
		Reference:   reference,
		Payload: &storepb.AttachmentPayload{Payload: &storepb.AttachmentPayload_S3Object_{
			S3Object: &storepb.AttachmentPayload_S3Object{
				Key:               key,
				StorageId:         storageID,
				S3Config:          embeddedConfig,
				LastPresignedTime: lastPresignedTime,
			},
		}},
	})
	require.NoError(t, err)
}

func getAttachmentByUID(ctx context.Context, t *testing.T, ts *store.Store, uid string) *store.Attachment {
	t.Helper()
	attachment, err := ts.GetAttachment(ctx, &store.FindAttachment{UID: &uid})
	require.NoError(t, err)
	require.NotNil(t, attachment)
	return attachment
}

func TestCloneAttachmentPayloadPreservesMotionMedia(t *testing.T) {
	payload := &storepb.AttachmentPayload{
		Payload: &storepb.AttachmentPayload_S3Object_{
			S3Object: &storepb.AttachmentPayload_S3Object{
				Key: "photos/live.jpg",
			},
		},
		MotionMedia: &storepb.MotionMedia{
			Family:  storepb.MotionMediaFamily_ANDROID_MOTION_PHOTO,
			Role:    storepb.MotionMediaRole_CONTAINER,
			GroupId: "motion-group",
		},
	}

	cloned := cloneAttachmentPayload(payload)
	require.NotNil(t, cloned)
	require.NotSame(t, payload, cloned)
	require.Equal(t, payload.MotionMedia, cloned.MotionMedia)

	cloned.GetS3Object().LastPresignedTime = timestamppb.Now()
	require.Nil(t, payload.GetS3Object().LastPresignedTime)
	require.Equal(t, "motion-group", cloned.MotionMedia.GroupId)
}