feat(files): enable space limit

src/handlers/fileshdr/handlers.go

@@ -129,17 +129,29 @@ func (h *FileHandlers) Create(c *gin.Context) {
 		return
 	}
 	role := c.MustGet(q.RoleParam).(string)
-	userID := c.MustGet(q.UserIDParam).(string)
 	userName := c.MustGet(q.UserParam).(string)
 	if !h.canAccess(userName, role, "create", req.Path) {
 		c.JSON(q.ErrResp(c, 403, q.ErrAccessDenied))
 		return
 	}
 
+	userID := c.MustGet(q.UserIDParam).(string)
+	userIDInt, err := strconv.ParseUint(userID, 10, 64)
+	if err != nil {
+		c.JSON(q.ErrResp(c, 500, err))
+		return
+	}
+
 	tmpFilePath := q.UploadPath(userName, req.Path)
 	locker := h.NewAutoLocker(c, lockName(tmpFilePath))
 	locker.Exec(func() {
-		err := h.deps.FS().Create(tmpFilePath)
+		err = h.deps.Users().SetUsed(userIDInt, true, req.FileSize)
+		if err != nil {
+			c.JSON(q.ErrResp(c, 500, err))
+			return
+		}
+
+		err = h.deps.FS().Create(tmpFilePath)
 		if err != nil {
 			if os.IsExist(err) {
 				c.JSON(q.ErrResp(c, 304, err))
@@ -177,13 +189,35 @@ func (h *FileHandlers) Delete(c *gin.Context) {
 		return
 	}
 
-	err := h.deps.FS().Remove(filePath)
+	userID := c.MustGet(q.UserIDParam).(string)
+	userIDInt, err := strconv.ParseUint(userID, 10, 64)
 	if err != nil {
 		c.JSON(q.ErrResp(c, 500, err))
 		return
 	}
 
-	c.JSON(q.Resp(200))
+	locker := h.NewAutoLocker(c, lockName(filePath))
+	locker.Exec(func() {
+		info, err := h.deps.FS().Stat(filePath)
+		if err != nil {
+			c.JSON(q.ErrResp(c, 500, err))
+			return
+		}
+
+		err = h.deps.Users().SetUsed(userIDInt, false, info.Size())
+		if err != nil {
+			c.JSON(q.ErrResp(c, 500, err))
+			return
+		}
+
+		err = h.deps.FS().Remove(filePath)
+		if err != nil {
+			c.JSON(q.ErrResp(c, 500, err))
+			return
+		}
+
+		c.JSON(q.Resp(200))
+	})
 }
 
 type MetadataResp struct {
@@ -648,13 +682,30 @@ func (h *FileHandlers) DelUploading(c *gin.Context) {
 		c.JSON(q.ErrResp(c, 400, errors.New("invalid file path")))
 		return
 	}
-	userID := c.MustGet(q.UserIDParam).(string)
 
-	var err error
+	userID := c.MustGet(q.UserIDParam).(string)
+	userIDInt, err := strconv.ParseUint(userID, 10, 64)
+	if err != nil {
+		c.JSON(q.ErrResp(c, 500, err))
+		return
+	}
+
 	userName := c.MustGet(q.UserParam).(string)
 	tmpFilePath := q.UploadPath(userName, filePath)
 	locker := h.NewAutoLocker(c, lockName(tmpFilePath))
 	locker.Exec(func() {
+		_, size, _, err := h.uploadMgr.GetInfo(userID, tmpFilePath)
+		if err != nil {
+			c.JSON(q.ErrResp(c, 500, err))
+			return
+		}
+
+		err = h.deps.Users().SetUsed(userIDInt, false, size)
+		if err != nil {
+			c.JSON(q.ErrResp(c, 500, err))
+			return
+		}
+
 		err = h.deps.FS().Remove(tmpFilePath)
 		if err != nil {
 			c.JSON(q.ErrResp(c, 500, err))
@@ -666,6 +717,7 @@ func (h *FileHandlers) DelUploading(c *gin.Context) {
 			c.JSON(q.ErrResp(c, 500, err))
 			return
 		}
+
 		c.JSON(q.Resp(200))
 	})
 }

src/handlers/multiusers/handlers.go

@@ -364,7 +364,7 @@ func (h *MultiUsersSvc) AddUser(c *gin.Context) {
 		Pwd:  string(pwdHash),
 		Role: req.Role,
 		Quota: &userstore.Quota{
-			SpaceLimit:         h.cfg.IntOr("Users.SpaceLimit", 1024),
+			SpaceLimit:         int64(h.cfg.IntOr("Users.SpaceLimit", 100*1024*1024)), // TODO: support int64
 			UploadSpeedLimit:   h.cfg.IntOr("Users.UploadSpeedLimit", 100*1024),
 			DownloadSpeedLimit: h.cfg.IntOr("Users.DownloadSpeedLimit", 100*1024),
 		},
@@ -556,9 +556,10 @@ func (h *MultiUsersSvc) isValidRole(role string) error {
 }
 
 type SelfResp struct {
-	ID   string `json:"id"`
+	ID        string `json:"id"`
-	Name string `json:"name"`
+	Name      string `json:"name"`
-	Role string `json:"role"`
+	Role      string `json:"role"`
+	UsedSpace int64  `json:"usedSpace,string"`
 }
 
 func (h *MultiUsersSvc) Self(c *gin.Context) {
@@ -568,9 +569,16 @@ func (h *MultiUsersSvc) Self(c *gin.Context) {
 		return
 	}
 
+	user, err := h.deps.Users().GetUserByName(claims[q.UserParam])
+	if err != nil {
+		c.JSON(q.ErrResp(c, 500, err))
+		return
+	}
+
 	c.JSON(200, &SelfResp{
-		ID:   claims[q.UserIDParam],
+		ID:        claims[q.UserIDParam],
-		Name: claims[q.UserParam],
+		Name:      claims[q.UserParam],
-		Role: claims[q.RoleParam],
+		Role:      claims[q.RoleParam],
+		UsedSpace: user.UsedSpace,
 	})
 }

src/server/server_concurrency_test.go

@@ -82,7 +82,11 @@ func TestConcurrency(t *testing.T) {
 		users[userName] = adResp.ID
 	}
 
-	filesSize := 10
+	getFilePath := func(name string, i int) string {
+		return fmt.Sprintf("%s/files/home_file_%d", name, i)
+	}
+
+	filesCount := 10
 	mockClient := func(name, pwd string, wg *sync.WaitGroup) {
 		usersCl := client.NewSingleUserClient(addr)
 		resp, _, errs := usersCl.Login(name, pwd)
@@ -95,8 +99,8 @@ func TestConcurrency(t *testing.T) {
 
 		files := map[string]string{}
 		content := "12345678"
-		for i := range make([]int, filesSize, filesSize) {
+		for i := range make([]int, filesCount, filesCount) {
-			files[fmt.Sprintf("%s/files/home_file_%d", name, i)] = content
+			files[getFilePath(name, i)] = content
 		}
 
 		for filePath, content := range files {
@@ -118,6 +122,34 @@ func TestConcurrency(t *testing.T) {
 			t.Fatalf("incorrct metadata size (%d)", len(lsResp.Metadatas))
 		}
 
+		resp, selfResp, errs := usersCl.Self(token)
+		if len(errs) > 0 {
+			t.Fatal(errs)
+		} else if resp.StatusCode != 200 {
+			t.Fatal("failed to self")
+		}
+		if selfResp.UsedSpace != int64(filesCount*len(content)) {
+			t.Fatalf("usedSpace(%d) doesn't match (%d)", selfResp.UsedSpace, filesCount*len(content))
+		}
+
+		resp, _, errs = filesCl.Delete(getFilePath(name, 0))
+		if len(errs) > 0 {
+			t.Fatal(errs)
+		} else if resp.StatusCode != 200 {
+			t.Fatal("failed to add user")
+		}
+
+		resp, selfResp, errs = usersCl.Self(token)
+		if len(errs) > 0 {
+			t.Fatal(errs)
+		} else if resp.StatusCode != 200 {
+			t.Fatal("failed to self")
+		}
+		if selfResp.UsedSpace != int64((filesCount-1)*len(content)) {
+			t.Fatalf("usedSpace(%d) doesn't match (%d)", selfResp.UsedSpace, int64((filesCount-1)*len(content)))
+		}
+		fmt.Println("\n\n\n", selfResp.UsedSpace, int64((filesCount-1)*len(content)))
+
 		wg.Done()
 	}
 

src/userstore/user_store.go

@@ -26,17 +26,18 @@ const (
 )
 
 type Quota struct {
-	SpaceLimit         int `json:"spaceLimit"`
+	SpaceLimit         int64 `json:"spaceLimit,space"`
-	UploadSpeedLimit   int `json:"uploadSpeedLimit"`
+	UploadSpeedLimit   int   `json:"uploadSpeedLimit"`
-	DownloadSpeedLimit int `json:"downloadSpeedLimit"`
+	DownloadSpeedLimit int   `json:"downloadSpeedLimit"`
 }
 
 type User struct {
-	ID    uint64 `json:"id,string"`
+	ID        uint64 `json:"id,string"`
-	Name  string `json:"name"`
+	Name      string `json:"name"`
-	Pwd   string `json:"pwd"`
+	Pwd       string `json:"pwd"`
-	Role  string `json:"role"`
+	Role      string `json:"role"`
-	Quota *Quota `json:"quota"`
+	UsedSpace int64  `json:"usedSpace,string"`
+	Quota     *Quota `json:"quota"`
 }
 
 type IUserStore interface {
@@ -47,6 +48,7 @@ type IUserStore interface {
 	GetUser(id uint64) (*User, error)
 	GetUserByName(name string) (*User, error)
 	SetInfo(id uint64, user *User) error
+	SetUsed(id uint64, incr bool, capacity int64) error
 	SetPwd(id uint64, pwd string) error
 	ListUsers() ([]*User, error)
 	AddRole(role string) error
@@ -215,31 +217,6 @@ func (us *KVUserStore) GetUserByName(name string) (*User, error) {
 
 }
 
-// func (us *KVUserStore) SetName(id uint64, name string) error {
-// 	us.mtx.Lock()
-// 	defer us.mtx.Unlock()
-
-// 	_, ok := us.store.GetStringIn(IDsNs, name)
-// 	if ok {
-// 		return fmt.Errorf("user name (%s) exists", name)
-// 	}
-
-// 	userID := fmt.Sprint(id)
-// 	_, ok = us.store.GetStringIn(UsersNs, userID)
-// 	if !ok {
-// 		return fmt.Errorf("Name (%d) does not exist", id)
-// 	}
-// 	if name == "" {
-// 		return fmt.Errorf("Name can not be empty")
-// 	}
-
-// 	err := us.store.SetStringIn(IDsNs, name, userID)
-// 	if err != nil {
-// 		return err
-// 	}
-// 	return us.store.SetStringIn(UsersNs, userID, name)
-// }
-
 func (us *KVUserStore) SetPwd(id uint64, pwd string) error {
 	us.mtx.Lock()
 	defer us.mtx.Unlock()
@@ -265,6 +242,44 @@ func (us *KVUserStore) SetPwd(id uint64, pwd string) error {
 	return us.store.SetStringIn(UsersNs, userID, string(infoBytes))
 }
 
+func (us *KVUserStore) SetUsed(id uint64, incr bool, capacity int64) error {
+	us.mtx.Lock()
+	defer us.mtx.Unlock()
+
+	userID := fmt.Sprint(id)
+	infoStr, ok := us.store.GetStringIn(UsersNs, userID)
+	if !ok {
+		return fmt.Errorf("user (%d) does not exist", id)
+	}
+	gotUser := &User{}
+	err := json.Unmarshal([]byte(infoStr), gotUser)
+	if err != nil {
+		return err
+	} else if gotUser.ID != id {
+		return fmt.Errorf("user id key(%d) info(%d) does match", id, gotUser.ID)
+	}
+
+	if incr && gotUser.UsedSpace+capacity > int64(gotUser.Quota.SpaceLimit) {
+		return fmt.Errorf(
+			"reached space limit (%d): file size(%d), used(%d)",
+			gotUser.Quota.SpaceLimit,
+			capacity,
+			gotUser.UsedSpace,
+		)
+	}
+
+	if incr {
+		gotUser.UsedSpace = gotUser.UsedSpace + capacity
+	} else {
+		gotUser.UsedSpace = gotUser.UsedSpace - capacity
+	}
+	infoBytes, err := json.Marshal(gotUser)
+	if err != nil {
+		return err
+	}
+	return us.store.SetStringIn(UsersNs, userID, string(infoBytes))
+}
+
 func (us *KVUserStore) SetInfo(id uint64, user *User) error {
 	us.mtx.Lock()
 	defer us.mtx.Unlock()
@@ -289,6 +304,9 @@ func (us *KVUserStore) SetInfo(id uint64, user *User) error {
 	if user.Quota != nil {
 		gotUser.Quota = user.Quota
 	}
+	if user.UsedSpace > 0 {
+		gotUser.UsedSpace = user.UsedSpace
+	}
 
 	infoBytes, err := json.Marshal(gotUser)
 	if err != nil {
@@ -297,19 +315,6 @@ func (us *KVUserStore) SetInfo(id uint64, user *User) error {
 	return us.store.SetStringIn(UsersNs, userID, string(infoBytes))
 }
 
-// func (us *KVUserStore) SetRole(id uint64, role string) error {
-// 	us.mtx.Lock()
-// 	defer us.mtx.Unlock()
-
-// 	userID := fmt.Sprint(id)
-// 	_, ok := us.store.GetStringIn(RolesNs, userID)
-// 	if !ok {
-// 		return fmt.Errorf("Role (%d) does not exist", id)
-// 	}
-
-// 	return us.store.SetStringIn(RolesNs, userID, role)
-// }
-
 func (us *KVUserStore) ListUsers() ([]*User, error) {
 	us.mtx.RLock()
 	defer us.mtx.RUnlock()

src/userstore/user_store_test.go

@@ -3,6 +3,7 @@ package userstore
 import (
 	"io/ioutil"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/ihexxa/quickshare/src/kvstore/boltdbpvd"
@@ -38,8 +39,8 @@ func TestUserStores(t *testing.T) {
 		id, name1 := uint64(1), "test_user1"
 		pwd1, pwd2 := "666", "888"
 		role1, role2 := UserRole, AdminRole
-		spaceLimit1, upLimit1, downLimit1 := 3, 5, 7
+		spaceLimit1, upLimit1, downLimit1 := 17, 5, 7
-		spaceLimit2, upLimit2, downLimit2 := 11, 13, 17
+		spaceLimit2, upLimit2, downLimit2 := 19, 13, 17
 
 		err = store.AddUser(&User{
 			ID:   id,
@@ -110,6 +111,22 @@ func TestUserStores(t *testing.T) {
 			},
 		})
 
+		usedIncr, usedDecr := int64(spaceLimit2), int64(7)
+		err = store.SetUsed(id, true, usedIncr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = store.SetUsed(id, false, usedDecr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = store.SetUsed(id, true, int64(spaceLimit2)-(usedIncr-usedDecr)+1)
+		if err == nil || !strings.Contains(err.Error(), "reached space limit") {
+			t.Fatal("should reject big file")
+		} else {
+			err = nil
+		}
+
 		user, err = store.GetUser(id)
 		if err != nil {
 			t.Fatal(err)
@@ -129,6 +146,9 @@ func TestUserStores(t *testing.T) {
 		if user.Quota.DownloadSpeedLimit != downLimit2 {
 			t.Fatalf("down limit not matched %d %d", downLimit2, user.Quota.DownloadSpeedLimit)
 		}
+		if user.UsedSpace != usedIncr-usedDecr {
+			t.Fatalf("used space not matched %d %d", user.UsedSpace, usedIncr-usedDecr)
+		}
 
 		user, err = store.GetUserByName(name1)
 		if err != nil {