Files
NianAIGC/backend/internal/httpapi/assets_test.go
2026-08-19 10:52:55 +08:00

217 lines
9.0 KiB
Go

package httpapi_test
import (
"bytes"
"context"
"encoding/json"
"io"
"mime/multipart"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets"
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/httpapi"
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/publicapi"
)
type assetCatalog struct{ values []assets.Asset }
func (c *assetCatalog) ListOwner(_ context.Context, owner string) ([]assets.Asset, error) {
var v []assets.Asset
for _, a := range c.values {
if a.OwnerID == owner {
v = append(v, a)
}
}
return v, nil
}
func (c *assetCatalog) GetOwner(_ context.Context, owner, id string) (assets.Asset, bool, error) {
for _, a := range c.values {
if a.OwnerID == owner && a.ID == id {
return a, true, nil
}
}
return assets.Asset{}, false, nil
}
func (c *assetCatalog) GetOwnerByStoragePath(_ context.Context, owner, key string) (assets.Asset, bool, error) {
for _, a := range c.values {
if a.OwnerID == owner && a.StoragePath == key {
return a, true, nil
}
}
return assets.Asset{}, false, nil
}
func (c *assetCatalog) ListPublic(ctx context.Context, owner, client string, _ int) ([]assets.Asset, error) {
all, _ := c.ListOwner(ctx, owner)
var v []assets.Asset
for _, a := range all {
for _, tag := range a.Tags {
if tag == assets.ClientTag(client) {
v = append(v, a)
}
}
}
return v, nil
}
func (c *assetCatalog) GetPublic(ctx context.Context, owner, client, id string, _ int) (assets.Asset, bool, error) {
v, _ := c.ListPublic(ctx, owner, client, 0)
for _, a := range v {
if a.ID == id {
return a, true, nil
}
}
return assets.Asset{}, false, nil
}
func (c *assetCatalog) Create(_ context.Context, a assets.Asset) (assets.Asset, error) {
c.values = append(c.values, a)
return a, nil
}
func (c *assetCatalog) DeleteOwner(_ context.Context, owner, id string) (assets.Asset, bool, error) {
for i, a := range c.values {
if a.OwnerID == owner && a.ID == id {
c.values = append(c.values[:i], c.values[i+1:]...)
return a, true, nil
}
}
return assets.Asset{}, false, nil
}
type assetBlobs struct{ values map[string][]byte }
type signedAssetBlobs struct {
*assetBlobs
url string
}
func (b *signedAssetBlobs) SignReadURL(string, time.Duration) (string, error) { return b.url, nil }
func (b *assetBlobs) Put(_ context.Context, key string, r io.Reader, _ int64, _ string) (assets.StoredObject, error) {
p, _ := io.ReadAll(r)
b.values[key] = p
return assets.StoredObject{Key: key, URL: "https://cdn.test/" + key}, nil
}
func (b *assetBlobs) Read(_ context.Context, key string) (assets.Blob, error) {
p, ok := b.values[key]
if !ok {
return assets.Blob{}, assets.ErrBlobNotFound
}
return assets.Blob{Body: io.NopCloser(bytes.NewReader(p)), ContentType: "image/png", Size: int64(len(p))}, nil
}
func (b *assetBlobs) Delete(_ context.Context, key string) error { delete(b.values, key); return nil }
func TestAssetsPlatformAndPublicHTTP(t *testing.T) {
cat := &assetCatalog{}
blobs := &assetBlobs{values: map[string][]byte{}}
svc := assets.NewService(cat, blobs, nil, func() time.Time { return time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC) }, func(prefix string) string { return prefix + "-1" })
platform, _ := httpapi.NewPlatformAuthorizer(httpapi.AuthState{}, nil)
h, err := httpapi.NewAssetsHandler(svc, platform, publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret"}), httpapi.AssetsConfig{})
if err != nil {
t.Fatal(err)
}
post := request(t, h, http.MethodPost, "/api/assets", strings.NewReader(`{"url":"https://example.test/a.png"}`), map[string]string{"Content-Type": "application/json"})
if post.Code != 201 || !strings.Contains(post.Body.String(), `"name":"外部图片"`) {
t.Fatalf("post=%d %s", post.Code, post.Body.String())
}
list := request(t, h, http.MethodGet, "/api/assets", nil, nil)
if list.Code != 200 || !strings.Contains(list.Body.String(), `"assets"`) {
t.Fatalf("list=%d %s", list.Code, list.Body.String())
}
pubUnauthorized := request(t, h, http.MethodGet, "/api/v1/assets", nil, nil)
if pubUnauthorized.Code != 401 {
t.Fatalf("public unauth=%d %s", pubUnauthorized.Code, pubUnauthorized.Body.String())
}
pub := request(t, h, http.MethodPost, "/api/v1/assets", strings.NewReader(`{"url":"https://example.test/p.png"}`), map[string]string{"Authorization": "Bearer secret", "Content-Type": "application/json"})
if pub.Code != 201 || !strings.Contains(pub.Body.String(), "api-client:agent-a") {
t.Fatalf("public post=%d %s", pub.Code, pub.Body.String())
}
}
func TestAssetsMultipartDownloadServingAndMethods(t *testing.T) {
cat := &assetCatalog{}
blobs := &assetBlobs{values: map[string][]byte{}}
svc := assets.NewService(cat, blobs, nil, time.Now, func(prefix string) string { return prefix + "-x" })
platform, _ := httpapi.NewPlatformAuthorizer(httpapi.AuthState{}, nil)
h, _ := httpapi.NewAssetsHandler(svc, platform, publicapi.NewAuthenticator(publicapi.Config{APIKeys: "a:k"}), httpapi.AssetsConfig{MaxUploadBytes: 1024})
var body bytes.Buffer
mw := multipart.NewWriter(&body)
part, _ := mw.CreateFormFile("files", "a.png")
_, _ = part.Write([]byte("png"))
_ = mw.Close()
upload := request(t, h, http.MethodPost, "/api/assets/upload", &body, map[string]string{"Content-Type": mw.FormDataContentType()})
if upload.Code != 201 {
t.Fatalf("upload=%d %s", upload.Code, upload.Body.String())
}
var payload struct {
Assets []assets.Asset `json:"assets"`
}
_ = json.Unmarshal(upload.Body.Bytes(), &payload)
id := payload.Assets[0].ID
dl := request(t, h, http.MethodGet, "/api/assets/"+id+"/download", nil, nil)
if dl.Code != 200 || dl.Body.String() != "png" || dl.Header().Get("Cache-Control") != "private, no-store" || !strings.Contains(dl.Header().Get("Content-Disposition"), "attachment") {
t.Fatalf("download=%d %#v %q", dl.Code, dl.Header(), dl.Body.String())
}
served := request(t, h, http.MethodGet, "/uploads/"+strings.TrimPrefix(payload.Assets[0].StoragePath, "uploads/"), nil, nil)
if served.Code != 200 || served.Header().Get("Cache-Control") != "public, max-age=31536000, immutable" {
t.Fatalf("served=%d %#v", served.Code, served.Header())
}
head := request(t, h, http.MethodHead, "/api/assets", nil, nil)
if head.Code != 200 || head.Body.Len() != 0 {
t.Fatalf("head=%d %q", head.Code, head.Body.String())
}
options := request(t, h, http.MethodOptions, "/api/assets", nil, nil)
if options.Code != 204 || options.Header().Get("Allow") != "GET, HEAD, POST, OPTIONS" {
t.Fatalf("options=%d allow=%q", options.Code, options.Header().Get("Allow"))
}
bad := request(t, h, http.MethodPatch, "/api/assets", nil, nil)
if bad.Code != 405 || bad.Body.Len() != 0 {
t.Fatalf("bad=%d %q", bad.Code, bad.Body.String())
}
}
func TestAssetInlineReadRedirectsToShortLivedSignedURL(t *testing.T) {
cat := &assetCatalog{values: []assets.Asset{{
ID: "asset-private", OwnerID: "demo-merchant", Name: "private.png", URL: "https://private.test/private.png", StoragePath: "uploads/private.png",
}}}
blobs := &signedAssetBlobs{
assetBlobs: &assetBlobs{values: map[string][]byte{"uploads/private.png": []byte("png")}},
url: "https://private.test/private.png?OSSAccessKeyId=test&Expires=1&Signature=signed",
}
svc := assets.NewService(cat, blobs, nil, time.Now, nil)
platform, _ := httpapi.NewPlatformAuthorizer(httpapi.AuthState{}, nil)
h, _ := httpapi.NewAssetsHandler(svc, platform, publicapi.NewAuthenticator(publicapi.Config{APIKeys: "a:k"}), httpapi.AssetsConfig{})
inline := request(t, h, http.MethodGet, "/api/assets/asset-private/download?inline=1", nil, nil)
if inline.Code != http.StatusTemporaryRedirect || inline.Header().Get("Location") != blobs.url || inline.Header().Get("Cache-Control") != "private, no-store" || inline.Header().Get("Content-Disposition") != "" {
t.Fatalf("inline response = %d %#v", inline.Code, inline.Header())
}
}
func TestAssetsLimitsAndInfrastructureErrorsDoNotLeak(t *testing.T) {
cat := &assetCatalog{}
svc := assets.NewService(cat, &assetBlobs{values: map[string][]byte{}}, nil, time.Now, nil)
platform, _ := httpapi.NewPlatformAuthorizer(httpapi.AuthState{}, nil)
h, _ := httpapi.NewAssetsHandler(svc, platform, publicapi.NewAuthenticator(publicapi.Config{APIKeys: "a:k"}), httpapi.AssetsConfig{MaxJSONBytes: 8, MaxUploadBytes: 8})
r := request(t, h, http.MethodPost, "/api/assets", strings.NewReader(`{"url":"https://secret.example.test"}`), map[string]string{"Content-Type": "application/json"})
if r.Code != 413 || strings.Contains(r.Body.String(), "secret") {
t.Fatalf("limit=%d %s", r.Code, r.Body.String())
}
missing := request(t, h, http.MethodGet, "/api/assets/missing/download", nil, nil)
if missing.Code != 404 || missing.Body.String() != "{\"error\":\"资产不存在\"}\n" {
t.Fatalf("missing=%d %q", missing.Code, missing.Body.String())
}
}
func request(t *testing.T, h http.Handler, method, path string, body io.Reader, headers map[string]string) *httptest.ResponseRecorder {
t.Helper()
r := httptest.NewRequest(method, path, body)
for k, v := range headers {
r.Header.Set(k, v)
}
w := httptest.NewRecorder()
h.ServeHTTP(w, r)
return w
}