refactor virtual folders

The same virtual folder can now be shared among users and different folder quota limits for each user are supported. Fixes #120
2025-12-07 14:50:55 +03:00 · 2020-06-07 23:30:18 +02:00
parent dc011af90d
commit 8306b6bde6
56 changed files with 6969 additions and 1018 deletions
--- a/dataprovider/bolt.go
+++ b/dataprovider/bolt.go
@@ -14,15 +14,17 @@ import (

 	"github.com/drakkan/sftpgo/logger"
 	"github.com/drakkan/sftpgo/utils"
+	"github.com/drakkan/sftpgo/vfs"
 )

 const (
-	boltDatabaseVersion = 3
+	boltDatabaseVersion = 4
 )

 var (
 	usersBucket      = []byte("users")
 	usersIDIdxBucket = []byte("users_id_idx")
+	foldersBucket    = []byte("folders")
 	dbVersionBucket  = []byte("db_version")
 	dbVersionKey     = []byte("version")
 )
@@ -90,6 +92,14 @@ func initializeBoltProvider(basePath string) error {
 			providerLog(logger.LevelWarn, "error creating username idx bucket: %v", err)
 			return err
 		}
+		err = dbHandle.Update(func(tx *bolt.Tx) error {
+			_, e := tx.CreateBucketIfNotExists(foldersBucket)
+			return e
+		})
+		if err != nil {
+			providerLog(logger.LevelWarn, "error creating username idx bucket: %v", err)
+			return err
+		}
 		err = dbHandle.Update(func(tx *bolt.Tx) error {
 			_, e := tx.CreateBucketIfNotExists(dbVersionBucket)
 			return e
@@ -106,7 +116,7 @@ func initializeBoltProvider(basePath string) error {
 }

 func (p BoltProvider) checkAvailability() error {
-	_, err := p.getUsers(1, 0, "ASC", "")
+	_, err := getBoltDatabaseVersion(p.dbHandle)
 	return err
 }

@@ -152,7 +162,12 @@ func (p BoltProvider) getUserByID(ID int64) (User, error) {
 		if u == nil {
 			return &RecordNotFoundError{err: fmt.Sprintf("username %#v and ID: %v does not exist", string(username), ID)}
 		}
-		return json.Unmarshal(u, &user)
+		folderBucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		user, err = joinUserAndFolders(u, folderBucket)
+		return err
 	})

 	return user, err
@@ -215,7 +230,10 @@ func (p BoltProvider) updateQuota(username string, filesAdd int, sizeAdd int64,
 		if err != nil {
 			return err
 		}
-		return bucket.Put([]byte(username), buf)
+		err = bucket.Put([]byte(username), buf)
+		providerLog(logger.LevelDebug, "quota updated for user %#v, files increment: %v size increment: %v is reset? %v",
+			username, filesAdd, sizeAdd, reset)
+		return err
 	})
 }

@@ -239,7 +257,12 @@ func (p BoltProvider) userExists(username string) (User, error) {
 		if u == nil {
 			return &RecordNotFoundError{err: fmt.Sprintf("username %v does not exist", username)}
 		}
-		return json.Unmarshal(u, &user)
+		folderBucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		user, err = joinUserAndFolders(u, folderBucket)
+		return err
 	})
 	return user, err
 }
@@ -254,6 +277,10 @@ func (p BoltProvider) addUser(user User) error {
 		if err != nil {
 			return err
 		}
+		folderBucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
 		if u := bucket.Get([]byte(user.Username)); u != nil {
 			return fmt.Errorf("username %v already exists", user.Username)
 		}
@@ -262,15 +289,21 @@ func (p BoltProvider) addUser(user User) error {
 			return err
 		}
 		user.ID = int64(id)
+		for _, folder := range user.VirtualFolders {
+			err = addUserToFolderMapping(folder, user, folderBucket)
+			if err != nil {
+				return err
+			}
+		}
 		buf, err := json.Marshal(user)
 		if err != nil {
 			return err
 		}
-		userIDAsBytes := itob(user.ID)
 		err = bucket.Put([]byte(user.Username), buf)
 		if err != nil {
 			return err
 		}
+		userIDAsBytes := itob(user.ID)
 		return idxBucket.Put(userIDAsBytes, []byte(user.Username))
 	})
 }
