memos/server/service/resource_presign/resource_presign.go
2024-04-18 23:34:35 +08:00

110 lines
3.2 KiB
Go

package resourcepresign
import (
"context"
"log/slog"
"strings"
"time"
"github.com/pkg/errors"
"github.com/usememos/memos/plugin/storage/s3"
storepb "github.com/usememos/memos/proto/gen/store"
"github.com/usememos/memos/store"
)
// RunPreSignLinks is a background runner that pre-signs external links stored in the database.
// It uses S3 client to generate presigned URLs and updates the corresponding resources in the store.
func RunPreSignLinks(ctx context.Context, dataStore *store.Store) {
for {
if err := signExternalLinks(ctx, dataStore); err != nil {
slog.Error("failed to pre-sign links", err)
} else {
slog.Debug("pre-signed links")
}
select {
case <-time.After(s3.LinkLifetime / 2):
case <-ctx.Done():
return
}
}
}
func signExternalLinks(ctx context.Context, dataStore *store.Store) error {
objectStore, err := findObjectStorage(ctx, dataStore)
if err != nil {
return errors.Wrapf(err, "find object storage")
}
if objectStore == nil || !objectStore.Config.PreSign {
// object storage not set or not supported
return nil
}
resources, err := dataStore.ListResources(ctx, &store.FindResource{
GetBlob: false,
})
if err != nil {
return errors.Wrapf(err, "list resources")
}
for _, resource := range resources {
if resource.ExternalLink == "" {
// not for object store
continue
}
if strings.Contains(resource.ExternalLink, "?") && time.Since(time.Unix(resource.UpdatedTs, 0)) < s3.LinkLifetime/2 {
// resource not signed (hack for migration)
// resource was recently updated - skipping
continue
}
newLink, err := objectStore.PreSignLink(ctx, resource.ExternalLink)
if err != nil {
slog.Error("failed to pre-sign link", err)
continue
}
now := time.Now().Unix()
if _, err := dataStore.UpdateResource(ctx, &store.UpdateResource{
ID: resource.ID,
UpdatedTs: &now,
ExternalLink: &newLink,
}); err != nil {
return errors.Wrapf(err, "update resource %d link to %q", resource.ID, newLink)
}
}
return nil
}
// findObjectStorage returns current default storage if it's S3-compatible or nil otherwise.
// Returns error only in case of internal problems (ie: database or configuration issues).
// May return nil client and nil error.
func findObjectStorage(ctx context.Context, dataStore *store.Store) (*s3.Client, error) {
workspaceStorageSetting, err := dataStore.GetWorkspaceStorageSetting(ctx)
if err != nil {
return nil, errors.Wrap(err, "Failed to find workspaceStorageSetting")
}
if workspaceStorageSetting.StorageType != storepb.WorkspaceStorageSetting_STORAGE_TYPE_EXTERNAL || workspaceStorageSetting.ActivedExternalStorageId == nil {
return nil, nil
}
storage, err := dataStore.GetStorage(ctx, &store.FindStorage{ID: workspaceStorageSetting.ActivedExternalStorageId})
if err != nil {
return nil, errors.Wrap(err, "Failed to find storage")
}
if storage == nil || storage.Type != storepb.Storage_S3 {
return nil, nil
}
s3Config := storage.Config.GetS3Config()
return s3.NewClient(ctx, &s3.Config{
AccessKey: s3Config.AccessKey,
SecretKey: s3Config.SecretKey,
EndPoint: s3Config.EndPoint,
Region: s3Config.Region,
Bucket: s3Config.Bucket,
URLPrefix: s3Config.UrlPrefix,
URLSuffix: s3Config.UrlSuffix,
PreSign: s3Config.PreSign,
})
}