package boltstore import ( "encoding/json" "errors" "fmt" "github.com/boltdb/bolt" "github.com/ihexxa/quickshare/src/db" ) type BoltStore struct { boltdb *bolt.DB } func NewBoltStore(boltdb *bolt.DB) (*BoltStore, error) { bs := &BoltStore{ boltdb: boltdb, } return bs, nil } func (bs *BoltStore) getUserInfo(tx *bolt.Tx, userID uint64) (*db.User, error) { var err error usersBucket := tx.Bucket([]byte(db.UsersNs)) if usersBucket == nil { return nil, db.ErrBucketNotFound } uidStr := fmt.Sprint(userID) infoBytes := usersBucket.Get([]byte(uidStr)) if infoBytes == nil { return nil, db.ErrKeyNotFound } userInfo := &db.User{} err = json.Unmarshal(infoBytes, userInfo) if err != nil { return nil, err } else if userInfo.ID != userID { return nil, fmt.Errorf("user id key(%d) info(%d) does match", userID, userInfo.ID) } return userInfo, nil } func (bs *BoltStore) setUserInfo(tx *bolt.Tx, userID uint64, userInfo *db.User) error { var err error usersBucket := tx.Bucket([]byte(db.UsersNs)) if usersBucket == nil { return db.ErrBucketNotFound } userInfoBytes, err := json.Marshal(userInfo) if err != nil { return err } uidStr := fmt.Sprint(userID) return usersBucket.Put([]byte(uidStr), userInfoBytes) } func (bs *BoltStore) getUploadInfo(tx *bolt.Tx, userID uint64, itemPath string) (*db.UploadInfo, error) { var err error uidStr := fmt.Sprint(userID) userUploadNS := db.UploadNS(uidStr) uploadInfoBucket := tx.Bucket([]byte(userUploadNS)) if uploadInfoBucket == nil { return nil, bolt.ErrBucketNotFound } uploadInfoBytes := uploadInfoBucket.Get([]byte(itemPath)) if uploadInfoBytes == nil { return nil, db.ErrKeyNotFound } uploadInfo := &db.UploadInfo{} err = json.Unmarshal(uploadInfoBytes, uploadInfo) if err != nil { return nil, err } return uploadInfo, nil } func (bs *BoltStore) setUploadInfo(tx *bolt.Tx, userID uint64, uploadPath string, uploadInfo *db.UploadInfo, overWrite bool) error { var err error uidStr := fmt.Sprint(userID) userUploadNS := db.UploadNS(uidStr) uploadInfoBucket := tx.Bucket([]byte(userUploadNS)) if uploadInfoBucket == nil { uploadInfoBucket, err = tx.CreateBucket([]byte(userUploadNS)) if err != nil { return err } } existingInfo := uploadInfoBucket.Get([]byte(uploadPath)) if existingInfo != nil && !overWrite { return db.ErrCreateExisting } uploadInfoBytes, err := json.Marshal(uploadInfo) if err != nil { return err } return uploadInfoBucket.Put([]byte(uploadPath), uploadInfoBytes) } func (bs *BoltStore) getFileInfo(tx *bolt.Tx, userID uint64, itemPath string) (*db.FileInfo, error) { var err error fileInfoBucket := tx.Bucket([]byte(db.InfoNs)) if fileInfoBucket == nil { return nil, db.ErrBucketNotFound } fileInfoBytes := fileInfoBucket.Get([]byte(itemPath)) if fileInfoBytes == nil { return nil, db.ErrKeyNotFound } fileInfo := &db.FileInfo{} err = json.Unmarshal(fileInfoBytes, fileInfo) if err != nil { return nil, err } return fileInfo, nil } func (bs *BoltStore) setFileInfo(tx *bolt.Tx, userID uint64, itemPath string, fileInfo *db.FileInfo) error { var err error fileInfoBucket := tx.Bucket([]byte(db.InfoNs)) if fileInfoBucket == nil { return db.ErrBucketNotFound } fileInfoBytes, err := json.Marshal(fileInfo) if err != nil { return err } return fileInfoBucket.Put([]byte(itemPath), fileInfoBytes) } func (bs *BoltStore) AddUploadInfos(userID uint64, tmpPath, filePath string, info *db.FileInfo) error { return bs.boltdb.Update(func(tx *bolt.Tx) error { var err error // check space quota userInfo, err := bs.getUserInfo(tx, userID) if err != nil { return err } if userInfo.UsedSpace+info.Size > int64(userInfo.Quota.SpaceLimit) { return errors.New("space limit is reached") } // update used space userInfo.UsedSpace += info.Size err = bs.setUserInfo(tx, userID, userInfo) if err != nil { return err } // add upload info uploadInfo := &db.UploadInfo{ RealFilePath: filePath, Size: info.Size, Uploaded: 0, } return bs.setUploadInfo(tx, userID, tmpPath, uploadInfo, false) }) } func (bs *BoltStore) DelUploadingInfos(userID uint64, uploadPath string) error { return bs.boltdb.Update(func(tx *bolt.Tx) error { var err error // delete upload info uidStr := fmt.Sprint(userID) userUploadNS := db.UploadNS(uidStr) uploadInfoBucket := tx.Bucket([]byte(userUploadNS)) if uploadInfoBucket == nil { return db.ErrBucketNotFound } uploadInfoBytes := uploadInfoBucket.Get([]byte(uploadPath)) if uploadInfoBytes == nil { return db.ErrKeyNotFound } uploadInfo := &db.UploadInfo{} err = json.Unmarshal(uploadInfoBytes, uploadInfo) if err != nil { return err } err = uploadInfoBucket.Delete([]byte(uploadPath)) if err != nil { return err } // decr used space userInfo, err := bs.getUserInfo(tx, userID) if err != nil { return err } userInfo.UsedSpace -= uploadInfo.Size return bs.setUserInfo(tx, userID, userInfo) }) } func (bs *BoltStore) MoveUploadingInfos(userID uint64, uploadPath, itemPath string) error { return bs.boltdb.Update(func(tx *bolt.Tx) error { var err error // delete upload info uidStr := fmt.Sprint(userID) userUploadNS := db.UploadNS(uidStr) uploadInfoBucket := tx.Bucket([]byte(userUploadNS)) if uploadInfoBucket == nil { return db.ErrBucketNotFound } uploadInfoBytes := uploadInfoBucket.Get([]byte(uploadPath)) if uploadInfoBytes == nil { return db.ErrKeyNotFound } uploadInfo := &db.UploadInfo{} err = json.Unmarshal(uploadInfoBytes, uploadInfo) if err != nil { return err } err = uploadInfoBucket.Delete([]byte(uploadPath)) if err != nil { return err } // create file info fileInfo := &db.FileInfo{ IsDir: false, Size: uploadInfo.Size, } return bs.setFileInfo(tx, userID, itemPath, fileInfo) }) } func (bs *BoltStore) delShareID(tx *bolt.Tx, itemPath string) error { var err error shareIDBucket := tx.Bucket([]byte(db.ShareIDNs)) if shareIDBucket == nil { return db.ErrBucketNotFound } shareIDtoDir := map[string]string{} shareIDBucket.ForEach(func(k, v []byte) error { shareIDtoDir[string(k)] = string(v) return nil }) // because before this version, shareIDs are not removed correctly // so it iterates all shareIDs and cleans remaining entries for shareID, shareDir := range shareIDtoDir { if shareDir == itemPath { err = shareIDBucket.Delete([]byte(shareID)) if err != nil { return err } } } return nil } func (bs *BoltStore) DelInfos(userID uint64, itemPath string) error { return bs.boltdb.Update(func(tx *bolt.Tx) error { var err error // delete file info fileInfoBucket := tx.Bucket([]byte(db.InfoNs)) if fileInfoBucket == nil { return db.ErrBucketNotFound } fileInfoBytes := fileInfoBucket.Get([]byte(itemPath)) if fileInfoBucket == nil { return db.ErrKeyNotFound } fileInfo := &db.FileInfo{} err = json.Unmarshal(fileInfoBytes, fileInfo) if err != nil { return err } err = fileInfoBucket.Delete([]byte(itemPath)) if err != nil { return err } // decr used space userInfo, err := bs.getUserInfo(tx, userID) if err != nil { return err } userInfo.UsedSpace -= fileInfo.Size err = bs.setUserInfo(tx, userID, userInfo) if err != nil { return err } // delete share id if fileInfo.IsDir { err = bs.delShareID(tx, itemPath) if err != nil { return err } } return nil }) } func (bs *BoltStore) MoveInfos(userID uint64, oldPath, newPath string, isDir bool) error { return bs.boltdb.Update(func(tx *bolt.Tx) error { var err error fileInfoBucket := tx.Bucket([]byte(db.InfoNs)) if fileInfoBucket == nil { return db.ErrBucketNotFound } fileInfo, err := bs.getFileInfo(tx, userID, oldPath) if err != nil { if errors.Is(err, db.ErrKeyNotFound) && isDir { // the file info for the dir does not exist return nil } return err } // delete old shareID if isDir { fileInfo.Shared = false fileInfo.ShareID = "" err = bs.delShareID(tx, oldPath) if err != nil { return err } } // delete old info err = fileInfoBucket.Delete([]byte(oldPath)) if err != nil { return err } fmt.Println("\n\n\n4", err) // add new info return bs.setFileInfo(tx, userID, newPath, fileInfo) }) }