package services import ( "context" "encoding/base64" "fmt" "io" "mime" "net/http" "strconv" "strings" "github.com/divyam234/teldrive/internal/cache" "github.com/divyam234/teldrive/internal/category" "github.com/divyam234/teldrive/internal/config" "github.com/divyam234/teldrive/internal/database" "github.com/divyam234/teldrive/internal/http_range" "github.com/divyam234/teldrive/internal/md5" "github.com/divyam234/teldrive/internal/reader" "github.com/divyam234/teldrive/internal/tgc" "github.com/divyam234/teldrive/internal/utils" "github.com/divyam234/teldrive/pkg/logging" "github.com/divyam234/teldrive/pkg/mapper" "github.com/divyam234/teldrive/pkg/models" "github.com/divyam234/teldrive/pkg/schemas" "github.com/divyam234/teldrive/pkg/types" "github.com/gin-gonic/gin" "github.com/gotd/td/tg" "go.uber.org/zap" "github.com/jackc/pgx/v5/pgtype" "gorm.io/gorm" "gorm.io/gorm/clause" ) type FileService struct { db *gorm.DB cnf *config.TGConfig worker *tgc.StreamWorker } func NewFileService(db *gorm.DB, cnf *config.Config, worker *tgc.StreamWorker) *FileService { return &FileService{db: db, cnf: &cnf.TG, worker: worker} } func (fs *FileService) CreateFile(c *gin.Context, userId int64, fileIn *schemas.FileIn) (*schemas.FileOut, *types.AppError) { var fileDB models.File fileIn.Path = strings.TrimSpace(fileIn.Path) if fileIn.Path != "" { pathId, err := fs.getPathId(fileIn.Path, userId) if err != nil || pathId == "" { return nil, &types.AppError{Error: err, Code: http.StatusNotFound} } fileDB.ParentID = pathId } if fileIn.Type == "folder" { fileDB.MimeType = "drive/folder" var fullPath string if fileIn.Path == "/" { fullPath = "/" + fileIn.Name } else { fullPath = fileIn.Path + "/" + fileIn.Name } fileDB.Path = fullPath fileDB.Depth = utils.IntPointer(len(strings.Split(fileIn.Path, "/")) - 1) } else if fileIn.Type == "file" { fileDB.Path = "" channelId := fileIn.ChannelID if fileIn.ChannelID == 0 { var err error channelId, err = GetDefaultChannel(c, fs.db, userId) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusNotFound} } } fileDB.ChannelID = &channelId fileDB.MimeType = fileIn.MimeType fileDB.Category = string(category.GetCategory(fileIn.Name)) parts := models.Parts{} for _, part := range fileIn.Parts { parts = append(parts, models.Part{ ID: part.ID, Salt: part.Salt, }) } fileDB.Parts = &parts fileDB.Starred = false fileDB.Size = &fileIn.Size } fileDB.Name = fileIn.Name fileDB.Type = fileIn.Type fileDB.UserID = userId fileDB.Status = "active" fileDB.Encrypted = fileIn.Encrypted if err := fs.db.Create(&fileDB).Error; err != nil { if database.IsKeyConflictErr(err) { return nil, &types.AppError{Error: database.ErrKeyConflict, Code: http.StatusConflict} } return nil, &types.AppError{Error: err} } res := mapper.ToFileOut(fileDB) return res, nil } func (fs *FileService) UpdateFile(id string, userId int64, update *schemas.FileUpdate) (*schemas.FileOut, *types.AppError) { var ( files []models.File chain *gorm.DB ) if update.Type == "folder" && update.Name != "" { chain = fs.db.Raw("select * from teldrive.update_folder(?, ?, ?)", id, update.Name, userId).Scan(&files) } else { chain = fs.db.Model(&files).Clauses(clause.Returning{}).Where("id = ?", id).Updates(update) } if chain.Error != nil { return nil, &types.AppError{Error: chain.Error} } if chain.RowsAffected == 0 { return nil, &types.AppError{Error: database.ErrNotFound, Code: http.StatusNotFound} } return mapper.ToFileOut(files[0]), nil } func (fs *FileService) GetFileByID(id string) (*schemas.FileOutFull, *types.AppError) { var file models.File if err := fs.db.Where("id = ?", id).First(&file).Error; err != nil { if database.IsRecordNotFoundErr(err) { return nil, &types.AppError{Error: database.ErrNotFound, Code: http.StatusNotFound} } return nil, &types.AppError{Error: err} } return mapper.ToFileOutFull(file), nil } func (fs *FileService) ListFiles(userId int64, fquery *schemas.FileQuery) (*schemas.FileResponse, *types.AppError) { var ( pathId string err error ) if fquery.Path != "" { pathId, err = fs.getPathId(fquery.Path, userId) if err != nil { return nil, &types.AppError{Error: err, Code: http.StatusNotFound} } } query := fs.db.Limit(fquery.PerPage) filter := &models.File{UserID: userId, Status: "active"} setOrderFilter(query, fquery) if fquery.Op == "list" { query.Order("type DESC").Order(getOrder(fquery)).Where("parent_id = ?", pathId). Model(filter).Where(&filter) } else if fquery.Op == "find" { filter.Name = fquery.Name filter.Type = fquery.Type filter.ParentID = fquery.ParentID filter.Category = fquery.Category filter.Path = fquery.Path filter.Type = fquery.Type if fquery.Starred != nil { filter.Starred = *fquery.Starred } if fquery.Path != "" && fquery.Name != "" { filter.ParentID = pathId filter.Path = "" } query.Order("type DESC").Order(getOrder(fquery)). Model(&filter).Where(&filter) } else if fquery.Op == "search" { query.Where("teldrive.get_tsquery(?) @@ teldrive.get_tsvector(name)", fquery.Search) query.Order(getOrder(fquery)). Model(&filter).Where(&filter) } if fquery.Path == "" { query.Select("*,(select path from teldrive.files as f where f.id = files.parent_id) as parent_path") } files := []schemas.FileOut{} query.Scan(&files) token := "" if len(files) == fquery.PerPage { lastItem := files[len(files)-1] token = utils.GetField(&lastItem, utils.CamelToPascalCase(fquery.Sort)) token = base64.StdEncoding.EncodeToString([]byte(token)) } res := &schemas.FileResponse{Files: files, NextPageToken: token} return res, nil } func (fs *FileService) getPathId(path string, userId int64) (string, error) { var file models.File if err := fs.db.Model(&models.File{}).Select("id").Where("path = ?", path).Where("user_id = ?", userId). First(&file).Error; database.IsRecordNotFoundErr(err) { return "", database.ErrNotFound } return file.ID, nil } func (fs *FileService) MakeDirectory(userId int64, payload *schemas.MkDir) (*schemas.FileOut, *types.AppError) { var files []models.File if err := fs.db.Raw("select * from teldrive.create_directories(?, ?)", userId, payload.Path). Scan(&files).Error; err != nil { return nil, &types.AppError{Error: err} } file := mapper.ToFileOut(files[0]) return file, nil } func (fs *FileService) MoveFiles(userId int64, payload *schemas.FileOperation) (*schemas.Message, *types.AppError) { items := pgtype.Array[string]{ Elements: payload.Files, Valid: true, Dims: []pgtype.ArrayDimension{{Length: int32(len(payload.Files)), LowerBound: 1}}, } if err := fs.db.Exec("select * from teldrive.move_items(? , ? , ?)", items, payload.Destination, userId).Error; err != nil { return nil, &types.AppError{Error: err} } return &schemas.Message{Message: "files moved"}, nil } func (fs *FileService) DeleteFiles(userId int64, payload *schemas.FileOperation) (*schemas.Message, *types.AppError) { if err := fs.db.Exec("call teldrive.delete_files($1)", payload.Files).Error; err != nil { return nil, &types.AppError{Error: err} } return &schemas.Message{Message: "files deleted"}, nil } func (fs *FileService) MoveDirectory(userId int64, payload *schemas.DirMove) (*schemas.Message, *types.AppError) { if err := fs.db.Exec("select * from teldrive.move_directory(? , ? , ?)", payload.Source, payload.Destination, userId).Error; err != nil { return nil, &types.AppError{Error: err} } return &schemas.Message{Message: "directory moved"}, nil } func (fs *FileService) GetCategoryStats(userId int64) ([]schemas.FileCategoryStats, *types.AppError) { var stats []schemas.FileCategoryStats if err := fs.db.Model(&models.File{}).Select("category", "COUNT(*) as total_files", "coalesce(SUM(size),0) as total_size"). Where(&models.File{UserID: userId, Type: "file", Status: "active"}). Order("category ASC").Group("category").Find(&stats).Error; err != nil { return nil, &types.AppError{Error: err} } return stats, nil } func (fs *FileService) CopyFile(c *gin.Context) (*schemas.FileOut, *types.AppError) { var payload schemas.Copy if err := c.ShouldBindJSON(&payload); err != nil { return nil, &types.AppError{Error: err, Code: http.StatusBadRequest} } userId, session := GetUserAuth(c) client, _ := tgc.AuthClient(c, fs.cnf, session) var res []models.File if err := fs.db.Model(&models.File{}).Where("id = ?", payload.ID).Find(&res).Error; err != nil { return nil, &types.AppError{Error: err} } file := mapper.ToFileOutFull(res[0]) newIds := models.Parts{} channelId, err := GetDefaultChannel(c, fs.db, userId) if err != nil { return nil, &types.AppError{Error: err} } err = tgc.RunWithAuth(c, client, "", func(ctx context.Context) error { user := strconv.FormatInt(userId, 10) messages, err := getTGMessages(c, client, file.Parts, file.ChannelID, user) if err != nil { return err } channel, err := GetChannelById(ctx, client, channelId, user) if err != nil { return err } for i, message := range messages { item := message.(*tg.Message) media := item.Media.(*tg.MessageMediaDocument) document := media.Document.(*tg.Document) id, _ := randInt64() request := tg.MessagesSendMediaRequest{ Silent: true, Peer: &tg.InputPeerChannel{ChannelID: channel.ChannelID, AccessHash: channel.AccessHash}, Media: &tg.InputMediaDocument{ID: document.AsInput()}, RandomID: id, } res, err := client.API().MessagesSendMedia(c, &request) if err != nil { return err } updates := res.(*tg.Updates) var msg *tg.Message for _, update := range updates.Updates { channelMsg, ok := update.(*tg.UpdateNewChannelMessage) if ok { msg = channelMsg.Message.(*tg.Message) break } } newIds = append(newIds, models.Part{ID: int64(msg.ID), Salt: file.Parts[i].Salt}) } return nil }) if err != nil { return nil, &types.AppError{Error: err} } var destRes []models.File if err := fs.db.Raw("select * from teldrive.create_directories(?, ?)", userId, payload.Destination).Scan(&destRes).Error; err != nil { return nil, &types.AppError{Error: err} } dest := destRes[0] dbFile := models.File{} dbFile.Name = payload.Name dbFile.Size = &file.Size dbFile.Type = file.Type dbFile.MimeType = file.MimeType dbFile.Parts = &newIds dbFile.UserID = userId dbFile.Starred = false dbFile.Status = "active" dbFile.ParentID = dest.ID dbFile.ChannelID = &channelId dbFile.Encrypted = file.Encrypted dbFile.Category = file.Category if err := fs.db.Create(&dbFile).Error; err != nil { return nil, &types.AppError{Error: err} } return mapper.ToFileOut(dbFile), nil } func (fs *FileService) GetFileStream(c *gin.Context) { w := c.Writer r := c.Request fileID := c.Param("fileID") authHash := c.Query("hash") if authHash == "" { http.Error(w, "missing hash param", http.StatusBadRequest) return } cache := cache.FromContext(c) session, err := getSessionByHash(fs.db, cache, authHash) if err != nil { http.Error(w, "invalid hash", http.StatusBadRequest) return } file := &schemas.FileOutFull{} key := fmt.Sprintf("files:%s", fileID) err = cache.Get(key, file) var appErr *types.AppError if err != nil { file, appErr = fs.GetFileByID(fileID) if appErr != nil { http.Error(w, appErr.Error.Error(), http.StatusBadRequest) return } cache.Set(key, file, 0) } c.Header("Accept-Ranges", "bytes") var start, end int64 rangeHeader := r.Header.Get("Range") if rangeHeader == "" { start = 0 end = file.Size - 1 w.WriteHeader(http.StatusOK) } else { ranges, err := http_range.Parse(rangeHeader, file.Size) if err == http_range.ErrNoOverlap { w.Header().Set("Content-Range", fmt.Sprintf("bytes */%d", file.Size)) http.Error(w, http_range.ErrNoOverlap.Error(), http.StatusRequestedRangeNotSatisfiable) return } if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } if len(ranges) > 1 { http.Error(w, "multiple ranges are not supported", http.StatusRequestedRangeNotSatisfiable) return } start = ranges[0].Start end = ranges[0].End c.Header("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, file.Size)) w.WriteHeader(http.StatusPartialContent) } contentLength := end - start + 1 mimeType := file.MimeType if mimeType == "" { mimeType = "application/octet-stream" } c.Header("Content-Type", mimeType) c.Header("Content-Length", strconv.FormatInt(contentLength, 10)) c.Header("E-Tag", fmt.Sprintf("\"%s\"", md5.FromString(file.ID+strconv.FormatInt(file.Size, 10)))) c.Header("Last-Modified", file.UpdatedAt.UTC().Format(http.TimeFormat)) disposition := "inline" if c.Query("d") == "1" { disposition = "attachment" } c.Header("Content-Disposition", mime.FormatMediaType(disposition, map[string]string{"filename": file.Name})) tokens, err := getBotsToken(c, fs.db, session.UserId, file.ChannelID) logger := logging.FromContext(c) if err != nil { logger.Error("failed to get bots", zap.Error(err)) http.Error(w, err.Error(), http.StatusInternalServerError) return } var ( channelUser string lr io.ReadCloser ) var client *tgc.Client if fs.cnf.DisableStreamBots || len(tokens) == 0 { tgClient, _ := tgc.AuthClient(c, fs.cnf, session.Session) client, err = fs.worker.UserWorker(tgClient, session.UserId) if err != nil { logger.Error("file stream", zap.Error(err)) http.Error(w, err.Error(), http.StatusInternalServerError) return } channelUser = strconv.FormatInt(session.UserId, 10) logger.Debugw("requesting file", "name", file.Name, "bot", channelUser, "user", channelUser, "start", start, "end", end, "fileSize", file.Size) } else { var index int limit := min(len(tokens), fs.cnf.BgBotsLimit) fs.worker.Set(tokens[:limit], file.ChannelID) client, index, err = fs.worker.Next(file.ChannelID) if err != nil { logger.Error("file stream", zap.Error(err)) http.Error(w, err.Error(), http.StatusInternalServerError) return } channelUser = strings.Split(tokens[index], ":")[0] logger.Debugw("requesting file", "name", file.Name, "bot", channelUser, "botNo", index, "start", start, "end", end, "fileSize", file.Size) } if r.Method != "HEAD" { parts, err := getParts(c, client.Tg, file, channelUser) if err != nil { logger.Error("file stream", err) http.Error(w, err.Error(), http.StatusInternalServerError) return } if file.Encrypted { lr, err = reader.NewDecryptedReader(c, client.Tg, parts, start, end, fs.cnf.Uploads.EncryptionKey) } else { lr, err = reader.NewLinearReader(c, client.Tg, parts, start, end) } if err != nil { logger.Error("file stream", err) http.Error(w, err.Error(), http.StatusInternalServerError) return } if lr == nil { http.Error(w, "failed to initialise reader", http.StatusInternalServerError) return } io.CopyN(w, lr, contentLength) } } func setOrderFilter(query *gorm.DB, fquery *schemas.FileQuery) *gorm.DB { if fquery.NextPageToken != "" { sortColumn := utils.CamelToSnake(fquery.Sort) tokenValue, err := base64.StdEncoding.DecodeString(fquery.NextPageToken) if err == nil { if fquery.Order == "asc" { return query.Where(fmt.Sprintf("%s > ?", sortColumn), string(tokenValue)) } else { return query.Where(fmt.Sprintf("%s < ?", sortColumn), string(tokenValue)) } } } return query } func getOrder(fquery *schemas.FileQuery) clause.OrderByColumn { sortColumn := utils.CamelToSnake(fquery.Sort) return clause.OrderByColumn{Column: clause.Column{Name: sortColumn}, Desc: fquery.Order == "desc"} }