512 lines
16 KiB
Go
512 lines
16 KiB
Go
// Package assets owns asset visibility, metadata, and blob lifecycle policy.
|
|
package assets
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"html"
|
|
"io"
|
|
"path"
|
|
"regexp"
|
|
"strings"
|
|
"time"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/publicapi"
|
|
)
|
|
|
|
const PublicJobLookupLimit = 200
|
|
|
|
var (
|
|
ErrNotFound = errors.New("asset not found")
|
|
ErrBlobNotFound = errors.New("asset blob not found")
|
|
ErrUnsafeBlobKey = errors.New("unsafe blob key")
|
|
ErrRemoteProtocol = errors.New("remote asset protocol is not allowed")
|
|
ErrRemoteTooLarge = errors.New("remote asset exceeds size limit")
|
|
)
|
|
|
|
type Kind string
|
|
|
|
const (
|
|
KindImage Kind = "image"
|
|
KindVideo Kind = "video"
|
|
KindMask Kind = "mask"
|
|
KindReference Kind = "reference"
|
|
KindOther Kind = "other"
|
|
)
|
|
|
|
type Source string
|
|
|
|
const (
|
|
SourceUpload Source = "upload"
|
|
SourceGenerated Source = "generated"
|
|
SourceEdited Source = "edited"
|
|
SourceUpscaled Source = "upscaled"
|
|
SourceExternal Source = "external"
|
|
SourceSeed Source = "seed"
|
|
)
|
|
|
|
type Asset struct {
|
|
ID string `json:"id"`
|
|
OwnerID string `json:"ownerId"`
|
|
Kind Kind `json:"kind"`
|
|
Name string `json:"name"`
|
|
URL string `json:"url"`
|
|
StoragePath string `json:"storagePath,omitempty"`
|
|
Source Source `json:"source"`
|
|
Tags []string `json:"tags"`
|
|
Metadata map[string]any `json:"metadata"`
|
|
CreatedAt time.Time `json:"createdAt"`
|
|
UpdatedAt time.Time `json:"updatedAt"`
|
|
}
|
|
|
|
// MarshalJSON preserves JavaScript Date.toISOString() millisecond precision.
|
|
func (a Asset) MarshalJSON() ([]byte, error) {
|
|
type wireAsset struct {
|
|
ID string `json:"id"`
|
|
OwnerID string `json:"ownerId"`
|
|
Kind Kind `json:"kind"`
|
|
Name string `json:"name"`
|
|
URL string `json:"url"`
|
|
StoragePath string `json:"storagePath,omitempty"`
|
|
Source Source `json:"source"`
|
|
Tags []string `json:"tags"`
|
|
Metadata map[string]any `json:"metadata"`
|
|
CreatedAt string `json:"createdAt"`
|
|
UpdatedAt string `json:"updatedAt"`
|
|
}
|
|
return json.Marshal(wireAsset{
|
|
ID: a.ID, OwnerID: a.OwnerID, Kind: a.Kind, Name: a.Name, URL: a.URL,
|
|
StoragePath: a.StoragePath, Source: a.Source, Tags: a.Tags, Metadata: a.Metadata,
|
|
CreatedAt: a.CreatedAt.UTC().Format("2006-01-02T15:04:05.000Z"),
|
|
UpdatedAt: a.UpdatedAt.UTC().Format("2006-01-02T15:04:05.000Z"),
|
|
})
|
|
}
|
|
|
|
type scopeKind uint8
|
|
|
|
const (
|
|
platformScope scopeKind = iota
|
|
publicScope
|
|
)
|
|
|
|
type Scope struct {
|
|
kind scopeKind
|
|
ownerID, clientID string
|
|
}
|
|
|
|
func PlatformScope(ownerID string) Scope { return Scope{kind: platformScope, ownerID: ownerID} }
|
|
func PublicScope(clientID string) Scope {
|
|
return Scope{kind: publicScope, ownerID: publicapi.OwnerID(clientID), clientID: clientID}
|
|
}
|
|
func ClientTag(clientID string) string { return "api-client:" + clientID }
|
|
|
|
type Catalog interface {
|
|
ListOwner(context.Context, string) ([]Asset, error)
|
|
GetOwner(context.Context, string, string) (Asset, bool, error)
|
|
ListPublic(context.Context, string, string, int) ([]Asset, error)
|
|
GetPublic(context.Context, string, string, string, int) (Asset, bool, error)
|
|
Create(context.Context, Asset) (Asset, error)
|
|
DeleteOwner(context.Context, string, string) (Asset, bool, error)
|
|
}
|
|
|
|
type storagePathCatalog interface {
|
|
GetOwnerByStoragePath(context.Context, string, string) (Asset, bool, error)
|
|
}
|
|
|
|
type StoredObject struct{ Key, URL string }
|
|
type Blob struct {
|
|
Body io.ReadCloser
|
|
ContentType string
|
|
Size int64
|
|
}
|
|
type BlobStore interface {
|
|
Put(context.Context, string, io.Reader, int64, string) (StoredObject, error)
|
|
Read(context.Context, string) (Blob, error)
|
|
Delete(context.Context, string) error
|
|
}
|
|
type RemoteFetcher interface {
|
|
Fetch(context.Context, string) (Blob, error)
|
|
}
|
|
|
|
type Service struct {
|
|
catalog Catalog
|
|
blobs BlobStore
|
|
remote RemoteFetcher
|
|
now func() time.Time
|
|
id func(string) string
|
|
}
|
|
|
|
func NewService(c Catalog, b BlobStore, r RemoteFetcher, now func() time.Time, id func(string) string) *Service {
|
|
if now == nil {
|
|
now = time.Now
|
|
}
|
|
if id == nil {
|
|
id = func(prefix string) string { return fmt.Sprintf("%s-%d", prefix, now().UnixNano()) }
|
|
}
|
|
return &Service{catalog: c, blobs: b, remote: r, now: now, id: id}
|
|
}
|
|
|
|
func (s *Service) List(ctx context.Context, scope Scope) ([]Asset, error) {
|
|
if err := validScope(scope); err != nil {
|
|
return nil, err
|
|
}
|
|
if scope.kind == publicScope {
|
|
return s.catalog.ListPublic(ctx, scope.ownerID, scope.clientID, PublicJobLookupLimit)
|
|
}
|
|
return s.catalog.ListOwner(ctx, scope.ownerID)
|
|
}
|
|
func (s *Service) Get(ctx context.Context, scope Scope, id string) (Asset, error) {
|
|
if err := validScope(scope); err != nil {
|
|
return Asset{}, err
|
|
}
|
|
var a Asset
|
|
var found bool
|
|
var err error
|
|
if scope.kind == publicScope {
|
|
a, found, err = s.catalog.GetPublic(ctx, scope.ownerID, scope.clientID, id, PublicJobLookupLimit)
|
|
} else {
|
|
a, found, err = s.catalog.GetOwner(ctx, scope.ownerID, id)
|
|
}
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
if !found {
|
|
return Asset{}, ErrNotFound
|
|
}
|
|
return a, nil
|
|
}
|
|
|
|
type CreateExternalCommand struct {
|
|
URL, Name string
|
|
Kind Kind
|
|
Source Source
|
|
Tags []string
|
|
}
|
|
|
|
func (s *Service) CreateExternal(ctx context.Context, scope Scope, cmd CreateExternalCommand) (Asset, error) {
|
|
if err := validScope(scope); err != nil {
|
|
return Asset{}, err
|
|
}
|
|
if strings.TrimSpace(cmd.URL) == "" {
|
|
return Asset{}, errors.New("url is required")
|
|
}
|
|
now := s.now().UTC()
|
|
kind := cmd.Kind
|
|
if kind == "" {
|
|
kind = KindImage
|
|
}
|
|
source := cmd.Source
|
|
if source == "" {
|
|
source = SourceExternal
|
|
}
|
|
name := cmd.Name
|
|
tags := cloneStrings(cmd.Tags)
|
|
metadata := map[string]any{}
|
|
if scope.kind == publicScope {
|
|
name = defaultString(name, "外部素材")
|
|
tags = appendUnique(tags, "external", "public-api", ClientTag(scope.clientID))
|
|
metadata["registeredFrom"] = "public-api"
|
|
metadata["externalClientId"] = scope.clientID
|
|
} else {
|
|
name = defaultString(name, "外部图片")
|
|
if cmd.Tags == nil {
|
|
tags = []string{"external"}
|
|
}
|
|
metadata["registeredFrom"] = "api"
|
|
}
|
|
a := Asset{ID: s.id("asset"), OwnerID: scope.ownerID, Kind: kind, Name: name, URL: cmd.URL, Source: source, Tags: tags, Metadata: metadata, CreatedAt: now, UpdatedAt: now}
|
|
return s.catalog.Create(ctx, a)
|
|
}
|
|
|
|
type UploadCommand struct {
|
|
Bytes []byte
|
|
FileName, ContentType, Origin string
|
|
Kind Kind
|
|
Tags []string
|
|
}
|
|
|
|
type ImportGeneratedCommand struct {
|
|
URL, Name, Capability, JobID string
|
|
Kind Kind
|
|
Source Source
|
|
Tags []string
|
|
Metadata map[string]any
|
|
}
|
|
|
|
// ImportMockCommand describes a locally generated development result. Unlike
|
|
// ImportGeneratedCommand it carries no URL, so mock output never enters the
|
|
// remote-fetch/SSRF boundary.
|
|
type ImportMockCommand struct {
|
|
Name, Capability, JobID string
|
|
Kind Kind
|
|
Tags []string
|
|
Metadata map[string]any
|
|
}
|
|
|
|
func (s *Service) Upload(ctx context.Context, scope Scope, cmd UploadCommand) (Asset, error) {
|
|
if err := validScope(scope); err != nil {
|
|
return Asset{}, err
|
|
}
|
|
if s.blobs == nil {
|
|
return Asset{}, errors.New("blob store is unavailable")
|
|
}
|
|
name := sanitizeFileName(cmd.FileName)
|
|
key := path.Join("uploads", s.now().UTC().Format("2006-01-02"), s.id("file")+"-"+name)
|
|
stored, err := s.blobs.Put(ctx, key, bytes.NewReader(cmd.Bytes), int64(len(cmd.Bytes)), cmd.ContentType)
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
kind := cmd.Kind
|
|
if kind == "" {
|
|
kind = inferKind(cmd.ContentType)
|
|
}
|
|
tags := cloneStrings(cmd.Tags)
|
|
if scope.kind == publicScope {
|
|
tags = appendUnique(tags, "upload", "public-api", ClientTag(scope.clientID))
|
|
} else if cmd.Tags == nil {
|
|
tags = []string{"upload"}
|
|
}
|
|
now := s.now().UTC()
|
|
a := Asset{ID: s.id("asset"), OwnerID: scope.ownerID, Kind: kind, Name: cmd.FileName, URL: stored.URL, StoragePath: stored.Key, Source: SourceUpload, Tags: tags, Metadata: map[string]any{"contentType": cmd.ContentType, "size": len(cmd.Bytes)}, CreatedAt: now, UpdatedAt: now}
|
|
created, err := s.catalog.Create(ctx, a)
|
|
if err != nil {
|
|
_ = s.blobs.Delete(context.WithoutCancel(ctx), stored.Key)
|
|
return Asset{}, err
|
|
}
|
|
return created, nil
|
|
}
|
|
|
|
// ImportGenerated fetches a provider result through the bounded RemoteFetcher,
|
|
// stores it in the configured BlobStore, and only then creates relational
|
|
// metadata. Upload already owns blob compensation when catalog persistence
|
|
// fails, so provider outputs share the same crash-safe write ordering.
|
|
func (s *Service) ImportGenerated(ctx context.Context, scope Scope, cmd ImportGeneratedCommand) (Asset, error) {
|
|
if strings.TrimSpace(cmd.URL) == "" || s.remote == nil {
|
|
return Asset{}, ErrBlobNotFound
|
|
}
|
|
blob, err := s.remote.Fetch(ctx, cmd.URL)
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
defer blob.Body.Close()
|
|
maximum := blob.Size
|
|
if maximum < 0 {
|
|
maximum = 64 << 20
|
|
}
|
|
content, err := io.ReadAll(io.LimitReader(blob.Body, maximum+1))
|
|
if err != nil || int64(len(content)) > maximum {
|
|
return Asset{}, ErrRemoteTooLarge
|
|
}
|
|
name := cmd.Name
|
|
if strings.TrimSpace(name) == "" {
|
|
name = path.Base(strings.SplitN(cmd.URL, "?", 2)[0])
|
|
}
|
|
if s.blobs == nil {
|
|
return Asset{}, errors.New("blob store is unavailable")
|
|
}
|
|
cleanName := sanitizeFileName(name)
|
|
key := path.Join("generated-results", s.now().UTC().Format("2006-01-02"), s.id("file")+"-"+cleanName)
|
|
stored, err := s.blobs.Put(ctx, key, bytes.NewReader(content), int64(len(content)), blob.ContentType)
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
kind := cmd.Kind
|
|
if kind == "" {
|
|
kind = inferKind(blob.ContentType)
|
|
}
|
|
now := s.now().UTC()
|
|
source := cmd.Source
|
|
if source == "" {
|
|
source = SourceGenerated
|
|
}
|
|
metadata := cloneMetadata(cmd.Metadata)
|
|
metadata["contentType"] = blob.ContentType
|
|
metadata["size"] = len(content)
|
|
metadata["capability"] = cmd.Capability
|
|
metadata["jobId"] = cmd.JobID
|
|
metadata["importedFrom"] = cmd.URL
|
|
asset := Asset{ID: s.id("asset"), OwnerID: scope.ownerID, Kind: kind, Name: name, URL: stored.URL, StoragePath: stored.Key, Source: source, Tags: cloneStrings(cmd.Tags), Metadata: metadata, CreatedAt: now, UpdatedAt: now}
|
|
created, err := s.catalog.Create(ctx, asset)
|
|
if err != nil {
|
|
_ = s.blobs.Delete(context.WithoutCancel(ctx), stored.Key)
|
|
return Asset{}, err
|
|
}
|
|
return created, nil
|
|
}
|
|
|
|
// ImportMock stores a small deterministic placeholder in the configured blob
|
|
// store and registers it as a generated asset. It is intentionally a separate
|
|
// path from remote imports: relative mock provider URLs are never interpreted
|
|
// as fetch destinations.
|
|
func (s *Service) ImportMock(ctx context.Context, scope Scope, cmd ImportMockCommand) (Asset, error) {
|
|
if err := validScope(scope); err != nil {
|
|
return Asset{}, err
|
|
}
|
|
if s.blobs == nil {
|
|
return Asset{}, errors.New("blob store is unavailable")
|
|
}
|
|
kind := cmd.Kind
|
|
if kind == "" {
|
|
kind = KindImage
|
|
}
|
|
name, contentType, content := mockOutput(cmd.Name, kind, cmd.JobID)
|
|
key := path.Join("generated-results", s.now().UTC().Format("2006-01-02"), s.id("file")+"-"+sanitizeFileName(name))
|
|
stored, err := s.blobs.Put(ctx, key, bytes.NewReader(content), int64(len(content)), contentType)
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
now := s.now().UTC()
|
|
metadata := cloneMetadata(cmd.Metadata)
|
|
metadata["contentType"] = contentType
|
|
metadata["size"] = len(content)
|
|
metadata["capability"] = cmd.Capability
|
|
metadata["jobId"] = cmd.JobID
|
|
metadata["mock"] = true
|
|
asset := Asset{ID: s.id("asset"), OwnerID: scope.ownerID, Kind: kind, Name: name, URL: stored.URL, StoragePath: stored.Key, Source: SourceGenerated, Tags: cloneStrings(cmd.Tags), Metadata: metadata, CreatedAt: now, UpdatedAt: now}
|
|
created, err := s.catalog.Create(ctx, asset)
|
|
if err != nil {
|
|
_ = s.blobs.Delete(context.WithoutCancel(ctx), stored.Key)
|
|
return Asset{}, err
|
|
}
|
|
return created, nil
|
|
}
|
|
|
|
func mockOutput(name string, kind Kind, jobID string) (string, string, []byte) {
|
|
if kind == KindVideo {
|
|
return defaultString(name, "mock-video.mp4"), "video/mp4", []byte("mock video result for " + jobID + "\n")
|
|
}
|
|
name = defaultString(name, "mock-image.svg")
|
|
content := `<svg xmlns="http://www.w3.org/2000/svg" width="1024" height="1024" viewBox="0 0 1024 1024"><rect width="1024" height="1024" fill="#111827"/><text x="512" y="492" text-anchor="middle" fill="#f9fafb" font-family="sans-serif" font-size="64">Mock image</text><text x="512" y="558" text-anchor="middle" fill="#9ca3af" font-family="sans-serif" font-size="28">` + html.EscapeString(jobID) + `</text></svg>`
|
|
return name, "image/svg+xml", []byte(content)
|
|
}
|
|
|
|
func (s *Service) Delete(ctx context.Context, scope Scope, id string) (Asset, error) {
|
|
a, err := s.Get(ctx, scope, id)
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
if a.StoragePath != "" && s.blobs != nil {
|
|
if err := s.blobs.Delete(ctx, a.StoragePath); err != nil && !errors.Is(err, ErrBlobNotFound) {
|
|
return Asset{}, err
|
|
}
|
|
}
|
|
deleted, found, err := s.catalog.DeleteOwner(ctx, scope.ownerID, id)
|
|
if err != nil {
|
|
return Asset{}, err
|
|
}
|
|
if !found {
|
|
return Asset{}, ErrNotFound
|
|
}
|
|
return deleted, nil
|
|
}
|
|
func (s *Service) Download(ctx context.Context, scope Scope, id string) (Blob, error) {
|
|
a, err := s.Get(ctx, scope, id)
|
|
if err != nil {
|
|
return Blob{}, err
|
|
}
|
|
if a.StoragePath != "" && s.blobs != nil {
|
|
b, e := s.blobs.Read(ctx, a.StoragePath)
|
|
if e == nil {
|
|
return b, nil
|
|
}
|
|
if !errors.Is(e, ErrBlobNotFound) {
|
|
return Blob{}, e
|
|
}
|
|
}
|
|
if s.remote == nil {
|
|
return Blob{}, ErrBlobNotFound
|
|
}
|
|
return s.remote.Fetch(ctx, a.URL)
|
|
}
|
|
|
|
// DownloadPath serves a stored object only after its metadata has been found
|
|
// in the requesting owner's catalog. Catalogs that support HTTP file serving
|
|
// implement the optional storage-path lookup without widening other callers.
|
|
func (s *Service) DownloadPath(ctx context.Context, scope Scope, storagePath string) (Blob, error) {
|
|
if err := validScope(scope); err != nil {
|
|
return Blob{}, err
|
|
}
|
|
lookup, ok := s.catalog.(storagePathCatalog)
|
|
if !ok {
|
|
return Blob{}, ErrNotFound
|
|
}
|
|
a, found, err := lookup.GetOwnerByStoragePath(ctx, scope.ownerID, storagePath)
|
|
if err != nil {
|
|
return Blob{}, err
|
|
}
|
|
if !found || a.StoragePath == "" {
|
|
return Blob{}, ErrNotFound
|
|
}
|
|
if s.blobs == nil {
|
|
return Blob{}, ErrBlobNotFound
|
|
}
|
|
return s.blobs.Read(ctx, a.StoragePath)
|
|
}
|
|
|
|
func validScope(s Scope) error {
|
|
if strings.TrimSpace(s.ownerID) == "" {
|
|
return errors.New("asset owner is required")
|
|
}
|
|
if s.kind == publicScope && strings.TrimSpace(s.clientID) == "" {
|
|
return errors.New("public client is required")
|
|
}
|
|
return nil
|
|
}
|
|
func defaultString(v, d string) string {
|
|
if strings.TrimSpace(v) == "" {
|
|
return d
|
|
}
|
|
return v
|
|
}
|
|
func appendUnique(values []string, add ...string) []string {
|
|
for _, v := range add {
|
|
if !contains(values, v) {
|
|
values = append(values, v)
|
|
}
|
|
}
|
|
return values
|
|
}
|
|
func cloneStrings(values []string) []string {
|
|
if values == nil {
|
|
return nil
|
|
}
|
|
return append([]string{}, values...)
|
|
}
|
|
func cloneMetadata(values map[string]any) map[string]any {
|
|
result := make(map[string]any, len(values)+5)
|
|
for key, value := range values {
|
|
result[key] = value
|
|
}
|
|
return result
|
|
}
|
|
func contains(values []string, want string) bool {
|
|
for _, v := range values {
|
|
if v == want {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
var unsafeName = regexp.MustCompile(`[^A-Za-z0-9._-]+`)
|
|
|
|
func sanitizeFileName(v string) string {
|
|
v = strings.Trim(unsafeName.ReplaceAllString(v, "-"), "-")
|
|
if v == "" {
|
|
return "asset.bin"
|
|
}
|
|
return v
|
|
}
|
|
func inferKind(contentType string) Kind {
|
|
if strings.HasPrefix(contentType, "image/") {
|
|
return KindImage
|
|
}
|
|
if strings.HasPrefix(contentType, "video/") {
|
|
return KindVideo
|
|
}
|
|
return KindOther
|
|
}
|