Feat: user storage pack

models/migration.go

@@ -29,7 +29,7 @@ func migration() {
 	if conf.DatabaseConfig.Type == "mysql" {
 		DB = DB.Set("gorm:table_options", "ENGINE=InnoDB")
 	}
-	DB.AutoMigrate(&User{}, &Setting{}, &Group{}, &Policy{}, &Folder{}, &File{})
+	DB.AutoMigrate(&User{}, &Setting{}, &Group{}, &Policy{}, &Folder{}, &File{}, &StoragePack{})
 
 	// 创建初始存储策略
 	addDefaultPolicy()

models/storage_pack.go

@@ -0,0 +1,61 @@
+package model
+
+import (
+	"github.com/HFO4/cloudreve/pkg/cache"
+	"github.com/jinzhu/gorm"
+	"strconv"
+	"time"
+)
+
+// StoragePack 容量包模型
+type StoragePack struct {
+	// 表字段
+	gorm.Model
+	Name        string
+	UserID      uint
+	ActiveTime  *time.Time
+	ExpiredTime *time.Time `gorm:"index:expired"`
+	Size        uint64
+}
+
+// GetAvailablePackSize 返回给定用户当前可用的容量包总容量
+func (user *User) GetAvailablePackSize() uint64 {
+	var (
+		packs       []StoragePack
+		total       uint64
+		firstExpire *time.Time
+		timeNow     = time.Now()
+		ttl         int64
+	)
+
+	// 尝试从缓存中读取
+	cacheKey := "pack_size_" + strconv.FormatUint(uint64(user.ID), 10)
+	if total, ok := cache.Get(cacheKey); ok {
+		return total.(uint64)
+	}
+
+	// 查找所有有效容量包
+	DB.Where("expired_time > ? AND user_id = ?", timeNow, user.ID).Find(&packs)
+
+	// 计算总容量, 并找到其中最早的过期时间
+	for _, v := range packs {
+		total += v.Size
+		if firstExpire == nil {
+			firstExpire = v.ExpiredTime
+			continue
+		}
+		if v.ExpiredTime != nil && firstExpire.After(*v.ExpiredTime) {
+			firstExpire = v.ExpiredTime
+		}
+	}
+
+	// 用最早的过期时间计算缓存TTL，并写入缓存
+	if firstExpire != nil {
+		ttl = firstExpire.Unix() - timeNow.Unix()
+		if ttl > 0 {
+			_ = cache.Set(cacheKey, total, int(ttl))
+		}
+	}
+
+	return total
+}

models/storage_pack_test.go

@@ -0,0 +1,39 @@
+package model
+
+import (
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+	"testing"
+	"time"
+)
+
+func TestUser_GetAvailablePackSize(t *testing.T) {
+	asserts := assert.New(t)
+	user := User{
+		Model: gorm.Model{ID: 1},
+	}
+
+	// 未命中缓存
+	{
+		mock.ExpectQuery("SELECT(.+)").
+			WithArgs(sqlmock.AnyArg(), 1).
+			WillReturnRows(
+				sqlmock.NewRows([]string{"id", "size", "expired_time"}).
+					AddRow(1, 10, time.Now().Add(time.Second*time.Duration(20))).
+					AddRow(3, 0, nil).
+					AddRow(2, 20, time.Now().Add(time.Second*time.Duration(10))),
+			)
+		total := user.GetAvailablePackSize()
+		asserts.NoError(mock.ExpectationsWereMet())
+		asserts.EqualValues(30, total)
+	}
+
+	// 命中缓存
+	{
+		total := user.GetAvailablePackSize()
+		asserts.NoError(mock.ExpectationsWereMet())
+		asserts.EqualValues(30, total)
+	}
+
+}

models/user.go

@@ -105,10 +105,11 @@ func (user *User) IncreaseStorageWithoutCheck(size uint64) {
 
 // GetRemainingCapacity 获取剩余配额
 func (user *User) GetRemainingCapacity() uint64 {
-	if user.Group.MaxStorage <= user.Storage {
+	total := user.Group.MaxStorage + user.GetAvailablePackSize()
+	if total <= user.Storage {
 		return 0
 	}
-	return user.Group.MaxStorage - user.Storage
+	return total - user.Storage
 }
 
 // GetPolicyID 获取用户当前的存储策略ID

models/user_test.go

@@ -3,6 +3,7 @@ package model
 import (
 	"encoding/json"
 	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/jinzhu/gorm"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
@@ -171,6 +172,7 @@ func TestUser_GetPolicyID(t *testing.T) {
 func TestUser_GetRemainingCapacity(t *testing.T) {
 	asserts := assert.New(t)
 	newUser := NewUser()
+	cache.Set("pack_size_0", uint64(0), 0)
 
 	newUser.Group.MaxStorage = 100
 	asserts.Equal(uint64(100), newUser.GetRemainingCapacity())
@@ -186,6 +188,11 @@ func TestUser_GetRemainingCapacity(t *testing.T) {
 	newUser.Group.MaxStorage = 100
 	newUser.Storage = 200
 	asserts.Equal(uint64(0), newUser.GetRemainingCapacity())
+
+	cache.Set("pack_size_0", uint64(10), 0)
+	newUser.Group.MaxStorage = 100
+	newUser.Storage = 101
+	asserts.Equal(uint64(9), newUser.GetRemainingCapacity())
 }
 
 func TestUser_DeductionCapacity(t *testing.T) {
@@ -204,6 +211,7 @@ func TestUser_DeductionCapacity(t *testing.T) {
 
 	newUser, err := GetUserByID(1)
 	newUser.Group.MaxStorage = 100
+	cache.Set("pack_size_1", uint64(0), 0)
 	asserts.NoError(err)
 	asserts.NoError(mock.ExpectationsWereMet())
 
@@ -219,6 +227,10 @@ func TestUser_DeductionCapacity(t *testing.T) {
 	asserts.Equal(false, newUser.IncreaseStorage(1))
 	asserts.Equal(uint64(100), newUser.Storage)
 
+	cache.Set("pack_size_1", uint64(1), 0)
+	asserts.Equal(true, newUser.IncreaseStorage(1))
+	asserts.Equal(uint64(101), newUser.Storage)
+
 	asserts.True(newUser.IncreaseStorage(0))
 }
 

pkg/filesystem/file.go

@@ -99,7 +99,7 @@ func (fs *FileSystem) Preview(ctx context.Context, path string, isText bool) (*r
 
 	// 如果是文本文件预览，需要检查大小限制
 	sizeLimit := model.GetIntSetting("maxEditSize", 2<<20)
-	if fs.FileTarget[0].Size > uint64(sizeLimit) {
+	if isText && fs.FileTarget[0].Size > uint64(sizeLimit) {
 		return nil, ErrFileSizeTooBig
 	}
 

pkg/filesystem/hooks_test.go

@@ -28,6 +28,7 @@ func TestGenericBeforeUpload(t *testing.T) {
 		Size: 5,
 		Name: "1.txt",
 	}
+	cache.Set("pack_size_0", uint64(0), 0)
 	ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, file)
 	fs := FileSystem{
 		User: &model.User{
@@ -266,6 +267,7 @@ func TestHookIsFileExist(t *testing.T) {
 
 func TestHookValidateCapacity(t *testing.T) {
 	asserts := assert.New(t)
+	cache.Set("pack_size_1", uint64(0), 0)
 	fs := &FileSystem{User: &model.User{
 		Model:   gorm.Model{ID: 1},
 		Storage: 0,
@@ -313,6 +315,7 @@ func TestHookResetPolicy(t *testing.T) {
 
 func TestHookChangeCapacity(t *testing.T) {
 	asserts := assert.New(t)
+	cache.Set("pack_size_1", uint64(0), 0)
 
 	// 容量增加 失败
 	{

pkg/filesystem/manage_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"github.com/DATA-DOG/go-sqlmock"
 	model "github.com/HFO4/cloudreve/models"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/HFO4/cloudreve/pkg/conf"
 	"github.com/HFO4/cloudreve/pkg/serializer"
 	"github.com/jinzhu/gorm"
@@ -275,6 +276,7 @@ func TestFileSystem_ListDeleteDirs(t *testing.T) {
 func TestFileSystem_Delete(t *testing.T) {
 	conf.DatabaseConfig.Type = "mysql"
 	asserts := assert.New(t)
+	cache.Set("pack_size_1", uint64(0), 0)
 	fs := &FileSystem{User: &model.User{
 		Model: gorm.Model{
 			ID: 1,
@@ -381,6 +383,7 @@ func TestFileSystem_Delete(t *testing.T) {
 
 func TestFileSystem_Copy(t *testing.T) {
 	asserts := assert.New(t)
+	cache.Set("pack_size_1", uint64(0), 0)
 	fs := &FileSystem{User: &model.User{
 		Model: gorm.Model{
 			ID: 1,
@@ -431,6 +434,7 @@ func TestFileSystem_Copy(t *testing.T) {
 
 func TestFileSystem_Move(t *testing.T) {
 	asserts := assert.New(t)
+	cache.Set("pack_size_1", uint64(0), 0)
 	fs := &FileSystem{User: &model.User{
 		Model: gorm.Model{
 			ID: 1,

pkg/filesystem/validator_test.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"github.com/DATA-DOG/go-sqlmock"
 	model "github.com/HFO4/cloudreve/models"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/assert"
 	"testing"
@@ -42,6 +43,7 @@ func TestFileSystem_ValidateLegalName(t *testing.T) {
 func TestFileSystem_ValidateCapacity(t *testing.T) {
 	asserts := assert.New(t)
 	ctx := context.Background()
+	cache.Set("pack_size_0", uint64(0), 0)
 	fs := FileSystem{
 		User: &model.User{
 			Storage: 10,
@@ -57,6 +59,11 @@ func TestFileSystem_ValidateCapacity(t *testing.T) {
 	fs.User.Storage = 5
 	asserts.False(fs.ValidateCapacity(ctx, 10))
 	asserts.Equal(uint64(5), fs.User.Storage)
+
+	fs.User.Storage = 5
+	cache.Set("pack_size_0", uint64(15), 0)
+	asserts.True(fs.ValidateCapacity(ctx, 10))
+	asserts.Equal(uint64(15), fs.User.Storage)
 }
 
 func TestFileSystem_ValidateFileSize(t *testing.T) {

pkg/serializer/user.go

@@ -87,13 +87,14 @@ func BuildUserResponse(user model.User) Response {
 
 // BuildUserStorageResponse 序列化用户存储概况响应
 func BuildUserStorageResponse(user model.User) Response {
+	total := user.Group.MaxStorage + user.GetAvailablePackSize()
 	storageResp := storage{
 		Used:  user.Storage,
-		Free:  user.Group.MaxStorage - user.Storage,
-		Total: user.Group.MaxStorage,
+		Free:  total - user.Storage,
+		Total: total,
 	}
 
-	if user.Group.MaxStorage < user.Storage {
+	if total < user.Storage {
 		storageResp.Free = 0
 	}
 

pkg/serializer/user_test.go

@@ -2,6 +2,7 @@ package serializer
 
 import (
 	model "github.com/HFO4/cloudreve/models"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/stretchr/testify/assert"
 	"testing"
 )
@@ -27,6 +28,7 @@ func TestBuildUserResponse(t *testing.T) {
 
 func TestBuildUserStorageResponse(t *testing.T) {
 	asserts := assert.New(t)
+	cache.Set("pack_size_0", uint64(0), 0)
 
 	{
 		user := model.User{
@@ -58,4 +60,15 @@ func TestBuildUserStorageResponse(t *testing.T) {
 		asserts.Equal(uint64(10), res.Data.(storage).Total)
 		asserts.Equal(uint64(0), res.Data.(storage).Free)
 	}
+	{
+		cache.Set("pack_size_0", uint64(1), 0)
+		user := model.User{
+			Storage: 6,
+			Group:   model.Group{MaxStorage: 10},
+		}
+		res := BuildUserStorageResponse(user)
+		asserts.Equal(uint64(6), res.Data.(storage).Used)
+		asserts.Equal(uint64(11), res.Data.(storage).Total)
+		asserts.Equal(uint64(5), res.Data.(storage).Free)
+	}
 }