improved session creation for bots and other changes

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**

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);

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

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"`
 }

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)

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)

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])
 

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}

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 {

ui/teldrive-ui

@@ -1 +1 @@
-Subproject commit a2d12d64157b72ef3d7daa359db0f6a07917d5ec
+Subproject commit 46d5af0386597fa793128c4d933aafcdbdb9d63a

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
 }