// Copyright (C) 2019 Nicola Murino // // This program is free software: you can redistribute it and/or modify // it under the terms of the GNU Affero General Public License as published // by the Free Software Foundation, version 3. // // This program is distributed in the hope that it will be useful, // but WITHOUT ANY WARRANTY; without even the implied warranty of // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the // GNU Affero General Public License for more details. // // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . package common import ( "fmt" "os" "path" "path/filepath" "strconv" "strings" "testing" "time" "github.com/rs/xid" "github.com/sftpgo/sdk" "github.com/stretchr/testify/assert" "github.com/drakkan/sftpgo/v2/internal/dataprovider" "github.com/drakkan/sftpgo/v2/internal/util" "github.com/drakkan/sftpgo/v2/internal/vfs" ) func TestTransfersCheckerDiskQuota(t *testing.T) { username := "transfers_check_username" folderName := "test_transfers_folder" groupName := "test_transfers_group" vdirPath := "/vdir" group := dataprovider.Group{ BaseGroup: sdk.BaseGroup{ Name: groupName, }, UserSettings: dataprovider.GroupUserSettings{ BaseGroupUserSettings: sdk.BaseGroupUserSettings{ QuotaSize: 120, }, }, } folder := vfs.BaseVirtualFolder{ Name: folderName, MappedPath: filepath.Join(os.TempDir(), folderName), } user := dataprovider.User{ BaseUser: sdk.BaseUser{ Username: username, Password: "testpwd", HomeDir: filepath.Join(os.TempDir(), username), Status: 1, QuotaSize: 0, // the quota size defined for the group is used Permissions: map[string][]string{ "/": {dataprovider.PermAny}, }, }, VirtualFolders: []vfs.VirtualFolder{ { BaseVirtualFolder: vfs.BaseVirtualFolder{ Name: folderName, }, VirtualPath: vdirPath, QuotaSize: 100, }, }, Groups: []sdk.GroupMapping{ { Name: groupName, Type: sdk.GroupTypePrimary, }, }, } err := dataprovider.AddGroup(&group, "", "", "") assert.NoError(t, err) group, err = dataprovider.GroupExists(groupName) assert.NoError(t, err) err = dataprovider.AddFolder(&folder, "", "", "") assert.NoError(t, err) assert.Equal(t, int64(120), group.UserSettings.QuotaSize) err = dataprovider.AddUser(&user, "", "", "") assert.NoError(t, err) user, err = dataprovider.GetUserWithGroupSettings(username, "") assert.NoError(t, err) connID1 := xid.New().String() fsUser, err := user.GetFilesystemForPath("/file1", connID1) assert.NoError(t, err) conn1 := NewBaseConnection(connID1, ProtocolSFTP, "", "", user) fakeConn1 := &fakeConnection{ BaseConnection: conn1, } transfer1 := NewBaseTransfer(nil, conn1, nil, filepath.Join(user.HomeDir, "file1"), filepath.Join(user.HomeDir, "file1"), "/file1", TransferUpload, 0, 0, 120, 0, true, fsUser, dataprovider.TransferQuota{}) transfer1.BytesReceived.Store(150) err = Connections.Add(fakeConn1) assert.NoError(t, err) // the transferschecker will do nothing if there is only one ongoing transfer Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) connID2 := xid.New().String() conn2 := NewBaseConnection(connID2, ProtocolSFTP, "", "", user) fakeConn2 := &fakeConnection{ BaseConnection: conn2, } transfer2 := NewBaseTransfer(nil, conn2, nil, filepath.Join(user.HomeDir, "file2"), filepath.Join(user.HomeDir, "file2"), "/file2", TransferUpload, 0, 0, 120, 40, true, fsUser, dataprovider.TransferQuota{}) transfer1.BytesReceived.Store(50) transfer2.BytesReceived.Store(60) err = Connections.Add(fakeConn2) assert.NoError(t, err) connID3 := xid.New().String() conn3 := NewBaseConnection(connID3, ProtocolSFTP, "", "", user) fakeConn3 := &fakeConnection{ BaseConnection: conn3, } transfer3 := NewBaseTransfer(nil, conn3, nil, filepath.Join(user.HomeDir, "file3"), filepath.Join(user.HomeDir, "file3"), "/file3", TransferDownload, 0, 0, 120, 0, true, fsUser, dataprovider.TransferQuota{}) transfer3.BytesReceived.Store(60) // this value will be ignored, this is a download err = Connections.Add(fakeConn3) assert.NoError(t, err) // the transfers are not overquota Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) transfer1.BytesReceived.Store(80) // truncated size will be subtracted, we are not overquota Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) transfer1.BytesReceived.Store(120) // we are now overquota // if another check is in progress nothing is done Connections.transfersCheckStatus.Store(true) Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) Connections.transfersCheckStatus.Store(false) Connections.checkTransfers() assert.True(t, conn1.IsQuotaExceededError(transfer1.errAbort), transfer1.errAbort) assert.True(t, conn2.IsQuotaExceededError(transfer2.errAbort), transfer2.errAbort) assert.True(t, conn1.IsQuotaExceededError(transfer1.GetAbortError())) assert.Nil(t, transfer3.errAbort) assert.True(t, conn3.IsQuotaExceededError(transfer3.GetAbortError())) // update the user quota size group.UserSettings.QuotaSize = 1000 err = dataprovider.UpdateGroup(&group, []string{username}, "", "", "") assert.NoError(t, err) transfer1.errAbort = nil transfer2.errAbort = nil Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) group.UserSettings.QuotaSize = 0 err = dataprovider.UpdateGroup(&group, []string{username}, "", "", "") assert.NoError(t, err) Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) // now check a public folder transfer1.BytesReceived.Store(0) transfer2.BytesReceived.Store(0) connID4 := xid.New().String() fsFolder, err := user.GetFilesystemForPath(path.Join(vdirPath, "/file1"), connID4) assert.NoError(t, err) conn4 := NewBaseConnection(connID4, ProtocolSFTP, "", "", user) fakeConn4 := &fakeConnection{ BaseConnection: conn4, } transfer4 := NewBaseTransfer(nil, conn4, nil, filepath.Join(os.TempDir(), folderName, "file1"), filepath.Join(os.TempDir(), folderName, "file1"), path.Join(vdirPath, "/file1"), TransferUpload, 0, 0, 100, 0, true, fsFolder, dataprovider.TransferQuota{}) err = Connections.Add(fakeConn4) assert.NoError(t, err) connID5 := xid.New().String() conn5 := NewBaseConnection(connID5, ProtocolSFTP, "", "", user) fakeConn5 := &fakeConnection{ BaseConnection: conn5, } transfer5 := NewBaseTransfer(nil, conn5, nil, filepath.Join(os.TempDir(), folderName, "file2"), filepath.Join(os.TempDir(), folderName, "file2"), path.Join(vdirPath, "/file2"), TransferUpload, 0, 0, 100, 0, true, fsFolder, dataprovider.TransferQuota{}) err = Connections.Add(fakeConn5) assert.NoError(t, err) transfer4.BytesReceived.Store(50) transfer5.BytesReceived.Store(40) Connections.checkTransfers() assert.Nil(t, transfer4.errAbort) assert.Nil(t, transfer5.errAbort) transfer5.BytesReceived.Store(60) Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) assert.True(t, conn1.IsQuotaExceededError(transfer4.errAbort)) assert.True(t, conn2.IsQuotaExceededError(transfer5.errAbort)) if dataprovider.GetProviderStatus().Driver != dataprovider.MemoryDataProviderName { providerConf := dataprovider.GetProviderConfig() err = dataprovider.Close() assert.NoError(t, err) transfer4.errAbort = nil transfer5.errAbort = nil Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) assert.Nil(t, transfer3.errAbort) assert.Nil(t, transfer4.errAbort) assert.Nil(t, transfer5.errAbort) err = dataprovider.Initialize(providerConf, configDir, true) assert.NoError(t, err) } err = transfer1.Close() assert.NoError(t, err) err = transfer2.Close() assert.NoError(t, err) err = transfer3.Close() assert.NoError(t, err) err = transfer4.Close() assert.NoError(t, err) err = transfer5.Close() assert.NoError(t, err) Connections.Remove(fakeConn1.GetID()) Connections.Remove(fakeConn2.GetID()) Connections.Remove(fakeConn3.GetID()) Connections.Remove(fakeConn4.GetID()) Connections.Remove(fakeConn5.GetID()) stats := Connections.GetStats("") assert.Len(t, stats, 0) err = dataprovider.DeleteUser(user.Username, "", "", "") assert.NoError(t, err) err = os.RemoveAll(user.GetHomeDir()) assert.NoError(t, err) err = dataprovider.DeleteFolder(folderName, "", "", "") assert.NoError(t, err) err = os.RemoveAll(filepath.Join(os.TempDir(), folderName)) assert.NoError(t, err) err = dataprovider.DeleteGroup(groupName, "", "", "") assert.NoError(t, err) } func TestTransferCheckerTransferQuota(t *testing.T) { username := "transfers_check_username" user := dataprovider.User{ BaseUser: sdk.BaseUser{ Username: username, Password: "test_pwd", HomeDir: filepath.Join(os.TempDir(), username), Status: 1, TotalDataTransfer: 1, Permissions: map[string][]string{ "/": {dataprovider.PermAny}, }, }, } err := dataprovider.AddUser(&user, "", "", "") assert.NoError(t, err) connID1 := xid.New().String() fsUser, err := user.GetFilesystemForPath("/file1", connID1) assert.NoError(t, err) conn1 := NewBaseConnection(connID1, ProtocolSFTP, "", "192.168.1.1", user) fakeConn1 := &fakeConnection{ BaseConnection: conn1, } transfer1 := NewBaseTransfer(nil, conn1, nil, filepath.Join(user.HomeDir, "file1"), filepath.Join(user.HomeDir, "file1"), "/file1", TransferUpload, 0, 0, 0, 0, true, fsUser, dataprovider.TransferQuota{AllowedTotalSize: 100}) transfer1.BytesReceived.Store(150) err = Connections.Add(fakeConn1) assert.NoError(t, err) // the transferschecker will do nothing if there is only one ongoing transfer Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) connID2 := xid.New().String() conn2 := NewBaseConnection(connID2, ProtocolSFTP, "", "127.0.0.1", user) fakeConn2 := &fakeConnection{ BaseConnection: conn2, } transfer2 := NewBaseTransfer(nil, conn2, nil, filepath.Join(user.HomeDir, "file2"), filepath.Join(user.HomeDir, "file2"), "/file2", TransferUpload, 0, 0, 0, 0, true, fsUser, dataprovider.TransferQuota{AllowedTotalSize: 100}) transfer2.BytesReceived.Store(150) err = Connections.Add(fakeConn2) assert.NoError(t, err) Connections.checkTransfers() assert.Nil(t, transfer1.errAbort) assert.Nil(t, transfer2.errAbort) // now test overquota transfer1.BytesReceived.Store(1024*1024 + 1) transfer2.BytesReceived.Store(0) Connections.checkTransfers() assert.True(t, conn1.IsQuotaExceededError(transfer1.errAbort), transfer1.errAbort) assert.Nil(t, transfer2.errAbort) transfer1.errAbort = nil transfer1.BytesReceived.Store(1024*1024 + 1) transfer2.BytesReceived.Store(1024) Connections.checkTransfers() assert.True(t, conn1.IsQuotaExceededError(transfer1.errAbort)) assert.True(t, conn2.IsQuotaExceededError(transfer2.errAbort)) transfer1.BytesReceived.Store(0) transfer2.BytesReceived.Store(0) transfer1.errAbort = nil transfer2.errAbort = nil err = transfer1.Close() assert.NoError(t, err) err = transfer2.Close() assert.NoError(t, err) Connections.Remove(fakeConn1.GetID()) Connections.Remove(fakeConn2.GetID()) connID3 := xid.New().String() conn3 := NewBaseConnection(connID3, ProtocolSFTP, "", "", user) fakeConn3 := &fakeConnection{ BaseConnection: conn3, } transfer3 := NewBaseTransfer(nil, conn3, nil, filepath.Join(user.HomeDir, "file1"), filepath.Join(user.HomeDir, "file1"), "/file1", TransferDownload, 0, 0, 0, 0, true, fsUser, dataprovider.TransferQuota{AllowedDLSize: 100}) transfer3.BytesSent.Store(150) err = Connections.Add(fakeConn3) assert.NoError(t, err) connID4 := xid.New().String() conn4 := NewBaseConnection(connID4, ProtocolSFTP, "", "", user) fakeConn4 := &fakeConnection{ BaseConnection: conn4, } transfer4 := NewBaseTransfer(nil, conn4, nil, filepath.Join(user.HomeDir, "file2"), filepath.Join(user.HomeDir, "file2"), "/file2", TransferDownload, 0, 0, 0, 0, true, fsUser, dataprovider.TransferQuota{AllowedDLSize: 100}) transfer4.BytesSent.Store(150) err = Connections.Add(fakeConn4) assert.NoError(t, err) Connections.checkTransfers() assert.Nil(t, transfer3.errAbort) assert.Nil(t, transfer4.errAbort) transfer3.BytesSent.Store(512 * 1024) transfer4.BytesSent.Store(512*1024 + 1) Connections.checkTransfers() if assert.Error(t, transfer3.errAbort) { assert.Contains(t, transfer3.errAbort.Error(), ErrReadQuotaExceeded.Error()) } if assert.Error(t, transfer4.errAbort) { assert.Contains(t, transfer4.errAbort.Error(), ErrReadQuotaExceeded.Error()) } Connections.Remove(fakeConn3.GetID()) Connections.Remove(fakeConn4.GetID()) stats := Connections.GetStats("") assert.Len(t, stats, 0) err = dataprovider.DeleteUser(user.Username, "", "", "") assert.NoError(t, err) err = os.RemoveAll(user.GetHomeDir()) assert.NoError(t, err) } func TestAggregateTransfers(t *testing.T) { checker := transfersCheckerMem{} checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "1", Username: "user", FolderName: "", TruncatedSize: 0, CurrentULSize: 100, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations := checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 0) assert.Len(t, aggregations, 1) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferDownload, ConnID: "2", Username: "user", FolderName: "", TruncatedSize: 0, CurrentULSize: 0, CurrentDLSize: 100, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 0) assert.Len(t, aggregations, 1) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "3", Username: "user", FolderName: "folder", TruncatedSize: 0, CurrentULSize: 10, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 0) assert.Len(t, aggregations, 2) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "4", Username: "user1", FolderName: "", TruncatedSize: 0, CurrentULSize: 100, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 0) assert.Len(t, aggregations, 3) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "5", Username: "user", FolderName: "", TruncatedSize: 0, CurrentULSize: 100, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 1) val, ok := usersToFetch["user"] assert.True(t, ok) assert.False(t, val) assert.Len(t, aggregations, 3) aggregate, ok := aggregations[0] assert.True(t, ok) assert.Len(t, aggregate, 2) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "6", Username: "user", FolderName: "", TruncatedSize: 0, CurrentULSize: 100, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 1) val, ok = usersToFetch["user"] assert.True(t, ok) assert.False(t, val) assert.Len(t, aggregations, 3) aggregate, ok = aggregations[0] assert.True(t, ok) assert.Len(t, aggregate, 3) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "7", Username: "user", FolderName: "folder", TruncatedSize: 0, CurrentULSize: 10, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 1) val, ok = usersToFetch["user"] assert.True(t, ok) assert.True(t, val) assert.Len(t, aggregations, 3) aggregate, ok = aggregations[0] assert.True(t, ok) assert.Len(t, aggregate, 3) aggregate, ok = aggregations[1] assert.True(t, ok) assert.Len(t, aggregate, 2) checker.AddTransfer(dataprovider.ActiveTransfer{ ID: 1, Type: TransferUpload, ConnID: "8", Username: "user", FolderName: "", TruncatedSize: 0, CurrentULSize: 100, CurrentDLSize: 0, CreatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), UpdatedAt: util.GetTimeAsMsSinceEpoch(time.Now()), }) usersToFetch, aggregations = checker.aggregateUploadTransfers() assert.Len(t, usersToFetch, 1) val, ok = usersToFetch["user"] assert.True(t, ok) assert.True(t, val) assert.Len(t, aggregations, 3) aggregate, ok = aggregations[0] assert.True(t, ok) assert.Len(t, aggregate, 4) aggregate, ok = aggregations[1] assert.True(t, ok) assert.Len(t, aggregate, 2) } func TestDataTransferExceeded(t *testing.T) { user := dataprovider.User{ BaseUser: sdk.BaseUser{ TotalDataTransfer: 1, }, } transfer := dataprovider.ActiveTransfer{ CurrentULSize: 0, CurrentDLSize: 0, } user.UsedDownloadDataTransfer = 1024 * 1024 user.UsedUploadDataTransfer = 512 * 1024 checker := transfersCheckerMem{} res := checker.isDataTransferExceeded(user, transfer, 100, 100) assert.False(t, res) transfer.CurrentULSize = 1 res = checker.isDataTransferExceeded(user, transfer, 100, 100) assert.True(t, res) user.UsedDownloadDataTransfer = 512*1024 - 100 user.UsedUploadDataTransfer = 512*1024 - 100 res = checker.isDataTransferExceeded(user, transfer, 100, 100) assert.False(t, res) res = checker.isDataTransferExceeded(user, transfer, 101, 100) assert.True(t, res) user.TotalDataTransfer = 0 user.DownloadDataTransfer = 1 user.UsedDownloadDataTransfer = 512 * 1024 transfer.CurrentULSize = 0 transfer.CurrentDLSize = 100 res = checker.isDataTransferExceeded(user, transfer, 0, 512*1024) assert.False(t, res) res = checker.isDataTransferExceeded(user, transfer, 0, 512*1024+1) assert.True(t, res) user.DownloadDataTransfer = 0 user.UploadDataTransfer = 1 user.UsedUploadDataTransfer = 512 * 1024 transfer.CurrentULSize = 0 transfer.CurrentDLSize = 0 res = checker.isDataTransferExceeded(user, transfer, 512*1024+1, 0) assert.False(t, res) transfer.CurrentULSize = 1 res = checker.isDataTransferExceeded(user, transfer, 512*1024+1, 0) assert.True(t, res) } func TestGetUsersForQuotaCheck(t *testing.T) { usersToFetch := make(map[string]bool) for i := 0; i < 70; i++ { usersToFetch[fmt.Sprintf("user%v", i)] = i%2 == 0 } users, err := dataprovider.GetUsersForQuotaCheck(usersToFetch) assert.NoError(t, err) assert.Len(t, users, 0) for i := 0; i < 60; i++ { folder := vfs.BaseVirtualFolder{ Name: fmt.Sprintf("f%v", i), MappedPath: filepath.Join(os.TempDir(), fmt.Sprintf("f%v", i)), } user := dataprovider.User{ BaseUser: sdk.BaseUser{ Username: fmt.Sprintf("user%v", i), Password: "pwd", HomeDir: filepath.Join(os.TempDir(), fmt.Sprintf("user%v", i)), Status: 1, QuotaSize: 120, Permissions: map[string][]string{ "/": {dataprovider.PermAny}, }, }, VirtualFolders: []vfs.VirtualFolder{ { BaseVirtualFolder: vfs.BaseVirtualFolder{ Name: folder.Name, }, VirtualPath: "/vfolder", QuotaSize: 100, }, }, } err = dataprovider.AddFolder(&folder, "", "", "") assert.NoError(t, err) err = dataprovider.AddUser(&user, "", "", "") assert.NoError(t, err) err = dataprovider.UpdateVirtualFolderQuota(&vfs.BaseVirtualFolder{Name: fmt.Sprintf("f%v", i)}, 1, 50, false) assert.NoError(t, err) } users, err = dataprovider.GetUsersForQuotaCheck(usersToFetch) assert.NoError(t, err) assert.Len(t, users, 60) for _, user := range users { userIdxStr := strings.Replace(user.Username, "user", "", 1) userIdx, err := strconv.Atoi(userIdxStr) assert.NoError(t, err) if userIdx%2 == 0 { if assert.Len(t, user.VirtualFolders, 1, user.Username) { assert.Equal(t, int64(100), user.VirtualFolders[0].QuotaSize) assert.Equal(t, int64(50), user.VirtualFolders[0].UsedQuotaSize) } } else { switch dataprovider.GetProviderStatus().Driver { case dataprovider.MySQLDataProviderName, dataprovider.PGSQLDataProviderName, dataprovider.CockroachDataProviderName, dataprovider.SQLiteDataProviderName: assert.Len(t, user.VirtualFolders, 0, user.Username) } } ul, dl, total := user.GetDataTransferLimits() assert.Equal(t, int64(0), ul) assert.Equal(t, int64(0), dl) assert.Equal(t, int64(0), total) } for i := 0; i < 60; i++ { err = dataprovider.DeleteUser(fmt.Sprintf("user%v", i), "", "", "") assert.NoError(t, err) err = dataprovider.DeleteFolder(fmt.Sprintf("f%v", i), "", "", "") assert.NoError(t, err) } users, err = dataprovider.GetUsersForQuotaCheck(usersToFetch) assert.NoError(t, err) assert.Len(t, users, 0) } func TestDBTransferChecker(t *testing.T) { if !isDbTransferCheckerSupported() { t.Skip("this test is not supported with the current database provider") } providerConf := dataprovider.GetProviderConfig() err := dataprovider.Close() assert.NoError(t, err) providerConf.IsShared = 1 err = dataprovider.Initialize(providerConf, configDir, true) assert.NoError(t, err) c := getTransfersChecker(1) checker, ok := c.(*transfersCheckerDB) assert.True(t, ok) assert.True(t, checker.lastCleanup.IsZero()) transfer1 := dataprovider.ActiveTransfer{ ID: 1, Type: TransferDownload, ConnID: xid.New().String(), Username: "user1", FolderName: "folder1", IP: "127.0.0.1", } checker.AddTransfer(transfer1) transfers, err := dataprovider.GetActiveTransfers(time.Now().Add(24 * time.Hour)) assert.NoError(t, err) assert.Len(t, transfers, 0) transfers, err = dataprovider.GetActiveTransfers(time.Now().Add(-periodicTimeoutCheckInterval * 2)) assert.NoError(t, err) var createdAt, updatedAt int64 if assert.Len(t, transfers, 1) { transfer := transfers[0] assert.Equal(t, transfer1.ID, transfer.ID) assert.Equal(t, transfer1.Type, transfer.Type) assert.Equal(t, transfer1.ConnID, transfer.ConnID) assert.Equal(t, transfer1.Username, transfer.Username) assert.Equal(t, transfer1.IP, transfer.IP) assert.Equal(t, transfer1.FolderName, transfer.FolderName) assert.Greater(t, transfer.CreatedAt, int64(0)) assert.Greater(t, transfer.UpdatedAt, int64(0)) assert.Equal(t, int64(0), transfer.CurrentDLSize) assert.Equal(t, int64(0), transfer.CurrentULSize) createdAt = transfer.CreatedAt updatedAt = transfer.UpdatedAt } time.Sleep(100 * time.Millisecond) checker.UpdateTransferCurrentSizes(100, 150, transfer1.ID, transfer1.ConnID) transfers, err = dataprovider.GetActiveTransfers(time.Now().Add(-periodicTimeoutCheckInterval * 2)) assert.NoError(t, err) if assert.Len(t, transfers, 1) { transfer := transfers[0] assert.Equal(t, int64(150), transfer.CurrentDLSize) assert.Equal(t, int64(100), transfer.CurrentULSize) assert.Equal(t, createdAt, transfer.CreatedAt) assert.Greater(t, transfer.UpdatedAt, updatedAt) } res := checker.GetOverquotaTransfers() assert.Len(t, res, 0) checker.RemoveTransfer(transfer1.ID, transfer1.ConnID) transfers, err = dataprovider.GetActiveTransfers(time.Now().Add(-periodicTimeoutCheckInterval * 2)) assert.NoError(t, err) assert.Len(t, transfers, 0) err = dataprovider.Close() assert.NoError(t, err) res = checker.GetOverquotaTransfers() assert.Len(t, res, 0) providerConf.IsShared = 0 err = dataprovider.Initialize(providerConf, configDir, true) assert.NoError(t, err) } func isDbTransferCheckerSupported() bool { // SQLite shares the implementation with other SQL-based provider but it makes no sense // to use it outside test cases switch dataprovider.GetProviderStatus().Driver { case dataprovider.MySQLDataProviderName, dataprovider.PGSQLDataProviderName, dataprovider.CockroachDataProviderName, dataprovider.SQLiteDataProviderName: return true default: return false } }