media: major refactor of downloading/direct URL

Signed-off-by: Sumner Evans <sumner.evans@automattic.com>
2024-07-09 23:03:15 -06:00
parent 7e680f1fee
commit 58cc638058
8 changed files with 319 additions and 310 deletions
@@ -1,17 +0,0 @@
-package media
-
-import (
-	"context"
-
-	"github.com/gotd/td/telegram/downloader"
-	"github.com/gotd/td/tg"
-)
-
-func DownloadDocument(ctx context.Context, client downloader.Client, document *tg.Document) ([]byte, error) {
-	data, _, err := DownloadFileLocation(ctx, client, &tg.InputDocumentFileLocation{
-		ID:            document.GetID(),
-		AccessHash:    document.GetAccessHash(),
-		FileReference: document.GetFileReference(),
-	})
-	return data, err
-}
@@ -1,40 +0,0 @@
-package media
-
-import (
-	"bytes"
-	"context"
-	"net/http"
-
-	"github.com/gotd/td/telegram/downloader"
-	"github.com/gotd/td/tg"
-)
-
-func DownloadFileLocation(ctx context.Context, client downloader.Client, loc tg.InputFileLocationClass) (data []byte, mimeType string, err error) {
-	// TODO convert entire function to streaming? Maybe at least stream to file?
-	var buf bytes.Buffer
-	storageFileTypeClass, err := downloader.NewDownloader().Download(client, loc).Stream(ctx, &buf)
-	if err != nil {
-		return nil, "", err
-	}
-	switch storageFileTypeClass.(type) {
-	case *tg.StorageFileJpeg:
-		mimeType = "image/jpeg"
-	case *tg.StorageFileGif:
-		mimeType = "image/gif"
-	case *tg.StorageFilePng:
-		mimeType = "image/png"
-	case *tg.StorageFilePdf:
-		mimeType = "application/pdf"
-	case *tg.StorageFileMp3:
-		mimeType = "audio/mp3"
-	case *tg.StorageFileMov:
-		mimeType = "video/quicktime"
-	case *tg.StorageFileMp4:
-		mimeType = "video/mp4"
-	case *tg.StorageFileWebp:
-		mimeType = "image/webp"
-	default:
-		mimeType = http.DetectContentType(buf.Bytes())
-	}
-	return buf.Bytes(), mimeType, nil
-}
@@ -1,69 +0,0 @@
-package media
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/gotd/td/telegram/downloader"
-	"github.com/gotd/td/tg"
-)
-
-type dimensionable interface {
-	GetW() int
-	GetH() int
-}
-
-func GetLargestPhotoSize(sizes []tg.PhotoSizeClass) (width, height int, largest tg.PhotoSizeClass) {
-	if len(sizes) == 0 {
-		panic("cannot get largest size from empty list of sizes")
-	}
-
-	var maxSize int
-	for _, s := range sizes {
-		var currentSize int
-		switch size := s.(type) {
-		case *tg.PhotoSize:
-			currentSize = size.GetSize()
-		case *tg.PhotoCachedSize:
-			currentSize = max(size.GetW(), size.GetH())
-		case *tg.PhotoSizeProgressive:
-			currentSize = max(size.GetW(), size.GetH())
-		case *tg.PhotoPathSize:
-			currentSize = len(size.GetBytes())
-		case *tg.PhotoStrippedSize:
-			currentSize = len(size.GetBytes())
-		}
-
-		if currentSize > maxSize {
-			maxSize = currentSize
-			largest = s
-			if d, ok := s.(dimensionable); ok {
-				width = d.GetW()
-				height = d.GetH()
-			}
-		}
-	}
-	return
-}
-
-func DownloadPhoto(ctx context.Context, client downloader.Client, photo *tg.Photo) (data []byte, width, height int, mimeType string, err error) {
-	var largest tg.PhotoSizeClass
-	width, height, largest = GetLargestPhotoSize(photo.GetSizes())
-	data, mimeType, err = DownloadFileLocation(ctx, client, &tg.InputPhotoFileLocation{
-		ID:            photo.GetID(),
-		AccessHash:    photo.GetAccessHash(),
-		FileReference: photo.GetFileReference(),
-		ThumbSize:     largest.GetType(),
-	})
-	return
-}
-
-func DownloadPhotoMedia(ctx context.Context, client downloader.Client, media *tg.MessageMediaPhoto) (data []byte, width, height int, mimeType string, err error) {
-	if p, ok := media.GetPhoto(); !ok {
-		return nil, 0, 0, "", fmt.Errorf("photo message sent without a photo")
-	} else if photo, ok := p.(*tg.Photo); !ok {
-		return nil, 0, 0, "", fmt.Errorf("unrecognized photo type %T", p)
-	} else {
-		return DownloadPhoto(ctx, client, photo)
-	}
-}
@@ -1,20 +1,62 @@
 package media

 import (
+	"bytes"
 	"context"
 	"fmt"
+	"net/http"

 	"github.com/gotd/td/telegram/downloader"
 	"github.com/gotd/td/tg"
+	"github.com/rs/zerolog"
 	"maunium.net/go/mautrix/bridgev2"
 	"maunium.net/go/mautrix/event"
 	"maunium.net/go/mautrix/id"

 	"go.mau.fi/util/lottie"

+	"go.mau.fi/mautrix-telegram/pkg/connector/ids"
 	"go.mau.fi/mautrix-telegram/pkg/connector/store"
 )

