perf: refactor

This commit is contained in:
ryan
2026-06-20 09:46:24 +08:00
parent 080be1e03a
commit cac6e88bc0
8 changed files with 244 additions and 120 deletions
+6 -6
View File
@@ -25,7 +25,7 @@ import (
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/diskcache"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/pkg/logger"
"github.com/gin-gonic/gin"
"golang.org/x/sync/singleflight"
@@ -66,14 +66,14 @@ func ServeFileByID(c *gin.Context) {
upload, err := GetUploadRecordByID(c)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if _, ok := err.(*strconv.NumError); ok {
c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid upload ID"})
response.AbortBadRequest(c, "无效的上传ID")
return
}
c.AbortWithStatus(http.StatusInternalServerError)
response.AbortInternal(c, "服务器内部错误")
return
}
@@ -263,7 +263,7 @@ func ImageCompressionCacheKey(upload *model.Upload, quality string) string {
func serveOriginal(c *gin.Context, upload *model.Upload) {
obj, err := uploadstorage.OpenStoredObject(c.Request.Context(), upload)
if err != nil {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件未找到")
return
}
defer func() { _ = obj.Body.Close() }()
@@ -313,4 +313,4 @@ func CheckFileAccessPermission(c *gin.Context, upload *model.Upload) error {
}
}
return nil
}
}
@@ -6,6 +6,7 @@ package filesrv
import (
"bytes"
"context"
"encoding/json"
"image"
"image/color"
@@ -13,6 +14,7 @@ import (
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/upload/cache"
@@ -20,12 +22,16 @@ import (
"github.com/Rain-kl/Wavelet/internal/apps/upload/util"
"github.com/Rain-kl/Wavelet/internal/common"
"github.com/Rain-kl/Wavelet/internal/common/response"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/diskcache"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/storage"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-contrib/sessions"
"github.com/gin-contrib/sessions/cookie"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
func TestServeFileByIDAccessControl(t *testing.T) {
@@ -33,8 +39,8 @@ func TestServeFileByIDAccessControl(t *testing.T) {
defer cleanup()
cache.ResetAccessCaches()
// Ensure uploads dir is cleaned up
defer func() { _ = os.RemoveAll("uploads") }()
tempDir := t.TempDir()
configureLocalStorageRoot(t, dbConn, tempDir)
// Create a user in DB
user := model.User{
@@ -60,33 +66,36 @@ func TestServeFileByIDAccessControl(t *testing.T) {
// Create two files: one in whitelist (avatar), one not in whitelist (attachment)
avatarFile := model.Upload{
ID: 8001,
UserID: user.ID,
FileName: "avatar.png",
FilePath: "uploads/avatar.png",
FileSize: 5,
MimeType: "image/png",
Extension: "png",
Type: "avatar",
Status: model.UploadStatusUsed,
AccessMode: 1,
ID: 8001,
UserID: user.ID,
FileName: "avatar.png",
FilePath: "avatar.png",
FileSize: 5,
MimeType: "image/png",
Extension: "png",
Type: "avatar",
Status: model.UploadStatusUsed,
AccessMode: 1,
}
attachmentFile := model.Upload{
ID: 8002,
UserID: user.ID,
FileName: "doc.pdf",
FilePath: "uploads/doc.pdf",
FileSize: 5,
MimeType: "application/pdf",
Extension: "pdf",
Type: "attachment",
Status: model.UploadStatusUsed,
AccessMode: 1,
ID: 8002,
UserID: user.ID,
FileName: "doc.pdf",
FilePath: "doc.pdf",
FileSize: 5,
MimeType: "application/pdf",
Extension: "pdf",
Type: "attachment",
Status: model.UploadStatusUsed,
AccessMode: 1,
}
_ = os.MkdirAll("uploads", 0755)
_ = os.WriteFile(avatarFile.FilePath, []byte("image"), 0644)
_ = os.WriteFile(attachmentFile.FilePath, []byte("bytes"), 0644)
if err := os.WriteFile(filepath.Join(tempDir, "avatar.png"), []byte("image"), 0644); err != nil {
t.Fatalf("failed to write avatar file: %v", err)
}
if err := os.WriteFile(filepath.Join(tempDir, "doc.pdf"), []byte("bytes"), 0644); err != nil {
t.Fatalf("failed to write attachment file: %v", err)
}
dbConn.Create(&avatarFile)
dbConn.Create(&attachmentFile)
@@ -159,15 +168,14 @@ func TestImageCompression(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
tempDir := t.TempDir()
configureLocalStorageRoot(t, dbConn, tempDir)
cache := diskcache.GetGlobalCache()
if err := cache.Clear(); err != nil {
t.Fatalf("failed to clear disk cache before test: %v", err)
}
// Ensure uploads dir is cleaned up
defer func() {
_ = os.RemoveAll("uploads")
}()
defer func() {
if err := cache.Clear(); err != nil {
t.Errorf("failed to clear disk cache after test: %v", err)
@@ -190,24 +198,23 @@ func TestImageCompression(t *testing.T) {
t.Fatalf("failed to encode test png: %v", err)
}
_ = os.MkdirAll("uploads", 0755)
filePath := "uploads/test_image.png"
filePath := filepath.Join(tempDir, "test_image.png")
if err := os.WriteFile(filePath, pngBuf.Bytes(), 0644); err != nil {
t.Fatalf("failed to write test png: %v", err)
}
// Save upload record to DB
uploadRecord := model.Upload{
ID: 3001,
UserID: user.ID,
FileName: "test_image.png",
FilePath: filePath,
FileSize: int64(pngBuf.Len()),
MimeType: "image/png",
Extension: "png",
Type: "avatar", // Whitelisted by default
Status: model.UploadStatusUsed,
AccessMode: 1,
ID: 3001,
UserID: user.ID,
FileName: "test_image.png",
FilePath: "test_image.png",
FileSize: int64(pngBuf.Len()),
MimeType: "image/png",
Extension: "png",
Type: "avatar", // Whitelisted by default
Status: model.UploadStatusUsed,
AccessMode: 1,
}
dbConn.Create(&uploadRecord)
@@ -344,3 +351,26 @@ func TestNormalizeImageQuality(t *testing.T) {
})
}
}
func configureLocalStorageRoot(t *testing.T, dbConn *gorm.DB, tempDir string) {
var sc model.SystemConfig
if err := dbConn.Where("key = ?", model.ConfigKeyStorageConfig).First(&sc).Error; err != nil {
t.Fatalf("failed to find storage config: %v", err)
}
var cfg storage.Config
if err := json.Unmarshal([]byte(sc.Value), &cfg); err != nil {
t.Fatalf("failed to unmarshal storage config: %v", err)
}
cfg.Local.Root = tempDir
newVal, err := json.Marshal(cfg)
if err != nil {
t.Fatalf("failed to marshal storage config: %v", err)
}
sc.Value = string(newVal)
if err := dbConn.Save(&sc).Error; err != nil {
t.Fatalf("failed to save storage config: %v", err)
}
_ = db.HSetJSON(context.Background(), repository.SystemConfigRedisHashKey, sc.Key, &sc)
repository.ResetSystemConfigRAMCacheForTest()
storage.ResetCache()
}
@@ -112,7 +112,7 @@ func DeleteFile(c *gin.Context) {
if _, err := softDeleteUpload(ctx, uploadID); err != nil {
if isRecordNotFound(err) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
response.AbortBadRequest(c, shared.ErrDeleteFileFailed)
@@ -233,11 +233,11 @@ func DeleteMyFile(c *gin.Context) {
if _, err := softDeleteOwnedUpload(ctx, currUser.ID, uploadID); err != nil {
if isRecordNotFound(err) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if err == ingest.ErrForbidden {
c.AbortWithStatus(http.StatusForbidden)
response.AbortForbidden(c, "无权操作")
return
}
response.AbortBadRequest(c, shared.ErrDeleteFileFailed)
@@ -287,11 +287,11 @@ func UpdateMyFile(c *gin.Context) {
upload, err := updateOwnedUpload(ctx, currUser.ID, uploadID, updateMyUploadInput(req))
if err != nil {
if isRecordNotFound(err) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if err == ingest.ErrForbidden {
c.AbortWithStatus(http.StatusForbidden)
response.AbortForbidden(c, "无权操作")
return
}
response.AbortBadRequest(c, "更新文件记录失败")
@@ -299,4 +299,4 @@ func UpdateMyFile(c *gin.Context) {
}
c.JSON(http.StatusOK, response.OK(upload))
}
}
+1 -1
View File
@@ -170,7 +170,7 @@ func DownloadFile(c *gin.Context) {
upload, err := filesrv.GetUploadRecordByID(c)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
c.AbortWithStatus(http.StatusNotFound)
response.AbortNotFound(c, "文件记录未找到")
return
}
if _, ok := err.(*strconv.NumError); ok {