feat(bot): add support for images (#163)

EverythingSuckz · Feb 20, 2025 · e9888df · e9888df
2 parents a8dceaf + e90279b
commit e9888df
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 19 deletions.
diff --git a/internal/cache/cache.go b/internal/cache/cache.go
@@ -23,6 +23,7 @@ func InitCache(log *zap.Logger) {
 	log = log.Named("cache")
 	gob.Register(types.File{})
 	gob.Register(tg.InputDocumentFileLocation{})
+	gob.Register(tg.InputPhotoFileLocation{})
 	defer log.Sugar().Info("Initialized")
 	cache = &Cache{cache: freecache.NewCache(10 * 1024 * 1024), log: log}
 }

diff --git a/internal/commands/stream.go b/internal/commands/stream.go
@@ -32,7 +32,7 @@ func supportedMediaFilter(m *types.Message) (bool, error) {
 	case *tg.MessageMediaDocument:
 		return true, nil
 	case *tg.MessageMediaPhoto:
-		return false, nil
+		return true, nil
 	case tg.MessageMediaClass:
 		return false, dispatcher.EndGroups
 	default:
@@ -65,11 +65,6 @@ func sendLink(ctx *ext.Context, u *ext.Update) error {
 		return dispatcher.EndGroups
 	}
 	messageID := update.Updates[0].(*tg.UpdateMessageID).ID
-	if err != nil {
-		utils.Logger.Sugar().Error(err)
-		ctx.Reply(u, fmt.Sprintf("Error - %s", err.Error()), nil)
-		return dispatcher.EndGroups
-	}
 	doc := update.Updates[1].(*tg.UpdateNewChannelMessage).Message.(*tg.Message).Media
 	file, err := utils.FileFromMedia(doc)
 	if err != nil {

diff --git a/internal/routes/stream.go b/internal/routes/stream.go
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"strconv"
 
+	"github.com/gotd/td/tg"
 	range_parser "github.com/quantumsheep/range-parser"
 	"go.uber.org/zap"
 
@@ -39,10 +40,6 @@ func getStreamRoute(ctx *gin.Context) {
 		return
 	}
 
-	ctx.Header("Accept-Ranges", "bytes")
-	var start, end int64
-	rangeHeader := r.Header.Get("Range")
-
 	worker := bot.GetNextWorker()
 
 	file, err := utils.FileFromMessage(ctx, worker.Client, messageID)
@@ -62,6 +59,34 @@ func getStreamRoute(ctx *gin.Context) {
 		return
 	}
 
+	// for photo messages
+	if file.FileSize == 0 {
+		res, err := worker.Client.API().UploadGetFile(ctx, &tg.UploadGetFileRequest{
+			Location: file.Location,
+			Offset:   0,
+			Limit:    1024 * 1024,
+		})
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		result, ok := res.(*tg.UploadFile)
+		if !ok {
+			http.Error(w, "unexpected response", http.StatusInternalServerError)
+			return
+		}
+		fileBytes := result.GetBytes()
+		ctx.Header("Content-Disposition", fmt.Sprintf("inline; filename=\"%s\"", file.FileName))
+		if r.Method != "HEAD" {
+			ctx.Data(http.StatusOK, file.MimeType, fileBytes)
+		}
+		return
+	}
+
+	ctx.Header("Accept-Ranges", "bytes")
+	var start, end int64
+	rangeHeader := r.Header.Get("Range")
+
 	if rangeHeader == "" {
 		start = 0
 		end = file.FileSize - 1
@@ -98,10 +123,6 @@ func getStreamRoute(ctx *gin.Context) {
 	ctx.Header("Content-Disposition", fmt.Sprintf("%s; filename=\"%s\"", disposition, file.FileName))
 
 	if r.Method != "HEAD" {
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
 		lr, _ := utils.NewTelegramReader(ctx, worker.Client, file.Location, start, end, contentLength)
 		if _, err := io.CopyN(w, lr, contentLength); err != nil {
 			log.Error("Error while copying stream", zap.Error(err))

diff --git a/internal/types/file.go b/internal/types/file.go
@@ -10,7 +10,7 @@ import (
 )
 
 type File struct {
-	Location *tg.InputDocumentFileLocation
+	Location tg.InputFileLocationClass
 	FileSize int64
 	FileName string
 	MimeType string

diff --git a/internal/utils/helpers.go b/internal/utils/helpers.go
@@ -67,7 +67,32 @@ func FileFromMedia(media tg.MessageMediaClass) (*types.File, error) {
 			MimeType: document.MimeType,
 			ID:       document.ID,
 		}, nil
-		// TODO: add photo support
+	case *tg.MessageMediaPhoto:
+		photo, ok := media.Photo.AsNotEmpty()
+		if !ok {
+			return nil, fmt.Errorf("unexpected type %T", media)
+		}
+		sizes := photo.Sizes
+		if len(sizes) == 0 {
+			return nil, errors.New("photo has no sizes")
+		}
+		photoSize := sizes[len(sizes)-1]
+		size, ok := photoSize.AsNotEmpty()
+		if !ok {
+			return nil, errors.New("photo size is empty")
+		}
+		location := new(tg.InputPhotoFileLocation)
+		location.ID = photo.GetID()
+		location.AccessHash = photo.GetAccessHash()
+		location.FileReference = photo.GetFileReference()
+		location.ThumbSize = size.GetType()
+		return &types.File{
+			Location: location,
+			FileSize: 0, // caller should judge if this is a photo or not
+			FileName: fmt.Sprintf("photo_%d.jpg", photo.GetID()),
+			MimeType: "image/jpeg",
+			ID:       photo.GetID(),
+		}, nil
 	}
 	return nil, fmt.Errorf("unexpected type %T", media)
 }
@@ -99,7 +124,6 @@ func FileFromMessage(ctx context.Context, client *gotgproto.Client, messageID in
 		return nil, err
 	}
 	return file, nil
-	// TODO: add photo support
 }
 
 func GetLogChannelPeer(ctx context.Context, api *tg.Client, peerStorage *storage.PeerStorage) (*tg.InputChannel, error) {

diff --git a/internal/utils/reader.go b/internal/utils/reader.go
@@ -14,7 +14,7 @@ type telegramReader struct {
 	ctx           context.Context
 	log           *zap.Logger
 	client        *gotgproto.Client
-	location      *tg.InputDocumentFileLocation
+	location      tg.InputFileLocationClass
 	start         int64
 	end           int64
 	next          func() ([]byte, error)
@@ -32,7 +32,7 @@ func (*telegramReader) Close() error {
 func NewTelegramReader(
 	ctx context.Context,
 	client *gotgproto.Client,
-	location *tg.InputDocumentFileLocation,
+	location tg.InputFileLocationClass,
 	start int64,
 	end int64,
 	contentLength int64,