diff --git a/src/handlers/fileshdr/handlers.go b/src/handlers/fileshdr/handlers.go index f90975f..820d1e3 100644 --- a/src/handlers/fileshdr/handlers.go +++ b/src/handlers/fileshdr/handlers.go @@ -147,7 +147,11 @@ func (h *FileHandlers) Create(c *gin.Context) { locker.Exec(func() { err = h.deps.Users().SetUsed(userIDInt, true, req.FileSize) if err != nil { - c.JSON(q.ErrResp(c, 500, err)) + if userstore.IsReachedLimitErr(err) { + c.JSON(q.ErrResp(c, 429, err)) + } else { + c.JSON(q.ErrResp(c, 500, err)) + } return } @@ -206,7 +210,11 @@ func (h *FileHandlers) Delete(c *gin.Context) { err = h.deps.Users().SetUsed(userIDInt, false, info.Size()) if err != nil { - c.JSON(q.ErrResp(c, 500, err)) + if userstore.IsReachedLimitErr(err) { + c.JSON(q.ErrResp(c, 429, err)) + } else { + c.JSON(q.ErrResp(c, 500, err)) + } return } @@ -702,7 +710,11 @@ func (h *FileHandlers) DelUploading(c *gin.Context) { err = h.deps.Users().SetUsed(userIDInt, false, size) if err != nil { - c.JSON(q.ErrResp(c, 500, err)) + if userstore.IsReachedLimitErr(err) { + c.JSON(q.ErrResp(c, 429, err)) + } else { + c.JSON(q.ErrResp(c, 500, err)) + } return } diff --git a/src/server/server_files_test.go b/src/server/server_files_test.go index f36dba6..6b10f8e 100644 --- a/src/server/server_files_test.go +++ b/src/server/server_files_test.go @@ -25,7 +25,7 @@ func TestFileHandlers(t *testing.T) { "captchaEnabled": false, "uploadSpeedLimit": 409600, "downloadSpeedLimit": 409600, - "spaceLimit": 1024, + "spaceLimit": 1000, "limiterCapacity": 1000, "limiterCyc": 1000 }, diff --git a/src/server/server_space_limit_test.go b/src/server/server_space_limit_test.go new file mode 100644 index 0000000..48bf894 --- /dev/null +++ b/src/server/server_space_limit_test.go @@ -0,0 +1,138 @@ +package server + +import ( + "fmt" + "os" + "testing" + + "github.com/ihexxa/quickshare/src/client" + q "github.com/ihexxa/quickshare/src/handlers" + "github.com/ihexxa/quickshare/src/userstore" +) + +func TestSpaceLimit(t *testing.T) { + addr := "http://127.0.0.1:8686" + root := "testData" + config := `{ + "users": { + "enableAuth": true, + "minUserNameLen": 2, + "minPwdLen": 4, + "captchaEnabled": false, + "uploadSpeedLimit": 409600, + "downloadSpeedLimit": 409600, + "spaceLimit": 100, + "limiterCapacity": 1000, + "limiterCyc": 1000 + }, + "server": { + "debug": true + }, + "fs": { + "root": "testData" + } + }` + + adminName := "qs" + adminPwd := "quicksh@re" + os.Setenv("DEFAULTADMIN", adminName) + os.Setenv("DEFAULTADMINPWD", adminPwd) + + os.RemoveAll(root) + err := os.MkdirAll(root, 0700) + if err != nil { + t.Fatal(err) + } + defer os.RemoveAll(root) + + srv := startTestServer(config) + defer srv.Shutdown() + // fs := srv.depsFS() + if !waitForReady(addr) { + t.Fatal("fail to start server") + } + + usersCl := client.NewSingleUserClient(addr) + resp, _, errs := usersCl.Login(adminName, adminPwd) + if len(errs) > 0 { + t.Fatal(errs) + } else if resp.StatusCode != 200 { + t.Fatal(resp.StatusCode) + } + token := client.GetCookie(resp.Cookies(), q.TokenCookie) + + userCount := 1 + userPwd := "1234" + users := map[string]string{} + getUserName := func(id int) string { + return fmt.Sprintf("space_limit_user_%d", id) + } + + for i := 0; i < userCount; i++ { + userName := getUserName(i) + + resp, adResp, errs := usersCl.AddUser(userName, userPwd, userstore.UserRole, token) + if len(errs) > 0 { + t.Fatal(errs) + } else if resp.StatusCode != 200 { + t.Fatal("failed to add user") + } + + users[userName] = adResp.ID + } + + resp, _, errs = usersCl.Logout(token) + if len(errs) > 0 { + t.Fatal(errs) + } else if resp.StatusCode != 200 { + t.Fatal(resp.StatusCode) + } + + t.Run("test space limitiong: Upload", func(t *testing.T) { + usersCl := client.NewSingleUserClient(addr) + resp, _, errs := usersCl.Login(getUserName(0), userPwd) + if len(errs) > 0 { + t.Fatal(errs) + } else if resp.StatusCode != 200 { + t.Fatal(resp.StatusCode) + } + token := client.GetCookie(resp.Cookies(), q.TokenCookie) + + fileContent := "" + for i := 0; i < 10; i++ { + fileContent += "0" + } + + for i := 0; i < 10; i++ { + ok := assertUploadOK(t, fmt.Sprintf("%s/spacelimit/f_%d", getUserName(0), 0), fileContent, addr, token) + if !ok { + t.Fatalf("space limit failed at %d", 0) + } + + resp, selfResp, errs := usersCl.Self(token) + if len(errs) > 0 { + t.Fatal(errs) + } else if resp.StatusCode != 200 { + t.Fatal("failed to add user") + } else if selfResp.UsedSpace != int64((i+1)*10) { + t.Fatal("incorrect used space") + } + } + + cl := client.NewFilesClient(addr, token) + filePath := fmt.Sprintf("%s/spacelimit/f_%d", getUserName(0), 11) + res, _, errs := cl.Create(filePath, 1) + if len(errs) > 0 { + t.Fatal(errs) + } else if res.StatusCode != 429 { + t.Fatal("(space limit): this request should be rejected") + } + }) + + resp, _, errs = usersCl.Logout(token) + if len(errs) > 0 { + t.Fatal(errs) + } else if resp.StatusCode != 200 { + t.Fatal(resp.StatusCode) + } +} diff --git a/src/userstore/user_store.go b/src/userstore/user_store.go index 94acad8..8a11aa1 100644 --- a/src/userstore/user_store.go +++ b/src/userstore/user_store.go @@ -25,6 +25,14 @@ const ( defaultDownloadSpeedLimit = 50 * 1024 * 1024 // 50MB ) +var ( + ErrReachedLimit = errors.New("reached space limit") +) + +func IsReachedLimitErr(err error) bool { + return err == ErrReachedLimit +} + type Quota struct { SpaceLimit int64 `json:"spaceLimit,space"` UploadSpeedLimit int `json:"uploadSpeedLimit"` @@ -260,12 +268,7 @@ func (us *KVUserStore) SetUsed(id uint64, incr bool, capacity int64) error { } 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, - ) + return ErrReachedLimit } if incr {