From 0b736ee74da38ac5aee5e79e3ce3937ab85321e2 Mon Sep 17 00:00:00 2001
From: Rogee <rogeecn@qq.com>
Date: Wed, 4 Sep 2024 19:24:10 +0800
Subject: [PATCH] fix: ext issues

---
 internal/client_channel.go | 59 ++++++++++++++++++++++++--------------
 main_test.go               |  1 +
 2 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/internal/client_channel.go b/internal/client_channel.go
index 902babf..46f3bb7 100644
--- a/internal/client_channel.go
+++ b/internal/client_channel.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"mime"
 	"os"
+	"strings"
 
 	"exporter/pkg/errorx"
 
@@ -159,6 +160,32 @@ func (t *TClient) savePhoto(ctx context.Context, cfg *DBChannel, photo *tg.Photo
 	return nil
 }
 
+func (t *TClient) getExtByFilenameAndMimeType(filename, mimeType string) string {
+	ext := ""
+	if filename != "" {
+		ext = filename[strings.LastIndex(filename, "."):]
+		return ext
+	}
+
+	exts, err := mime.ExtensionsByType(mimeType)
+	if err != nil {
+		logger.Error("get extension failed", zap.Error(err), zap.String("mime_type", mimeType))
+		return ext
+	}
+
+	if len(exts) == 0 {
+		logger.Warn("no extension found", zap.String("mime_type", mimeType))
+		switch mimeType {
+		case "application/rar":
+			ext = ".rar"
+		}
+	} else {
+		ext = exts[len(exts)-1]
+	}
+
+	return ext
+}
+
 // saveDocument
 func (t *TClient) saveDocument(ctx context.Context, cfg *DBChannel, doc *tg.Document) (ChannelMessageDocument, error) {
 	location := &tg.InputDocumentFileLocation{
@@ -166,31 +193,10 @@ func (t *TClient) saveDocument(ctx context.Context, cfg *DBChannel, doc *tg.Docu
 		AccessHash:    doc.GetAccessHash(),
 		FileReference: doc.GetFileReference(),
 	}
-	exts, err := mime.ExtensionsByType(doc.GetMimeType())
-	if err != nil {
-		logger.Error("get extension failed", zap.Error(err), zap.String("mime_type", doc.GetMimeType()))
-		return ChannelMessageDocument{}, err
-	}
 
-	if len(exts) == 0 {
-		logger.Warn("no extension found", zap.String("mime_type", doc.GetMimeType()))
-		switch doc.GetMimeType() {
-		case "application/rar":
-			exts = []string{".rar"}
-		}
-	}
-	ext := exts[len(exts)-1]
-
-	saveTo := cfg.Asset(doc.GetID(), ext)
-	downloader := downloader.NewDownloader()
-	_, err = downloader.Download(t.Client.API(), location).ToPath(ctx, saveTo)
-	if err != nil {
-		os.Remove(saveTo)
-		logger.Error("download failed", zap.Error(err))
-		return ChannelMessageDocument{}, err
-	}
 	docAttr := doc.GetAttributes()
 
+	ext := t.getExtByFilenameAndMimeType("", doc.GetMimeType())
 	data := ChannelMessageDocument{
 		Ext:      ext,
 		MimeType: doc.GetMimeType(),
@@ -201,6 +207,7 @@ func (t *TClient) saveDocument(ctx context.Context, cfg *DBChannel, doc *tg.Docu
 			switch attr.(type) {
 			case *tg.DocumentAttributeFilename:
 				data.Filename = attr.(*tg.DocumentAttributeFilename).GetFileName()
+				data.Ext = t.getExtByFilenameAndMimeType("", doc.GetMimeType())
 			case *tg.DocumentAttributeVideo:
 				m := attr.(*tg.DocumentAttributeVideo)
 				data.Video = &ChannelMessageDocumentVideo{
@@ -213,6 +220,14 @@ func (t *TClient) saveDocument(ctx context.Context, cfg *DBChannel, doc *tg.Docu
 		}
 	}
 
+	saveTo := cfg.Asset(doc.GetID(), ext)
+	downloader := downloader.NewDownloader()
+	_, err := downloader.Download(t.Client.API(), location).ToPath(ctx, saveTo)
+	if err != nil {
+		os.Remove(saveTo)
+		logger.Error("download failed", zap.Error(err))
+		return ChannelMessageDocument{}, err
+	}
 	logger.Info("download document success", zap.String("location", saveTo), zap.Any("document", data))
 	return data, nil
 }
diff --git a/main_test.go b/main_test.go
index 3c2a69b..ea03cd5 100644
--- a/main_test.go
+++ b/main_test.go
@@ -15,6 +15,7 @@ func Test_Mime(t *testing.T) {
 	// Test code here
 	m := "application/rar"
 	m = "video/mp4"
+	m = "video/quicktime"
 
 	t.Log(mime.ExtensionsByType(m))
 }