From 0a900e3f1f6b27f2c8fcf541372569e27d54646b Mon Sep 17 00:00:00 2001 From: divyam234 Date: Wed, 27 Sep 2023 17:08:33 +0530 Subject: [PATCH] improved session creation for bots and other changes --- README.md | 7 ++++- .../migrations/20230927152733_alter_bots.sql | 7 +++++ go.mod | 4 +-- models/bot.model.go | 1 + routes/user.go | 12 +++++++++ services/common.go | 17 ++++++++---- services/file.service.go | 4 +-- services/upload.service.go | 2 +- services/user.service.go | 27 ++++++++++++++++--- ui/teldrive-ui | 2 +- utils/config.go | 1 + 11 files changed, 68 insertions(+), 16 deletions(-) create mode 100644 database/migrations/20230927152733_alter_bots.sql diff --git a/README.md b/README.md index 88326cb..8015f83 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,7 @@ COOKIE_SAME_SITE=true JWT_SECRET=abc DATABASE_URL=abc RATE_LIMIT=true +LAZY_STREAM_BOTS=false ``` > **Warning** @@ -118,7 +119,11 @@ In addition to the mandatory variables, you can also set the following optional - `HTTPS` : Only needed when frontend is on other domain. - `PORT` : Change listen port default is 8080 - `ALLOWED_USERS` : Allow certain telegram usernames including yours to access the app.Enter comma seperated telegram usernames here.Its needed when your instance is on public cloud and you want to restrict other people to access you app. -- `COOKIE_SAME_SITE` : Only needed when frontend is deployed on vercel. +- `COOKIE_SAME_SITE` : Only needed when frontend is on other domain. + +- `LAZY_STREAM_BOTS` : If set to true start Bot session and close immediately when stream or download request is over otherwise run bots forever till server stops. + +- `BG_BOTS_LIMIT` : If LAZY_STREAM_BOTS is set to false it start atmost BG_BOTS_LIMIT no of bots in background to prevent connection recreation on every request(Default is 10). ### For making use of Multi Bots support > **Warning** diff --git a/database/migrations/20230927152733_alter_bots.sql b/database/migrations/20230927152733_alter_bots.sql new file mode 100644 index 0000000..25f5bf9 --- /dev/null +++ b/database/migrations/20230927152733_alter_bots.sql @@ -0,0 +1,7 @@ +-- +goose Up + +ALTER TABLE teldrive.bots ADD COLUMN channel_id BIGINT; + +ALTER TABLE teldrive.bots DROP CONSTRAINT IF EXISTS btoken_user_un; + +ALTER TABLE teldrive.bots ADD CONSTRAINT btoken_user_channel_un UNIQUE (user_id, token,channel_id); diff --git a/go.mod b/go.mod index 8d85225..b42eb76 100644 --- a/go.mod +++ b/go.mod @@ -53,10 +53,8 @@ require ( github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/json-iterator/go v1.1.12 // indirect - github.com/klauspost/compress v1.16.7 // indirect - github.com/klauspost/cpuid/v2 v2.2.5 // indirect github.com/klauspost/compress v1.17.0 // indirect - github.com/klauspost/cpuid/v2 v2.2.4 // indirect + github.com/klauspost/cpuid/v2 v2.2.5 // indirect github.com/leodido/go-urn v1.2.4 // indirect github.com/mattn/go-isatty v0.0.19 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect diff --git a/models/bot.model.go b/models/bot.model.go index ceb4764..b8a1a76 100644 --- a/models/bot.model.go +++ b/models/bot.model.go @@ -5,4 +5,5 @@ type Bot struct { UserID int64 `gorm:"type:bigint"` BotID int64 `gorm:"type:bigint"` BotUserName string `gorm:"type:text"` + ChannelID int64 `gorm:"type:bigint"` } diff --git a/routes/user.go b/routes/user.go index 4d02a6e..719ce5c 100644 --- a/routes/user.go +++ b/routes/user.go @@ -73,6 +73,18 @@ func addUserRoutes(rg *gin.RouterGroup) { c.JSON(http.StatusCreated, res) }) + r.DELETE("/bots", func(c *gin.Context) { + + res, err := userService.RemoveBots(c) + + if err != nil { + c.AbortWithError(err.Code, err.Error) + return + } + + c.JSON(http.StatusOK, res) + }) + r.GET("/bots/revoke", func(c *gin.Context) { res, err := userService.RevokeBotSession(c) diff --git a/services/common.go b/services/common.go index ff6e288..f8a400e 100644 --- a/services/common.go +++ b/services/common.go @@ -181,15 +181,22 @@ func GetDefaultChannel(ctx context.Context, userID int64) (int64, error) { return channelID, nil } -func GetBotsToken(userID int64) ([]string, error) { +func GetBotsToken(ctx context.Context, userID int64) ([]string, error) { var bots []string - key := kv.Key("users", strconv.FormatInt(userID, 10), "bots") - - err := kv.GetValue(database.KV, key, &bots) + channelId, err := GetDefaultChannel(ctx, userID) if err != nil { - if err := database.DB.Model(&models.Bot{}).Where("user_id = ?", userID).Pluck("token", &bots).Error; err != nil { + return nil, err + } + + key := kv.Key("users", strconv.FormatInt(userID, 10), strconv.FormatInt(channelId, 10), "bots") + + err = kv.GetValue(database.KV, key, &bots) + + if err != nil { + if err := database.DB.Model(&models.Bot{}).Where("user_id = ?", userID). + Where("channel_id = ?", channelId).Pluck("token", &bots).Error; err != nil { return nil, err } kv.SetValue(database.KV, key, &bots) diff --git a/services/file.service.go b/services/file.service.go index b853c51..235d81b 100644 --- a/services/file.service.go +++ b/services/file.service.go @@ -396,7 +396,7 @@ func (fs *FileService) GetFileStream(c *gin.Context) { userID, _ := strconv.ParseInt(jwtUser.Subject, 10, 64) - tokens, err := GetBotsToken(userID) + tokens, err := GetBotsToken(c, userID) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) @@ -434,7 +434,7 @@ func (fs *FileService) GetFileStream(c *gin.Context) { } } else { - limit := utils.Min(len(tokens), 10) + limit := utils.Min(len(tokens), config.BgBotsLimit) tgc.StreamWorkers.Set(tokens[:limit]) diff --git a/services/upload.service.go b/services/upload.service.go index 018a92f..9d7cbdb 100644 --- a/services/upload.service.go +++ b/services/upload.service.go @@ -89,7 +89,7 @@ func (us *UploadService) UploadFile(c *gin.Context) (*schemas.UploadPartOut, *ty userId, session := getUserAuth(c) - tokens, err := GetBotsToken(userId) + tokens, err := GetBotsToken(c, userId) if err != nil { return nil, &types.AppError{Error: errors.New("failed to fetch bots"), Code: http.StatusInternalServerError} diff --git a/services/user.service.go b/services/user.service.go index 024e440..15548ac 100644 --- a/services/user.service.go +++ b/services/user.service.go @@ -91,7 +91,7 @@ func (us *UserService) Stats(c *gin.Context) (*schemas.AccountStats, *types.AppE func (us *UserService) GetBots(c *gin.Context) ([]string, *types.AppError) { userID, _ := getUserAuth(c) - tokens, err := GetBotsToken(userID) + tokens, err := GetBotsToken(c, userID) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} @@ -178,6 +178,27 @@ func (us *UserService) AddBots(c *gin.Context) (*schemas.Message, *types.AppErro } +func (us *UserService) RemoveBots(c *gin.Context) (*schemas.Message, *types.AppError) { + userID, _ := getUserAuth(c) + + channelId, err := GetDefaultChannel(c, userID) + + if err != nil { + return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} + } + + if err := us.Db.Where("user_id = ?", userID).Where("channel_id = ?", channelId). + Delete(&models.Bot{}).Error; err != nil { + return nil, &types.AppError{Error: errors.New("failed to delete bots"), Code: http.StatusInternalServerError} + } + key := kv.Key("users", strconv.FormatInt(userID, 10), strconv.FormatInt(channelId, 10), "bots") + + database.KV.Delete(key) + + return &schemas.Message{Status: true, Message: "bots deleted"}, nil + +} + func (us *UserService) RevokeBotSession(c *gin.Context) (*schemas.Message, *types.AppError) { pattern := []byte("botsession:") @@ -334,11 +355,11 @@ func (us *UserService) addBots(c context.Context, client *telegram.Client, userI for _, info := range botInfo { payload = append(payload, models.Bot{UserID: userId, Token: info.Token, BotID: info.Id, - BotUserName: info.UserName, + BotUserName: info.UserName, ChannelID: channelId, }) } - key := kv.Key("users", strconv.FormatInt(userId, 10), "bots") + key := kv.Key("users", strconv.FormatInt(userId, 10), strconv.FormatInt(channelId, 10), "bots") database.KV.Delete(key) if err := us.Db.Clauses(clause.OnConflict{DoNothing: true}).Create(&payload).Error; err != nil { diff --git a/ui/teldrive-ui b/ui/teldrive-ui index a2d12d6..46d5af0 160000 --- a/ui/teldrive-ui +++ b/ui/teldrive-ui @@ -1 +1 @@ -Subproject commit a2d12d64157b72ef3d7daa359db0f6a07917d5ec +Subproject commit 46d5af0386597fa793128c4d933aafcdbdb9d63a diff --git a/utils/config.go b/utils/config.go index 2e15a58..250c235 100644 --- a/utils/config.go +++ b/utils/config.go @@ -29,6 +29,7 @@ type Config struct { RunMigrations bool `envconfig:"RUN_MIGRATIONS" default:"true"` Port int `envconfig:"PORT" default:"8080"` LazyStreamBots bool `envconfig:"LAZY_STREAM_BOTS" default:"true"` + BgBotsLimit int `envconfig:"BG_BOTS_LIMIT" default:"10"` ExecDir string }