hyapp-server/server/admin/internal/modules/upload/handler_test.go

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
	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.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 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)
	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()
}