memos/server/runner/s3presign/runner.go

package s3presign

import (
	"context"
	"log/slog"
	"time"

	"google.golang.org/protobuf/types/known/timestamppb"

	"github.com/usememos/memos/plugin/storage/s3"
	storepb "github.com/usememos/memos/proto/gen/store"
	"github.com/usememos/memos/store"
)

type Runner struct {
	Store *store.Store
}

func NewRunner(store *store.Store) *Runner {
	return &Runner{
		Store: store,
	}
}

// Schedule runner every 12 hours.
const runnerInterval = time.Hour * 12

func (r *Runner) Run(ctx context.Context) {
	ticker := time.NewTicker(runnerInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			r.RunOnce(ctx)
		case <-ctx.Done():
			return
		}
	}
}

func (r *Runner) RunOnce(ctx context.Context) {
	r.CheckAndPresign(ctx)
}

func (r *Runner) CheckAndPresign(ctx context.Context) {
	workspaceStorageSetting, err := r.Store.GetWorkspaceStorageSetting(ctx)
	if err != nil {
		return
	}

	s3StorageType := storepb.ResourceStorageType_S3
	resources, err := r.Store.ListResources(ctx, &store.FindResource{
		GetBlob:     false,
		StorageType: &s3StorageType,
	})
	if err != nil {
		return
	}

	for _, resource := range resources {
		s3ObjectPayload := resource.Payload.GetS3Object()
		if s3ObjectPayload == nil {
			continue
		}

		if s3ObjectPayload.LastPresignedTime != nil {
			// Skip if the presigned URL is still valid for the next 4 days.
			// The expiration time is set to 5 days.
			if time.Now().Before(s3ObjectPayload.LastPresignedTime.AsTime().Add(4 * 24 * time.Hour)) {
				continue
			}
		}

		s3Config := workspaceStorageSetting.GetS3Config()
		if s3ObjectPayload.S3Config != nil {
			s3Config = s3ObjectPayload.S3Config
		}
		if s3Config == nil {
			slog.Error("S3 config is not found")
			continue
		}

		s3Client, err := s3.NewClient(ctx, s3Config)
		if err != nil {
			slog.Error("Failed to create S3 client", "error", err)
			continue
		}

		presignURL, err := s3Client.PresignGetObject(ctx, s3ObjectPayload.Key)
		if err != nil {
			return
		}
		s3ObjectPayload.S3Config = s3Config
		s3ObjectPayload.LastPresignedTime = timestamppb.New(time.Now())
		if err := r.Store.UpdateResource(ctx, &store.UpdateResource{
			ID:        resource.ID,
			Reference: &presignURL,
			Payload: &storepb.ResourcePayload{
				Payload: &storepb.ResourcePayload_S3Object_{
					S3Object: s3ObjectPayload,
				},
			},
		}); err != nil {
			return
		}
	}
}