449 lines
14 KiB
Go
449 lines
14 KiB
Go
package common
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"path"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/rs/xid"
|
|
"github.com/sftpgo/sdk"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/drakkan/sftpgo/v2/dataprovider"
|
|
"github.com/drakkan/sftpgo/v2/util"
|
|
"github.com/drakkan/sftpgo/v2/vfs"
|
|
)
|
|
|
|
func TestTransfersCheckerDiskQuota(t *testing.T) {
|
|
username := "transfers_check_username"
|
|
folderName := "test_transfers_folder"
|
|
vdirPath := "/vdir"
|
|
user := dataprovider.User{
|
|
BaseUser: sdk.BaseUser{
|
|
Username: username,
|
|
Password: "testpwd",
|
|
HomeDir: filepath.Join(os.TempDir(), username),
|
|
Status: 1,
|
|
QuotaSize: 120,
|
|
Permissions: map[string][]string{
|
|
"/": {dataprovider.PermAny},
|
|
},
|
|
},
|
|
VirtualFolders: []vfs.VirtualFolder{
|
|
{
|
|
BaseVirtualFolder: vfs.BaseVirtualFolder{
|
|
Name: folderName,
|
|
MappedPath: filepath.Join(os.TempDir(), folderName),
|
|
},
|
|
VirtualPath: vdirPath,
|
|
QuotaSize: 100,
|
|
},
|
|
},
|
|
}
|
|
|
|
err := dataprovider.AddUser(&user, "", "")
|
|
assert.NoError(t, err)
|
|
user, err = dataprovider.UserExists(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)
|
|
transfer1.BytesReceived = 150
|
|
Connections.Add(fakeConn1)
|
|
// 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)
|
|
transfer1.BytesReceived = 50
|
|
transfer2.BytesReceived = 60
|
|
Connections.Add(fakeConn2)
|
|
|
|
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)
|
|
transfer3.BytesReceived = 60 // this value will be ignored, this is a download
|
|
Connections.Add(fakeConn3)
|
|
|
|
// the transfers are not overquota
|
|
Connections.checkTransfers()
|
|
assert.Nil(t, transfer1.errAbort)
|
|
assert.Nil(t, transfer2.errAbort)
|
|
assert.Nil(t, transfer3.errAbort)
|
|
|
|
transfer1.BytesReceived = 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 = 120
|
|
// we are now overquota
|
|
// if another check is in progress nothing is done
|
|
atomic.StoreInt32(&Connections.transfersCheckStatus, 1)
|
|
Connections.checkTransfers()
|
|
assert.Nil(t, transfer1.errAbort)
|
|
assert.Nil(t, transfer2.errAbort)
|
|
assert.Nil(t, transfer3.errAbort)
|
|
atomic.StoreInt32(&Connections.transfersCheckStatus, 0)
|
|
|
|
Connections.checkTransfers()
|
|
assert.True(t, conn1.IsQuotaExceededError(transfer1.errAbort))
|
|
assert.True(t, conn2.IsQuotaExceededError(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
|
|
user.QuotaSize = 1000
|
|
err = dataprovider.UpdateUser(&user, "", "")
|
|
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)
|
|
|
|
user.QuotaSize = 0
|
|
err = dataprovider.UpdateUser(&user, "", "")
|
|
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 = 0
|
|
transfer2.BytesReceived = 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)
|
|
Connections.Add(fakeConn4)
|
|
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)
|
|
|
|
Connections.Add(fakeConn5)
|
|
transfer4.BytesReceived = 50
|
|
transfer5.BytesReceived = 40
|
|
Connections.checkTransfers()
|
|
assert.Nil(t, transfer4.errAbort)
|
|
assert.Nil(t, transfer5.errAbort)
|
|
transfer5.BytesReceived = 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)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
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.aggregateTransfers()
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 0)
|
|
assert.Len(t, aggregations, 2)
|
|
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 0)
|
|
assert.Len(t, aggregations, 3)
|
|
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 0)
|
|
assert.Len(t, aggregations, 4)
|
|
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 1)
|
|
val, ok := usersToFetch["user"]
|
|
assert.True(t, ok)
|
|
assert.False(t, val)
|
|
assert.Len(t, aggregations, 4)
|
|
aggregate, ok := aggregations["user0"]
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 1)
|
|
val, ok = usersToFetch["user"]
|
|
assert.True(t, ok)
|
|
assert.False(t, val)
|
|
assert.Len(t, aggregations, 4)
|
|
aggregate, ok = aggregations["user0"]
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 1)
|
|
val, ok = usersToFetch["user"]
|
|
assert.True(t, ok)
|
|
assert.True(t, val)
|
|
assert.Len(t, aggregations, 4)
|
|
aggregate, ok = aggregations["user0"]
|
|
assert.True(t, ok)
|
|
assert.Len(t, aggregate, 3)
|
|
aggregate, ok = aggregations["userfolder0"]
|
|
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.aggregateTransfers()
|
|
assert.Len(t, usersToFetch, 1)
|
|
val, ok = usersToFetch["user"]
|
|
assert.True(t, ok)
|
|
assert.True(t, val)
|
|
assert.Len(t, aggregations, 4)
|
|
aggregate, ok = aggregations["user0"]
|
|
assert.True(t, ok)
|
|
assert.Len(t, aggregate, 4)
|
|
aggregate, ok = aggregations["userfolder0"]
|
|
assert.True(t, ok)
|
|
assert.Len(t, aggregate, 2)
|
|
}
|
|
|
|
func TestGetUsersForQuotaCheck(t *testing.T) {
|
|
usersToFetch := make(map[string]bool)
|
|
for i := 0; i < 50; 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 < 40; 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: fmt.Sprintf("f%v", i),
|
|
MappedPath: filepath.Join(os.TempDir(), fmt.Sprintf("f%v", i)),
|
|
},
|
|
VirtualPath: "/vfolder",
|
|
QuotaSize: 100,
|
|
},
|
|
},
|
|
}
|
|
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, 40)
|
|
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
|
|
for i := 0; i < 40; 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)
|
|
}
|