178 lines
7.7 KiB
Go
178 lines
7.7 KiB
Go
package httpapi_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"reflect"
|
|
"strconv"
|
|
"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 formatDownloadBlobs struct {
|
|
*assetBlobs
|
|
contentTypes map[string]string
|
|
reads int
|
|
}
|
|
|
|
type failingImageBody struct {
|
|
prefix []byte
|
|
closed bool
|
|
}
|
|
|
|
func (b *failingImageBody) Read(p []byte) (int, error) {
|
|
if len(b.prefix) > 0 {
|
|
n := copy(p, b.prefix)
|
|
b.prefix = b.prefix[n:]
|
|
return n, nil
|
|
}
|
|
return 0, errors.New("storage stream failed")
|
|
}
|
|
|
|
func (b *failingImageBody) Close() error {
|
|
b.closed = true
|
|
return nil
|
|
}
|
|
|
|
type failingDownloadBlobs struct {
|
|
*assetBlobs
|
|
body *failingImageBody
|
|
}
|
|
|
|
func (b *failingDownloadBlobs) Read(context.Context, string) (assets.Blob, error) {
|
|
return assets.Blob{Body: b.body, ContentType: "application/octet-stream", Size: 1000}, nil
|
|
}
|
|
|
|
func (b *formatDownloadBlobs) Read(ctx context.Context, key string) (assets.Blob, error) {
|
|
blob, err := b.assetBlobs.Read(ctx, key)
|
|
if err == nil {
|
|
b.reads++
|
|
blob.ContentType = b.contentTypes[key]
|
|
}
|
|
return blob, err
|
|
}
|
|
|
|
func TestLegacyImageDownloadUsesBytesForResponseNameAndType(t *testing.T) {
|
|
tests := []struct {
|
|
label, id, name, mediaType, suffix string
|
|
data []byte
|
|
}{
|
|
{"png", "legacy-png", "生成图片.image", "image/png", ".png", append([]byte("\x89PNG\r\n\x1a\n"), bytes.Repeat([]byte("p"), 650)...)},
|
|
{"jpeg", "legacy-jpeg", "生成图片.image", "image/jpeg", ".jpg", append([]byte("\xff\xd8\xff\xe0"), bytes.Repeat([]byte("j"), 650)...)},
|
|
{"webp", "legacy-webp", "生成图片.image", "image/webp", ".webp", append([]byte("RIFF\x12\x00\x00\x00WEBPVP8 "), bytes.Repeat([]byte("w"), 650)...)},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.label, func(t *testing.T) {
|
|
key := "uploads/" + tc.id + ".image"
|
|
original := assets.Asset{ID: tc.id, OwnerID: "demo-merchant", Name: tc.name, Kind: assets.KindImage, StoragePath: key, URL: "https://cdn.test/" + key}
|
|
catalog := &assetCatalog{values: []assets.Asset{original}}
|
|
blobs := &formatDownloadBlobs{
|
|
assetBlobs: &assetBlobs{values: map[string][]byte{key: bytes.Clone(tc.data)}},
|
|
contentTypes: map[string]string{key: "application/octet-stream"},
|
|
}
|
|
h := newImageDownloadHandler(t, catalog, blobs)
|
|
for _, inline := range []bool{false, true} {
|
|
path := "/api/assets/" + tc.id + "/download"
|
|
kind := "attachment"
|
|
if inline {
|
|
path += "?inline=1"
|
|
kind = "inline"
|
|
}
|
|
response := request(t, h, http.MethodGet, path, nil, nil)
|
|
wantName := strings.TrimSuffix(tc.name, ".image") + tc.suffix
|
|
if response.Code != http.StatusOK || !bytes.Equal(response.Body.Bytes(), tc.data) {
|
|
t.Fatalf("inline=%t status=%d body=%q", inline, response.Code, response.Body.Bytes())
|
|
}
|
|
if got := response.Header().Get("Content-Type"); got != tc.mediaType {
|
|
t.Errorf("Content-Type=%q, want %q", got, tc.mediaType)
|
|
}
|
|
if got := response.Header().Get("Content-Length"); got != strconv.Itoa(len(tc.data)) {
|
|
t.Errorf("Content-Length=%q, want %d", got, len(tc.data))
|
|
}
|
|
if got := response.Header().Get("Content-Disposition"); !strings.HasPrefix(got, kind+"; ") || !strings.Contains(got, "filename*=UTF-8''"+url.PathEscape(wantName)) {
|
|
t.Errorf("Content-Disposition=%q, want %s and UTF-8 filename %q", got, kind, wantName)
|
|
}
|
|
}
|
|
if !reflect.DeepEqual(catalog.values, []assets.Asset{original}) || !bytes.Equal(blobs.values[key], tc.data) || len(blobs.values) != 1 {
|
|
t.Fatal("download changed catalog metadata or stored bytes")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPublicLegacyImageDownloadRequiresClientReadAccess(t *testing.T) {
|
|
data := append([]byte("\x89PNG\r\n\x1a\n"), bytes.Repeat([]byte("x"), 600)...)
|
|
key := "uploads/public.image"
|
|
asset := assets.Asset{ID: "public-image", OwnerID: publicapi.OwnerID("agent-a"), Name: "公开图.image", Kind: assets.KindImage, StoragePath: key, Tags: []string{assets.ClientTag("agent-a")}}
|
|
catalog := &assetCatalog{values: []assets.Asset{asset}}
|
|
blobs := &formatDownloadBlobs{assetBlobs: &assetBlobs{values: map[string][]byte{key: data}}, contentTypes: map[string]string{key: "application/octet-stream"}}
|
|
h := newImageDownloadHandler(t, catalog, blobs)
|
|
path := "/api/v1/assets/public-image/download"
|
|
|
|
for _, headers := range []map[string]string{nil, {"Authorization": "Bearer key-b"}} {
|
|
response := request(t, h, http.MethodGet, path, nil, headers)
|
|
if response.Code != http.StatusUnauthorized && response.Code != http.StatusNotFound {
|
|
t.Fatalf("unreadable public download status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
}
|
|
if blobs.reads != 0 {
|
|
t.Fatalf("blob was read without access: %d reads", blobs.reads)
|
|
}
|
|
response := request(t, h, http.MethodGet, path, nil, map[string]string{"Authorization": "Bearer key-a"})
|
|
if response.Code != http.StatusOK || !bytes.Equal(response.Body.Bytes(), data) || response.Header().Get("Content-Type") != "image/png" || !strings.Contains(response.Header().Get("Content-Disposition"), "filename*=UTF-8''"+url.PathEscape("公开图.png")) {
|
|
t.Fatalf("public download status=%d headers=%v body=%q", response.Code, response.Header(), response.Body.Bytes())
|
|
}
|
|
if blobs.reads != 1 {
|
|
t.Fatalf("authorized blob reads=%d, want 1", blobs.reads)
|
|
}
|
|
}
|
|
|
|
func TestVideoDownloadKeepsOriginalNameAndType(t *testing.T) {
|
|
key := "uploads/video.mp4"
|
|
data := []byte("\x00\x00\x00\x18ftypmp42-video-content")
|
|
catalog := &assetCatalog{values: []assets.Asset{{ID: "video", OwnerID: "demo-merchant", Kind: assets.KindVideo, Name: "视频.image", StoragePath: key}}}
|
|
blobs := &formatDownloadBlobs{assetBlobs: &assetBlobs{values: map[string][]byte{key: data}}, contentTypes: map[string]string{key: "video/mp4"}}
|
|
h := newImageDownloadHandler(t, catalog, blobs)
|
|
response := request(t, h, http.MethodGet, "/api/assets/video/download", nil, nil)
|
|
if response.Code != http.StatusOK || !bytes.Equal(response.Body.Bytes(), data) || response.Header().Get("Content-Type") != "video/mp4" || !strings.Contains(response.Header().Get("Content-Disposition"), "filename*=UTF-8''"+url.PathEscape("视频.image")) {
|
|
t.Fatalf("video download status=%d headers=%v body=%q", response.Code, response.Header(), response.Body.Bytes())
|
|
}
|
|
}
|
|
|
|
func TestImageDownloadPrefixFailureDoesNotSendSuccessHeaders(t *testing.T) {
|
|
key := "uploads/failing.image"
|
|
catalog := &assetCatalog{values: []assets.Asset{{ID: "failing", OwnerID: "demo-merchant", Kind: assets.KindImage, Name: "失败.image", StoragePath: key}}}
|
|
body := &failingImageBody{prefix: []byte("\x89PNG\r\n\x1a\n")}
|
|
blobs := &failingDownloadBlobs{assetBlobs: &assetBlobs{values: map[string][]byte{}}, body: body}
|
|
h := newImageDownloadHandler(t, catalog, blobs)
|
|
response := request(t, h, http.MethodGet, "/api/assets/failing/download", nil, nil)
|
|
if response.Code != http.StatusInternalServerError || response.Header().Get("Content-Disposition") != "" || response.Header().Get("Content-Type") == "image/png" || !body.closed {
|
|
t.Fatalf("failed download status=%d headers=%v closed=%t body=%q", response.Code, response.Header(), body.closed, response.Body.String())
|
|
}
|
|
}
|
|
|
|
func newImageDownloadHandler(t *testing.T, catalog *assetCatalog, blobs assets.BlobStore) http.Handler {
|
|
t.Helper()
|
|
service := assets.NewService(catalog, blobs, nil, time.Now, nil)
|
|
platform, err := httpapi.NewPlatformAuthorizer(httpapi.AuthState{}, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
h, err := httpapi.NewAssetsHandler(service, platform, publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:key-a,agent-b:key-b"}), httpapi.AssetsConfig{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return h
|
|
}
|
|
|
|
var _ io.ReadCloser = (*failingImageBody)(nil)
|