+type dimensionable interface {
+	GetW() int
+	GetH() int
+}
+
+func getLargestPhotoSize(sizes []tg.PhotoSizeClass) (width, height int, largest tg.PhotoSizeClass) {
+	if len(sizes) == 0 {
+		panic("cannot get largest size from empty list of sizes")
+	}
+
+	var maxSize int
+	for _, s := range sizes {
+		var currentSize int
+		switch size := s.(type) {
+		case *tg.PhotoSize:
+			currentSize = size.GetSize()
+		case *tg.PhotoCachedSize:
+			currentSize = max(size.GetW(), size.GetH())
+		case *tg.PhotoSizeProgressive:
+			currentSize = max(size.GetW(), size.GetH())
+		case *tg.PhotoPathSize:
+			currentSize = len(size.GetBytes())
+		case *tg.PhotoStrippedSize:
+			currentSize = len(size.GetBytes())
+		}
+
+		if currentSize > maxSize {
+			maxSize = currentSize
+			largest = s
+			if d, ok := s.(dimensionable); ok {
+				width = d.GetW()
+				height = d.GetH()
+			}
+		}
+	}
+	return
+}
+
 // getLocationID converts a Telegram [tg.Document],
 // [tg.InputDocumentFileLocation], [tg.InputPeerPhotoFileLocation],
 // [tg.InputFileLocation], or [tg.InputPhotoFileLocation] into a [LocationID]
@@ -56,72 +98,234 @@ func (c AnimatedStickerConfig) WebmConvert() bool {
 	return c.ConvertFromWebm && c.Target != "webm"
 }

+// Transferer is a utility for downloading media from Telegram and uploading it
+// to Matrix.
+// TODO better name?
 type Transferer struct {
-	RoomID    id.RoomID
-	Filename  string
-	IsSticker bool
-	Config    AnimatedStickerConfig
+	client downloader.Client
+
+	roomID                id.RoomID
+	filename              string
+	animatedStickerConfig *AnimatedStickerConfig
+
+	fileInfo event.FileInfo
 }

