mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-01 06:36:38 +08:00
perf: refactor
This commit is contained in:
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user