package upload import ( "bytes" "context" "encoding/json" "io" "mime/multipart" "net/http" "net/http/httptest" "strings" "testing" "hyapp-admin-server/internal/model" "hyapp-admin-server/internal/modules/shared" "github.com/gin-gonic/gin" ) type fakeUploader struct { contentType string key string keys []string sizeBytes int64 body []byte } func (u *fakeUploader) PutObject(_ context.Context, key string, reader io.Reader, sizeBytes int64, contentType string) (string, error) { body, err := io.ReadAll(reader) if err != nil { return "", err } u.key = key u.keys = append(u.keys, key) u.body = body u.sizeBytes = sizeBytes u.contentType = contentType return "https://media.example.com/" + key, nil } type fakeAudit struct{} func (fakeAudit) CreateOperationLog(model.OperationLog) error { return nil } func (fakeAudit) LogOperation(shared.OperationEvent) error { return nil } func TestUploadImageUsesAdminPrefix(t *testing.T) { uploader := &fakeUploader{} handler := New(uploader, "admin", fakeAudit{}) router := ginRouter(handler) payload := append([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, bytes.Repeat([]byte{1}, 32)...) body, contentType := multipartUploadBody(t, "file", "avatar.png", "image/png", payload) req := httptest.NewRequest(http.MethodPost, "/upload-image", body) req.Header.Set("Content-Type", contentType) rec := httptest.NewRecorder() router.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d body=%s", rec.Code, rec.Body.String()) } if !strings.HasPrefix(uploader.key, "admin/images/7/") || !strings.HasSuffix(uploader.key, ".png") { t.Fatalf("unexpected key: %s", uploader.key) } if !bytes.Equal(uploader.body, payload) || uploader.sizeBytes != int64(len(payload)) || uploader.contentType != "image/png" { t.Fatalf("unexpected upload content: size=%d type=%s", uploader.sizeBytes, uploader.contentType) } var envelope struct { Code int `json:"code"` Data struct { ObjectKey string `json:"object_key"` URL string `json:"url"` } `json:"data"` } if err := json.NewDecoder(rec.Body).Decode(&envelope); err != nil { t.Fatalf("decode response: %v", err) } if envelope.Data.ObjectKey != uploader.key || envelope.Data.URL != "https://media.example.com/"+uploader.key { t.Fatalf("unexpected response: %+v", envelope.Data) } } func TestUploadImageRejectsSpoofedContent(t *testing.T) { uploader := &fakeUploader{} handler := New(uploader, "admin", fakeAudit{}) router := ginRouter(handler) body, contentType := multipartUploadBody(t, "file", "avatar.png", "image/png", []byte("not an image")) req := httptest.NewRequest(http.MethodPost, "/upload-image", body) req.Header.Set("Content-Type", contentType) rec := httptest.NewRecorder() router.ServeHTTP(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("status = %d body=%s", rec.Code, rec.Body.String()) } if uploader.key != "" { t.Fatalf("invalid image must not upload: %s", uploader.key) } } func TestUploadImageAllowsAnimatedAssetFormats(t *testing.T) { uploader := &fakeUploader{} handler := New(uploader, "admin", fakeAudit{}) router := ginRouter(handler) body, contentType := multipartUploadBody(t, "file", "cover.pag", "application/octet-stream", []byte("PAG animated asset")) req := httptest.NewRequest(http.MethodPost, "/upload-image", body) req.Header.Set("Content-Type", contentType) rec := httptest.NewRecorder() router.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d body=%s", rec.Code, rec.Body.String()) } if !strings.HasPrefix(uploader.key, "admin/images/7/") || !strings.HasSuffix(uploader.key, ".pag") { t.Fatalf("unexpected key: %s", uploader.key) } if uploader.contentType != "application/vnd.tencent.pag" { t.Fatalf("unexpected content type: %s", uploader.contentType) } } func TestUploadImageAllowsPayloadOverPreviousImageLimit(t *testing.T) { uploader := &fakeUploader{} handler := New(uploader, "admin", fakeAudit{}) router := ginRouter(handler) payload := append([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, bytes.Repeat([]byte{1}, (5<<20)+1)...) body, contentType := multipartUploadBody(t, "file", "large.png", "image/png", payload) req := httptest.NewRequest(http.MethodPost, "/upload-image", body) req.Header.Set("Content-Type", contentType) rec := httptest.NewRecorder() router.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d body=%s", rec.Code, rec.Body.String()) } if uploader.sizeBytes != int64(len(payload)) { t.Fatalf("unexpected upload size: %d", uploader.sizeBytes) } } func TestUploadImagesBatchUploadsSeriallyAndLimitsCount(t *testing.T) { uploader := &fakeUploader{} handler := New(uploader, "admin", fakeAudit{}) router := ginRouter(handler) body, contentType := multipartBatchUploadBody(t, []string{"one.png", "two.pag"}) req := httptest.NewRequest(http.MethodPost, "/upload-images-batch", body) req.Header.Set("Content-Type", contentType) rec := httptest.NewRecorder() router.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("status = %d body=%s", rec.Code, rec.Body.String()) } if len(uploader.keys) != 2 { t.Fatalf("expected 2 uploads, got %d", len(uploader.keys)) } var envelope struct { Code int `json:"code"` Data []struct { Filename string `json:"filename"` URL string `json:"url"` } `json:"data"` } if err := json.NewDecoder(rec.Body).Decode(&envelope); err != nil { t.Fatalf("decode response: %v", err) } if len(envelope.Data) != 2 || envelope.Data[0].Filename != "one.png" || envelope.Data[1].Filename != "two.pag" { t.Fatalf("unexpected response order: %+v", envelope.Data) } tooMany := make([]string, 11) for index := range tooMany { tooMany[index] = "asset.png" } body, contentType = multipartBatchUploadBody(t, tooMany) req = httptest.NewRequest(http.MethodPost, "/upload-images-batch", body) req.Header.Set("Content-Type", contentType) rec = httptest.NewRecorder() router.ServeHTTP(rec, req) if rec.Code != http.StatusBadRequest { t.Fatalf("too many status = %d body=%s", rec.Code, rec.Body.String()) } } func ginRouter(handler *Handler) http.Handler { router := gin.New() router.Use(func(c *gin.Context) { c.Set("userID", uint(7)) c.Set("username", "admin") c.Next() }) router.POST("/upload-image", handler.UploadImage) router.POST("/upload-images-batch", handler.UploadImagesBatch) return router } func multipartUploadBody(t *testing.T, fieldName string, filename string, contentType string, payload []byte) (*bytes.Buffer, string) { t.Helper() body := &bytes.Buffer{} writer := multipart.NewWriter(body) part, err := writer.CreatePart(map[string][]string{ "Content-Disposition": {`form-data; name="` + fieldName + `"; filename="` + filename + `"`}, "Content-Type": {contentType}, }) if err != nil { t.Fatalf("create multipart part: %v", err) } if _, err := part.Write(payload); err != nil { t.Fatalf("write multipart part: %v", err) } if err := writer.Close(); err != nil { t.Fatalf("close multipart writer: %v", err) } return body, writer.FormDataContentType() } func multipartBatchUploadBody(t *testing.T, filenames []string) (*bytes.Buffer, string) { t.Helper() body := &bytes.Buffer{} writer := multipart.NewWriter(body) for _, filename := range filenames { contentType := "image/png" payload := append([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, bytes.Repeat([]byte{1}, 8)...) if strings.HasSuffix(filename, ".pag") { contentType = "application/octet-stream" payload = []byte("PAG animated asset") } part, err := writer.CreatePart(map[string][]string{ "Content-Disposition": {`form-data; name="files"; filename="` + filename + `"`}, "Content-Type": {contentType}, }) if err != nil { t.Fatalf("create multipart part: %v", err) } if _, err := part.Write(payload); err != nil { t.Fatalf("write multipart part: %v", err) } } if err := writer.Close(); err != nil { t.Fatalf("close multipart writer: %v", err) } return body, writer.FormDataContentType() }