-func NewTransferer(cfg AnimatedStickerConfig) *Transferer {
-	return &Transferer{Config: cfg}
+type ReadyTransferer struct {
+	inner *Transferer
+	loc   tg.InputFileLocationClass
 }

+// NewTransferer creates a new [Transferer] with the given [downloader.Client].
+// The client is used to download the media from Telegram.
+func NewTransferer(client downloader.Client) *Transferer {
+	return &Transferer{client: client}
+}
+
+// WithRoomID sets the room ID for the [Transferer].
 func (t *Transferer) WithRoomID(roomID id.RoomID) *Transferer {
-	t.RoomID = roomID
+	t.roomID = roomID
 	return t
 }

+// WithFilename sets the filename for the [Transferer].
 func (t *Transferer) WithFilename(filename string) *Transferer {
-	t.Filename = filename
+	t.filename = filename
 	return t
 }

-func (t *Transferer) WithIsSticker(isSticker bool) *Transferer {
-	t.IsSticker = isSticker
+func (t *Transferer) WithMIMEType(mimeType string) *Transferer {
+	t.fileInfo.MimeType = mimeType
 	return t
 }

-func (t *Transferer) Transfer(ctx context.Context, store *store.Container, client downloader.Client, intent bridgev2.MatrixAPI, loc tg.InputFileLocationClass) (mxc id.ContentURIString, encryptedFileInfo *event.EncryptedFileInfo, size int, mimeType string, err error) {
-	locationID := getLocationID(loc)
+// WithStickerConfig sets the animated sticker config for the [Transferer].
+func (t *Transferer) WithStickerConfig(cfg AnimatedStickerConfig) *Transferer {
+	t.animatedStickerConfig = &cfg
+	return t
+}
+
+func (t *Transferer) WithThumbnail(uri id.ContentURIString, file *event.EncryptedFileInfo, info *event.FileInfo) *Transferer {
+	t.fileInfo.ThumbnailURL = uri
+	t.fileInfo.ThumbnailFile = file
+	t.fileInfo.ThumbnailInfo = info
+	return t
+}
+
+func (t *Transferer) WithVideo(attr *tg.DocumentAttributeVideo) *Transferer {
+	t.fileInfo.Width, t.fileInfo.Height = attr.W, attr.H
+	t.fileInfo.Duration = int(attr.Duration * 1000)
+	return t
+}
+
+// WithDocument transforms a [Transferer] to a [ReadyTransferer] by setting the
+// given document as the location that will be downloaded by the
+// [ReadyTransferer].
+func (t *Transferer) WithDocument(doc tg.DocumentClass, thumbnail bool) *ReadyTransferer {
+	document := doc.(*tg.Document)
+	documentFileLocation := tg.InputDocumentFileLocation{
+		ID:            document.GetID(),
+		AccessHash:    document.GetAccessHash(),
+		FileReference: document.GetFileReference(),
+	}
+	if thumbnail {
+		_, _, largestThumbnail := getLargestPhotoSize(document.Thumbs)
+		documentFileLocation.ThumbSize = largestThumbnail.GetType()
+	} else {
+		t.fileInfo.Size = int(document.Size)
+		t.fileInfo.MimeType = document.GetMimeType()
+	}
+	return &ReadyTransferer{t, &documentFileLocation}
+}
+
+// WithPhoto transforms a [Transferer] to a [ReadyTransferer] by setting the
+// given photo as the location that will be downloaded by the
+// [ReadyTransferer].
+func (t *Transferer) WithPhoto(pc tg.PhotoClass) *ReadyTransferer {
+	photo := pc.(*tg.Photo)
+	var largest tg.PhotoSizeClass
+	t.fileInfo.Width, t.fileInfo.Height, largest = getLargestPhotoSize(photo.GetSizes())
+	return &ReadyTransferer{
+		inner: t,
+		loc: &tg.InputPhotoFileLocation{
+			ID:            photo.GetID(),
+			AccessHash:    photo.GetAccessHash(),
+			FileReference: photo.GetFileReference(),
+			ThumbSize:     largest.GetType(),
+		},
+	}
+}
+
+// WithUser transforms a [Transferer] to a [ReadyTransferer] by setting the
+// given user's photo as the location that will be downloaded by the
+// [ReadyTransferer].
+func (t *Transferer) WithUserPhoto(user *tg.User, photoID int64) *ReadyTransferer {
+	return &ReadyTransferer{
+		inner: t,
+		loc: &tg.InputPeerPhotoFileLocation{
+			Peer:    &tg.InputPeerUser{UserID: user.GetID()},
+			PhotoID: photoID,
+			Big:     true,
+		},
+	}
+}
+
+// Transfer downloads the media from Telegram and uploads it to Matrix.
+//
+// If the file is already in the database, the MXC URI will be reused. The
+// file's MXC URI will only be cached if the room ID is unset or if the room is
+// not encrypted.
+//
+// If there is a sticker config on the [Transferer], this function converts
+// animated stickers to the target format specified by the specified
+// [AnimatedStickerConfig].
+func (t *ReadyTransferer) Transfer(ctx context.Context, store *store.Container, intent bridgev2.MatrixAPI) (mxc id.ContentURIString, encryptedFileInfo *event.EncryptedFileInfo, outFileInfo *event.FileInfo, err error) {
+	locationID := getLocationID(t.loc)
+	log := zerolog.Ctx(ctx).With().
+		Str("component", "media_transfer").
+		Str("location_id", string(locationID)).
+		Logger()
+
 	if file, err := store.TelegramFile.GetByLocationID(ctx, locationID); err != nil {
-		return "", nil, 0, "", fmt.Errorf("failed to search for Telegram file by location ID: %w", err)
+		return "", nil, nil, fmt.Errorf("failed to search for Telegram file by location ID: %w", err)
 	} else if file != nil {
-		return file.MXC, nil, file.Size, file.MIMEType, nil
+		t.inner.fileInfo.Size, t.inner.fileInfo.MimeType = file.Size, file.MIMEType
+		return file.MXC, nil, &t.inner.fileInfo, nil
 	}

-	var data []byte
-	data, mimeType, err = DownloadFileLocation(ctx, client, loc)
+	data, _, err := t.Download(ctx)
 	if err != nil {
-		return "", nil, 0, "", fmt.Errorf("downloading file failed: %w", err)
+		return "", nil, nil, fmt.Errorf("downloading file failed: %w", err)
 	}

-	if t.IsSticker {
-		if lottie.Supported() && t.Config.TGSConvert() && mimeType == "application/x-gzip" {
-			data, err = lottie.ConvertBytes(ctx, data, t.Config.Target, t.Config.Args.Width, t.Config.Args.Height, fmt.Sprintf("%d", t.Config.Args.FPS))
+	if t.inner.animatedStickerConfig != nil {
+		if lottie.Supported() && t.inner.animatedStickerConfig.TGSConvert() && t.inner.fileInfo.MimeType == "application/x-tgsticker" {
+			newData, err := lottie.ConvertBytes(ctx, data,
+				t.inner.animatedStickerConfig.Target,
+				t.inner.animatedStickerConfig.Args.Width,
+				t.inner.animatedStickerConfig.Args.Height,
+				fmt.Sprintf("%d", t.inner.animatedStickerConfig.Args.FPS))
 			if err != nil {
-				return "", nil, 0, "", err
+				log.Err(err).Msg("failed to convert animated sticker")
+			} else {
+				data = newData
+				t.inner.fileInfo.Size = len(data)
+				t.inner.fileInfo.MimeType = fmt.Sprintf("image/%s", t.inner.animatedStickerConfig.Target)
 			}
-			mimeType = fmt.Sprintf("image/%s", t.Config.Target)
 			// TODO support ffmpeg conversion
 			// } else if ffmpeg.Supported() && t.Config.WebmConvert() && mimeType == "video/webm" {
 		}
 	}

-	mxcURI, encryptedFileInfo, err := intent.UploadMedia(ctx, t.RoomID, data, t.Filename, mimeType)
+	mxc, encryptedFileInfo, err = intent.UploadMedia(ctx, t.inner.roomID, data, t.inner.filename, t.inner.fileInfo.MimeType)
 	if err != nil {
-		return "", nil, 0, "", err
+		return "", nil, nil, fmt.Errorf("failed to upload media to Matrix: %w", err)
 	}
-	if len(mxcURI) > 0 {
+
+	// If it's an unencrypted file, cache the MXC URI corresponding to the
+	// location ID.
+	if len(mxc) > 0 {
 		file := store.TelegramFile.New()
 		file.LocationID = locationID
-		file.MXC = mxcURI
-		file.Size = len(data)
-		file.MIMEType = mimeType
-		// TODO width, height, thumbnail?
+		file.MXC = mxc
+		file.Size = t.inner.fileInfo.Size
+		file.MIMEType = t.inner.fileInfo.MimeType
 		if err = file.Insert(ctx); err != nil {
-			return "", nil, 0, "", fmt.Errorf("failed to insert Telegram file into database: %w", err)
+			log.Err(err).Msg("failed to insert Telegram file into database")
 		}
 	}
-	return mxcURI, encryptedFileInfo, len(data), mimeType, nil
+	return mxc, encryptedFileInfo, &t.inner.fileInfo, nil
+}
+
+// Download downloads the media from Telegram.
+func (t *ReadyTransferer) Download(ctx context.Context) (data []byte, fileInfo *event.FileInfo, err error) {
+	// TODO convert entire function to streaming? Maybe at least stream to file?
+	var buf bytes.Buffer
+	storageFileTypeClass, err := downloader.NewDownloader().Download(t.inner.client, t.loc).Stream(ctx, &buf)
+	if err != nil {
+		return nil, nil, err
+	}
+	if t.inner.fileInfo.MimeType == "" {
+		switch storageFileTypeClass.(type) {
+		case *tg.StorageFileJpeg:
+			t.inner.fileInfo.MimeType = "image/jpeg"
+		case *tg.StorageFileGif:
+			t.inner.fileInfo.MimeType = "image/gif"
+		case *tg.StorageFilePng:
+			t.inner.fileInfo.MimeType = "image/png"
+		case *tg.StorageFilePdf:
+			t.inner.fileInfo.MimeType = "application/pdf"
+		case *tg.StorageFileMp3:
+			t.inner.fileInfo.MimeType = "audio/mp3"
+		case *tg.StorageFileMov:
+			t.inner.fileInfo.MimeType = "video/quicktime"
+		case *tg.StorageFileMp4:
+			t.inner.fileInfo.MimeType = "video/mp4"
+		case *tg.StorageFileWebp:
+			t.inner.fileInfo.MimeType = "image/webp"
+		default:
+			t.inner.fileInfo.MimeType = http.DetectContentType(buf.Bytes())
+		}
+	}
+	t.inner.fileInfo.Size = len(data)
+	return buf.Bytes(), &t.inner.fileInfo, nil
+}
+
+// DirectDownloadURL returns the direct download URL for the media.
+func (t *ReadyTransferer) DirectDownloadURL(ctx context.Context, portal *bridgev2.Portal, msgID int, thumbnail bool) (id.ContentURIString, *event.FileInfo, error) {
+	peerType, chatID, err := ids.ParsePortalID(portal.ID)
+	if err != nil {
+		return "", nil, err
+	}
+	mediaID, err := ids.DirectMediaInfo{
+		PeerType:  peerType,
+		ChatID:    chatID,
+		MessageID: int64(msgID),
+		Thumbnail: thumbnail,
+	}.AsMediaID()
+	if err != nil {
+		return "", nil, err
+	}
+	mxc, err := portal.Bridge.Matrix.GenerateContentURI(ctx, mediaID)
+	return mxc, &t.inner.fileInfo, err
 }