@@ -285,9 +318,35 @@ func (p BoltProvider) updateUser(user User) error {
 		if err != nil {
 			return err
 		}
-		if u := bucket.Get([]byte(user.Username)); u == nil {
+		folderBucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		var u []byte
+		if u = bucket.Get([]byte(user.Username)); u == nil {
 			return &RecordNotFoundError{err: fmt.Sprintf("username %v does not exist", user.Username)}
 		}
+		var oldUser User
+		err = json.Unmarshal(u, &oldUser)
+		if err != nil {
+			return err
+		}
+		for _, folder := range oldUser.VirtualFolders {
+			err = removeUserFromFolderMapping(folder, oldUser, folderBucket)
+			if err != nil {
+				return err
+			}
+		}
+		for _, folder := range user.VirtualFolders {
+			err = addUserToFolderMapping(folder, user, folderBucket)
+			if err != nil {
+				return err
+			}
+		}
+		user.LastQuotaUpdate = oldUser.LastQuotaUpdate
+		user.UsedQuotaSize = oldUser.UsedQuotaSize
+		user.UsedQuotaFiles = oldUser.UsedQuotaFiles
+		user.LastLogin = oldUser.LastLogin
 		buf, err := json.Marshal(user)
 		if err != nil {
 			return err
@@ -302,6 +361,18 @@ func (p BoltProvider) deleteUser(user User) error {
 		if err != nil {
 			return err
 		}
+		if len(user.VirtualFolders) > 0 {
+			folderBucket, err := getFolderBucket(tx)
+			if err != nil {
+				return err
+			}
+			for _, folder := range user.VirtualFolders {
+				err = removeUserFromFolderMapping(folder, user, folderBucket)
+				if err != nil {
+					return err
+				}
+			}
+		}
 		userIDAsBytes := itob(user.ID)
 		userName := idxBucket.Get(userIDAsBytes)
 		if userName == nil {
@@ -316,16 +387,19 @@ func (p BoltProvider) deleteUser(user User) error {
 }

 func (p BoltProvider) dumpUsers() ([]User, error) {
-	users := []User{}
+	users := make([]User, 0, 100)
 	err := p.dbHandle.View(func(tx *bolt.Tx) error {
 		bucket, _, err := getBuckets(tx)
 		if err != nil {
 			return err
 		}
+		folderBucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
 		cursor := bucket.Cursor()
 		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
-			var user User
-			err = json.Unmarshal(v, &user)
+			user, err := joinUserAndFolders(v, folderBucket)
 			if err != nil {
 				return err
 			}
@@ -355,7 +429,7 @@ func (p BoltProvider) getUserWithUsername(username string) ([]User, error) {
 }

 func (p BoltProvider) getUsers(limit int, offset int, order string, username string) ([]User, error) {
-	users := []User{}
+	users := make([]User, 0, limit)
 	var err error
 	if limit <= 0 {
 		return users, err
@@ -371,16 +445,19 @@ func (p BoltProvider) getUsers(limit int, offset int, order string, username str
 		if err != nil {
 			return err
 		}
+		folderBucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
 		cursor := bucket.Cursor()
 		itNum := 0
-		if order == "ASC" {
+		if order == OrderASC {
 			for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
 				itNum++
 				if itNum <= offset {
 					continue
 				}
-				var user User
-				err = json.Unmarshal(v, &user)
+				user, err := joinUserAndFolders(v, folderBucket)
 				if err == nil {
 					users = append(users, HideUserSensitiveData(&user))
 				}
@@ -394,8 +471,7 @@ func (p BoltProvider) getUsers(limit int, offset int, order string, username str
 				if itNum <= offset {
 					continue
 				}
-				var user User
-				err = json.Unmarshal(v, &user)
+				user, err := joinUserAndFolders(v, folderBucket)
 				if err == nil {
 					users = append(users, HideUserSensitiveData(&user))
 				}
@@ -409,6 +485,209 @@ func (p BoltProvider) getUsers(limit int, offset int, order string, username str
 	return users, err
 }

+func (p BoltProvider) dumpFolders() ([]vfs.BaseVirtualFolder, error) {
+	folders := make([]vfs.BaseVirtualFolder, 0, 50)
+	err := p.dbHandle.View(func(tx *bolt.Tx) error {
+		bucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		cursor := bucket.Cursor()
+		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
+			var folder vfs.BaseVirtualFolder
+			err = json.Unmarshal(v, &folder)
+			if err != nil {
+				return err
+			}
+			folders = append(folders, folder)
+		}
+		return err
+	})
+	return folders, err
+}
+
+func (p BoltProvider) getFolders(limit, offset int, order, folderPath string) ([]vfs.BaseVirtualFolder, error) {
+	folders := make([]vfs.BaseVirtualFolder, 0, limit)
+	var err error
+	if limit <= 0 {
+		return folders, err
+	}
+	if len(folderPath) > 0 {
+		if offset == 0 {
+			var folder vfs.BaseVirtualFolder
+			folder, err = p.getFolderByPath(folderPath)
+			if err == nil {
+				folders = append(folders, folder)
+			}
+		}
+		return folders, err
+	}
+	err = p.dbHandle.View(func(tx *bolt.Tx) error {
+		bucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		cursor := bucket.Cursor()
+		itNum := 0
+		if order == OrderASC {
+			for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
+				itNum++
+				if itNum <= offset {
+					continue
+				}
+				var folder vfs.BaseVirtualFolder
+				err = json.Unmarshal(v, &folder)
+				if err != nil {
+					return err
+				}
+				folders = append(folders, folder)
+				if len(folders) >= limit {
+					break
+				}
+			}
+		} else {
+			for k, v := cursor.Last(); k != nil; k, v = cursor.Prev() {
+				itNum++
+				if itNum <= offset {
+					continue
+				}
+				var folder vfs.BaseVirtualFolder
+				err = json.Unmarshal(v, &folder)
+				if err != nil {
+					return err
+				}
+				folders = append(folders, folder)
+				if len(folders) >= limit {
+					break
+				}
+			}
+		}
+		return err
+	})
+	return folders, err
+}
+
+func (p BoltProvider) getFolderByPath(name string) (vfs.BaseVirtualFolder, error) {
+	var folder vfs.BaseVirtualFolder
+	err := p.dbHandle.View(func(tx *bolt.Tx) error {
+		bucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		folder, err = folderExistsInternal(name, bucket)
+		return err
+	})
+	return folder, err
+}
+
+func (p BoltProvider) addFolder(folder vfs.BaseVirtualFolder) error {
+	err := validateFolder(&folder)
+	if err != nil {
+		return err
+	}
+	return p.dbHandle.Update(func(tx *bolt.Tx) error {
+		bucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		if f := bucket.Get([]byte(folder.MappedPath)); f != nil {
+			return fmt.Errorf("folder %v already exists", folder.MappedPath)
+		}
+		_, err = addFolderInternal(folder, bucket)
+		return err
+	})
+}
+
+func (p BoltProvider) deleteFolder(folder vfs.BaseVirtualFolder) error {
+	return p.dbHandle.Update(func(tx *bolt.Tx) error {
+		bucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		usersBucket, _, err := getBuckets(tx)
+		if err != nil {
+			return err
+		}
+		var f []byte
+		if f = bucket.Get([]byte(folder.MappedPath)); f == nil {
+			return &RecordNotFoundError{err: fmt.Sprintf("folder %v does not exist", folder.MappedPath)}
+		}
+		var folder vfs.BaseVirtualFolder
+		err = json.Unmarshal(f, &folder)
+		if err != nil {
+			return err
+		}
+		for _, username := range folder.Users {
+			var u []byte
+			if u = usersBucket.Get([]byte(username)); u == nil {
+				continue
+			}
+			var user User
+			err = json.Unmarshal(u, &user)
+			if err != nil {
+				return err
+			}
+			var folders []vfs.VirtualFolder
+			for _, userFolder := range user.VirtualFolders {
+				if folder.MappedPath != userFolder.MappedPath {
+					folders = append(folders, userFolder)
+				}
+			}
+			user.VirtualFolders = folders
+			buf, err := json.Marshal(user)
+			if err != nil {
+				return err
+			}
+			err = usersBucket.Put([]byte(user.Username), buf)
+			if err != nil {
+				return err
+			}
+		}
+
+		return bucket.Delete([]byte(folder.MappedPath))
+	})
+}
+
+func (p BoltProvider) updateFolderQuota(mappedPath string, filesAdd int, sizeAdd int64, reset bool) error {
+	return p.dbHandle.Update(func(tx *bolt.Tx) error {
+		bucket, err := getFolderBucket(tx)
+		if err != nil {
+			return err
+		}
+		var f []byte
+		if f = bucket.Get([]byte(mappedPath)); f == nil {
+			return &RecordNotFoundError{err: fmt.Sprintf("folder %v does not exist, unable to update quota", mappedPath)}
+		}
+		var folder vfs.BaseVirtualFolder
+		err = json.Unmarshal(f, &folder)
+		if err != nil {
+			return err
+		}
+		if reset {
+			folder.UsedQuotaSize = sizeAdd
+			folder.UsedQuotaFiles = filesAdd
+		} else {
+			folder.UsedQuotaSize += sizeAdd
+			folder.UsedQuotaFiles += filesAdd
+		}
+		folder.LastQuotaUpdate = utils.GetTimeAsMsSinceEpoch(time.Now())
+		buf, err := json.Marshal(folder)
+		if err != nil {
+			return err
+		}
+		return bucket.Put([]byte(folder.MappedPath), buf)
+	})
+}
+
+func (p BoltProvider) getUsedFolderQuota(mappedPath string) (int, int64, error) {
+	folder, err := p.getFolderByPath(mappedPath)
+	if err != nil {
+		providerLog(logger.LevelWarn, "unable to get quota for folder %#v error: %v", mappedPath, err)
+		return 0, 0, err
+	}
+	return folder.UsedQuotaFiles, folder.UsedQuotaSize, err
+}
+
 func (p BoltProvider) close() error {
 	return p.dbHandle.Close()
 }
@@ -437,9 +716,19 @@ func (p BoltProvider) migrateDatabase() error {
 		if err != nil {
 			return err
 		}
-		return updateDatabaseFrom2To3(p.dbHandle)
+		err = updateDatabaseFrom2To3(p.dbHandle)
+		if err != nil {
+			return err
+		}
+		return updateDatabaseFrom3To4(p.dbHandle)
 	case 2:
-		return updateDatabaseFrom2To3(p.dbHandle)
+		err = updateDatabaseFrom2To3(p.dbHandle)
+		if err != nil {
+			return err
+		}
+		return updateDatabaseFrom3To4(p.dbHandle)
+	case 3:
+		return updateDatabaseFrom3To4(p.dbHandle)
 	default:
 		return fmt.Errorf("Database version not handled: %v", dbVersion.Version)
 	}
@@ -452,6 +741,106 @@ func itob(v int64) []byte {
 	return b
 }

+func joinUserAndFolders(u []byte, foldersBucket *bolt.Bucket) (User, error) {
+	var user User
+	err := json.Unmarshal(u, &user)
+	if len(user.VirtualFolders) > 0 {
+		var folders []vfs.VirtualFolder
+		for _, folder := range user.VirtualFolders {
+			baseFolder, err := folderExistsInternal(folder.MappedPath, foldersBucket)
+			if err != nil {
+				continue
+			}
+			folder.UsedQuotaFiles = baseFolder.UsedQuotaFiles
+			folder.UsedQuotaSize = baseFolder.UsedQuotaSize
+			folder.LastQuotaUpdate = baseFolder.LastQuotaUpdate
+			folder.ID = baseFolder.ID
+			folders = append(folders, folder)
+		}
+		user.VirtualFolders = folders
+	}
+	return user, err
+}
+
+func folderExistsInternal(name string, bucket *bolt.Bucket) (vfs.BaseVirtualFolder, error) {
+	var folder vfs.BaseVirtualFolder
+	f := bucket.Get([]byte(name))
+	if f == nil {
+		err := &RecordNotFoundError{err: fmt.Sprintf("folder %v does not exist", name)}
+		return folder, err
+	}
+	err := json.Unmarshal(f, &folder)
+	return folder, err
+}
+
+func addFolderInternal(folder vfs.BaseVirtualFolder, bucket *bolt.Bucket) (vfs.BaseVirtualFolder, error) {
+	id, err := bucket.NextSequence()
+	if err != nil {
+		return folder, err
+	}
+	folder.ID = int64(id)
+	buf, err := json.Marshal(folder)
+	if err != nil {
+		return folder, err
+	}
+	err = bucket.Put([]byte(folder.MappedPath), buf)
+	return folder, err
+}
+
+func addUserToFolderMapping(folder vfs.VirtualFolder, user User, bucket *bolt.Bucket) error {
+	var baseFolder vfs.BaseVirtualFolder
+	var err error
+	if f := bucket.Get([]byte(folder.MappedPath)); f == nil {
+		// folder does not exists, try to create
+		baseFolder, err = addFolderInternal(folder.BaseVirtualFolder, bucket)
+	} else {
+		err = json.Unmarshal(f, &baseFolder)
+	}
+	if err != nil {
+		return err
+	}
+	if !utils.IsStringInSlice(user.Username, baseFolder.Users) {
+		baseFolder.Users = append(baseFolder.Users, user.Username)
+		buf, err := json.Marshal(baseFolder)
+		if err != nil {
+			return err
+		}
+		err = bucket.Put([]byte(folder.MappedPath), buf)
+		if err != nil {
+			return err
+		}
+	}
+	return err
+}
+
+func removeUserFromFolderMapping(folder vfs.VirtualFolder, user User, bucket *bolt.Bucket) error {
+	var f []byte
+	if f = bucket.Get([]byte(folder.MappedPath)); f == nil {
+		// the folder does not exists so there is no associated user
+		return nil
+	}
+	var baseFolder vfs.BaseVirtualFolder
+	err := json.Unmarshal(f, &baseFolder)
+	if err != nil {
+		return err
+	}
+	if utils.IsStringInSlice(user.Username, baseFolder.Users) {
+		var newUserMapping []string
+		for _, u := range baseFolder.Users {
+			if u != user.Username {
+				newUserMapping = append(newUserMapping, u)
+			}
+		}
+		baseFolder.Users = newUserMapping
+		buf, err := json.Marshal(baseFolder)
+		if err != nil {
+			return err
+		}
+		return bucket.Put([]byte(folder.MappedPath), buf)
+	}
+	return err
+}
+
 func getBuckets(tx *bolt.Tx) (*bolt.Bucket, *bolt.Bucket, error) {
 	var err error
 	bucket := tx.Bucket(usersBucket)
@@ -462,6 +851,15 @@ func getBuckets(tx *bolt.Tx) (*bolt.Bucket, *bolt.Bucket, error) {
 	return bucket, idxBucket, err
 }

+func getFolderBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
+	var err error
+	bucket := tx.Bucket(foldersBucket)
+	if bucket == nil {
+		err = fmt.Errorf("unable to find required buckets, bolt database structure not correcly defined")
+	}
+	return bucket, err
+}
+
 func updateDatabaseFrom1To2(dbHandle *bolt.DB) error {
 	providerLog(logger.LevelInfo, "updating bolt database version: 1 -> 2")
 	usernames, err := getBoltAvailableUsernames(dbHandle)
@@ -537,6 +935,69 @@ func updateDatabaseFrom2To3(dbHandle *bolt.DB) error {
 	return updateBoltDatabaseVersion(dbHandle, 3)
 }

+func updateDatabaseFrom3To4(dbHandle *bolt.DB) error {
+	providerLog(logger.LevelInfo, "updating bolt database version: 3 -> 4")
+	foldersToScan := []string{}
+	users := []userCompactVFolders{}
+	err := dbHandle.View(func(tx *bolt.Tx) error {
+		bucket, _, err := getBuckets(tx)
+		if err != nil {
+			return err
+		}
+		cursor := bucket.Cursor()
+		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
+			var compatUser userCompactVFolders
+			err = json.Unmarshal(v, &compatUser)
+			if err == nil && len(compatUser.VirtualFolders) > 0 {
+				users = append(users, compatUser)
+			}
+		}
+		return err
+	})
+	if err != nil {
+		return err
+	}
+	for _, u := range users {
+		user, err := provider.userExists(u.Username)
+		if err != nil {
+			return err
+		}
+		var folders []vfs.VirtualFolder
+		for _, f := range u.VirtualFolders {
+			providerLog(logger.LevelInfo, "restoring virtual folder: %+v for user %#v", f, user.Username)
+			quotaSize := int64(-1)
+			quotaFiles := -1
+			if f.ExcludeFromQuota {
+				quotaSize = 0
+				quotaFiles = 0
+			}
+			folder := vfs.VirtualFolder{
+				QuotaSize:   quotaSize,
+				QuotaFiles:  quotaFiles,
+				VirtualPath: f.VirtualPath,
+			}
+			folder.MappedPath = f.MappedPath
+			folders = append(folders, folder)
+			if !utils.IsStringInSlice(folder.MappedPath, foldersToScan) {
+				foldersToScan = append(foldersToScan, folder.MappedPath)
+			}
+		}
+		user.VirtualFolders = folders
+		err = provider.updateUser(user)
+		providerLog(logger.LevelInfo, "number of virtual folders to restore %v, user %#v, error: %v", len(user.VirtualFolders),
+			user.Username, err)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = updateBoltDatabaseVersion(dbHandle, 4)
+	if err == nil {
+		go updateVFoldersQuotaAfterRestore(foldersToScan)
+	}
+	return err
+}
+
 func getBoltAvailableUsernames(dbHandle *bolt.DB) ([]string, error) {
 	usernames := []string{}
 	err := dbHandle.View(func(tx *bolt.Tx) error {