From 3326b2aa2fdc1076884f5bc81e8d7616633c0c79 Mon Sep 17 00:00:00 2001
From: wdvxdr <wdvxdr1123@gmail.com>
Date: Fri, 6 Aug 2021 22:26:42 +0800
Subject: [PATCH] feat(coolq): mimetypes switch to array.

---
 coolq/bot.go    | 32 ++++++++++++++++++--------------
 coolq/cqcode.go | 13 ++++++++++---
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/coolq/bot.go b/coolq/bot.go
index 3c725a2..98e998b 100644
--- a/coolq/bot.go
+++ b/coolq/bot.go
@@ -81,9 +81,9 @@ var ForceFragmented = false
 // SkipMimeScan 是否跳过Mime扫描
 var SkipMimeScan bool
 
-var lawfulImageTypes = []string{"image/png", "image/jpeg", "image/gif", "image/bmp", "image/webp"}
+var lawfulImageTypes = [...]string{"image/png", "image/jpeg", "image/gif", "image/bmp", "image/webp"}
 
-var lawfulAudioTypes = []string{
+var lawfulAudioTypes = [...]string{
 	"audio/mpeg", "audio/flac", "audio/midi", "audio/ogg",
 	"audio/ape", "audio/amr", "audio/wav", "audio/aiff",
 	"audio/mp4", "audio/aac", "audio/x-m4a",
@@ -228,18 +228,18 @@ func (bot *CQBot) UploadLocalVideo(target int64, v *LocalVideoElement) (*message
 
 // UploadLocalImageAsPrivate 上传本地图片至私聊
 func (bot *CQBot) UploadLocalImageAsPrivate(userID int64, img *LocalImageElement) (i *message.FriendImageElement, err error) {
-	if img.Stream != nil {
-		i, err = bot.Client.UploadPrivateImage(userID, img.Stream)
-	} else {
-		// need update.
-		f, e := os.Open(img.File)
-		if e != nil {
-			return nil, e
+	if img.File != "" {
+		f, err := os.Open(img.File)
+		if err != nil {
+			return nil, errors.Wrap(err, "open image error")
 		}
-		defer f.Close()
-		i, err = bot.Client.UploadPrivateImage(userID, f)
+		defer func() { _ = f.Close() }()
+		img.Stream = f
 	}
-
+	if lawful, mime := IsLawfulImage(img.Stream); !lawful {
+		return nil, errors.New("image type error: " + mime)
+	}
+	i, err = bot.Client.UploadPrivateImage(userID, img.Stream)
 	if i != nil {
 		i.Flash = img.Flash
 	}
@@ -604,6 +604,10 @@ func IsLawfulImage(r io.ReadSeeker) (bool, string) {
 		log.Debugf("扫描 Mime 时出现问题: %v", err)
 		return false, ""
 	}
-	mime := t.String()
-	return mimetype.EqualsAny(mime, lawfulImageTypes...), mime
+	for _, lt := range lawfulImageTypes {
+		if t.Is(lt) {
+			return true, t.String()
+		}
+	}
+	return false, t.String()
 }
diff --git a/coolq/cqcode.go b/coolq/cqcode.go
index e7dce37..d7d4e4d 100644
--- a/coolq/cqcode.go
+++ b/coolq/cqcode.go
@@ -731,9 +731,16 @@ func (bot *CQBot) ToElement(t string, d map[string]string, isGroup bool) (m inte
 			return nil, err
 		}
 		if !SkipMimeScan && !global.IsAMRorSILK(data) {
-			mt := mimetype.Detect(data).String()
-			if !mimetype.EqualsAny(mt, lawfulAudioTypes...) {
-				return nil, errors.New("audio type error: " + mt)
+			mt := mimetype.Detect(data)
+			lawful := false
+			for _, lt := range lawfulAudioTypes {
+				if mt.Is(lt) {
+					lawful = true
+					break
+				}
+			}
+			if !lawful {
+				return nil, errors.New("audio type error: " + mt.String())
 			}
 		}
 		if !global.IsAMRorSILK(data) {