beenull/sftpgo-mirror
1
0
Fork 0
mirror of https://github.com/drakkan/sftpgo.git synced 2024-11-21 23:20:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

quota tracking: fix concurrent updates

Browse source 
added a test case to check quota size for upload that replace an existing file
This commit is contained in:
Nicola Murino 2019-07-28 13:24:46 +02:00
parent b728a06984
commit c547c5c9ee
6 changed files with 85 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
dataprovider/sqlcommon.go
View file

@ -96,39 +96,17 @@ func sqlCommonGetUserByID(ID int64) (User, error) {
}
func sqlCommonUpdateQuota(username string, filesAdd int, sizeAdd int64, reset bool, p Provider) error {
var usedFiles int
var usedSize int64
var err error
if reset {
usedFiles = 0
usedSize = 0
} else {
usedFiles, usedSize, err = p.getUsedQuota(username)
if err != nil {
return err
}
}
usedFiles += filesAdd
usedSize += sizeAdd
if usedFiles < 0 {
logger.Warn(logSender, "used files is negative, probably some files were added and not tracked, please rescan quota!")
usedFiles = 0
}
if usedSize < 0 {
logger.Warn(logSender, "used files is negative, probably some files were added and not tracked, please rescan quota!")
usedSize = 0
}
q := getUpdateQuotaQuery()
q := getUpdateQuotaQuery(reset)
stmt, err := dbHandle.Prepare(q)
if err != nil {
logger.Debug(logSender, "error preparing database query %v: %v", q, err)
return err
}
defer stmt.Close()
_, err = stmt.Exec(usedSize, usedFiles, utils.GetTimeAsMsSinceEpoch(time.Now()), username)
_, err = stmt.Exec(sizeAdd, filesAdd, utils.GetTimeAsMsSinceEpoch(time.Now()), username)
if err == nil {
logger.Debug(logSender, "quota updated for user %v, new files: %v new size: %v", username, usedFiles, usedSize)
logger.Debug(logSender, "quota updated for user %v, files increment: %v size increment: %v is reset? %v",
username, filesAdd, sizeAdd, reset)
} else {
logger.Warn(logSender, "error updating quota for username %v: %v", username, err)
}

6
dataprovider/sqlqueries.go
View file

@ -36,10 +36,14 @@ func getUsersQuery(order string, username string) string {
order, sqlPlaceholders[0], sqlPlaceholders[1])
}
func getUpdateQuotaQuery() string {
func getUpdateQuotaQuery(reset bool) string {
if reset {
return fmt.Sprintf(`UPDATE %v SET used_quota_size = %v,used_quota_files = %v,last_quota_scan = %v
WHERE username = %v`, config.UsersTable, sqlPlaceholders[0], sqlPlaceholders[1], sqlPlaceholders[2], sqlPlaceholders[3])
}
return fmt.Sprintf(`UPDATE %v SET used_quota_size = used_quota_size + %v,used_quota_files = used_quota_files + %v,last_quota_scan = %v
WHERE username = %v`, config.UsersTable, sqlPlaceholders[0], sqlPlaceholders[1], sqlPlaceholders[2], sqlPlaceholders[3])
}
func getQuotaQuery() string {
return fmt.Sprintf(`SELECT used_quota_size,used_quota_files FROM %v WHERE username = %v`, config.UsersTable,

5
sftpd/handler.go
View file

@ -168,8 +168,9 @@ func (c Connection) Filewrite(request *sftp.Request) (io.WriterAt, error) {
}
if trunc {
// the file is truncated so we need to decrease quota size but not quota files
logger.Debug(logSender, "file truncation requested update quota for user %v", c.User.Username)
dataprovider.UpdateUserQuota(dataProvider, c.User.Username, -1, -stat.Size(), false)
dataprovider.UpdateUserQuota(dataProvider, c.User.Username, 0, -stat.Size(), false)
}
utils.SetPathPermissions(p, c.User.GetUID(), c.User.GetGID())
@ -185,7 +186,7 @@ func (c Connection) Filewrite(request *sftp.Request) (io.WriterAt, error) {
user: c.User,
connectionID: c.ID,
transferType: transferUpload,
isNewFile: trunc,
isNewFile: !trunc,
}
addTransfer(&transfer)
return &transfer, nil

4
sftpd/sftpd.go
View file

@ -298,8 +298,8 @@ func executeAction(operation string, username string, path string, target string
if _, err = os.Stat(actions.Command); err == nil {
command := exec.Command(actions.Command, operation, username, path, target)
err = command.Start()
logger.Debug(logSender, "executed command \"%v\" with arguments: %v, %v, %v, error: %v",
actions.Command, operation, path, target, err)
logger.Debug(logSender, "executed command \"%v\" with arguments: %v, %v, %v, %v, error: %v",
actions.Command, operation, username, path, target, err)
} else {
logger.Warn(logSender, "Invalid action command \"%v\" : %v", actions.Command, err)
}

68
sftpd/sftpd_test.go
View file

@ -574,6 +574,66 @@ func TestMaxSessions(t *testing.T) {
}
}
func TestQuotaFileReplace(t *testing.T) {
usePubKey := false
user, err := api.AddUser(getTestUser(usePubKey), http.StatusOK)
if err != nil {
t.Errorf("unable to add user: %v", err)
}
client, err := getSftpClient(user, usePubKey)
if err != nil {
t.Errorf("unable to create sftp client: %v", err)
} else {
defer client.Close()
testFileSize := int64(65535)
expectedQuotaSize := user.UsedQuotaSize + testFileSize
expectedQuotaFiles := user.UsedQuotaFiles + 1
testFileName := "test_file.dat"
testFilePath := filepath.Join(homeBasePath, testFileName)
err = createTestFile(testFilePath, testFileSize)
if err != nil {
t.Errorf("unable to create test file: %v", err)
}
err = sftpUploadFile(testFilePath, testFileName, testFileSize, client)
if err != nil {
t.Errorf("file upload error: %v", err)
}
user, err = api.GetUserByID(user.ID, http.StatusOK)
if err != nil {
t.Errorf("error getting user: %v", err)
}
if expectedQuotaFiles != user.UsedQuotaFiles {
t.Errorf("quota files does not match, expected: %v, actual: %v", expectedQuotaFiles, user.UsedQuotaFiles)
}
if expectedQuotaSize != user.UsedQuotaSize {
t.Errorf("quota size does not match, expected: %v, actual: %v", expectedQuotaSize, user.UsedQuotaSize)
}
// now replace the same file, the quota must not change
err = sftpUploadFile(testFilePath, testFileName, testFileSize, client)
if err != nil {
t.Errorf("file upload error: %v", err)
}
user, err = api.GetUserByID(user.ID, http.StatusOK)
if err != nil {
t.Errorf("error getting user: %v", err)
}
if expectedQuotaFiles != user.UsedQuotaFiles {
t.Errorf("quota files does not match, expected: %v, actual: %v", expectedQuotaFiles, user.UsedQuotaFiles)
}
if expectedQuotaSize != user.UsedQuotaSize {
t.Errorf("quota size does not match, expected: %v, actual: %v", expectedQuotaSize, user.UsedQuotaSize)
}
err = client.Remove(testFileName)
if err != nil {
t.Errorf("error removing uploaded file: %v", err)
}
}
err = api.RemoveUser(user, http.StatusOK)
if err != nil {
t.Errorf("unable to remove user: %v", err)
}
}
func TestQuotaScan(t *testing.T) {
usePubKey := false
user, err := api.AddUser(getTestUser(usePubKey), http.StatusOK)
@ -670,14 +730,18 @@ func TestQuotaSize(t *testing.T) {
if err != nil {
t.Errorf("unable to create test file: %v", err)
}
err = sftpUploadFile(testFilePath, testFileName, testFileSize, client)
err = sftpUploadFile(testFilePath, testFileName+".quota", testFileSize, client)
if err != nil {
t.Errorf("file upload error: %v", err)
}
err = sftpUploadFile(testFilePath, testFileName+".1", testFileSize, client)
err = sftpUploadFile(testFilePath, testFileName+".quota.1", testFileSize, client)
if err == nil {
t.Errorf("user is over quota file upload must fail")
}
err = client.Remove(testFileName + ".quota")
if err != nil {
t.Errorf("error removing uploaded file: %v", err)
}
}
err = api.RemoveUser(user, http.StatusOK)
if err != nil {

6
sftpd/transfer.go
View file

@ -57,9 +57,11 @@ func (t *Transfer) Close() error {
executeAction(operationUpload, t.user.Username, t.path, "")
}
removeTransfer(t)
if t.transferType == transferUpload && t.bytesReceived > 0 && t.isNewFile {
if t.transferType == transferUpload {
numFiles := 0
numFiles++
if t.isNewFile {
numFiles = 1
}
dataprovider.UpdateUserQuota(dataProvider, t.user.Username, numFiles, t.bytesReceived, false)
}
return err