From cdada929a505a892a108ec11318f8c95f0c79e42 Mon Sep 17 00:00:00 2001
From: Deng Junhai <DengJunHai@foxmail.com>
Date: Sat, 22 Jun 2024 11:43:44 +0800
Subject: [PATCH] fix: fix image extract ext (#208)

Co-Authored-By: Minghan Zhang <112773885+zmh-program@users.noreply.github.com>
---
 .gitignore        |  1 +
 manager/images.go |  7 ++++---
 utils/char.go     | 17 +++++++++++++++--
 utils/image.go    |  5 +++--
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index af5ece9d..57af2a09 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,6 +14,7 @@ sdk
 logs
 
 chat
+*.exe
 chat.exe
 
 # for reverse engine
diff --git a/manager/images.go b/manager/images.go
index 228d21ee..1f036dc6 100644
--- a/manager/images.go
+++ b/manager/images.go
@@ -1,17 +1,18 @@
 package manager
 
 import (
-	"chat/adapter/common"
+	adaptercommon "chat/adapter/common"
 	"chat/admin"
 	"chat/auth"
 	"chat/channel"
 	"chat/globals"
 	"chat/utils"
 	"fmt"
-	"github.com/gin-gonic/gin"
 	"net/http"
 	"strings"
 	"time"
+
+	"github.com/gin-gonic/gin"
 )
 
 func ImagesRelayAPI(c *gin.Context) {
@@ -74,7 +75,7 @@ func getImageProps(form RelayImageForm, messages []globals.Message, buffer *util
 func getUrlFromBuffer(buffer *utils.Buffer) string {
 	content := buffer.Read()
 
-	_, urls := utils.ExtractImages(content, true)
+	urls := utils.ExtractImagesFromMarkdown(content)
 	if len(urls) > 0 {
 		return urls[len(urls)-1]
 	}
diff --git a/utils/char.go b/utils/char.go
index a9129e60..9ccff39b 100644
--- a/utils/char.go
+++ b/utils/char.go
@@ -2,11 +2,12 @@ package utils
 
 import (
 	"fmt"
-	"github.com/goccy/go-json"
 	"regexp"
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/goccy/go-json"
 )
 
 func GetRandomInt(min int, max int) int {
@@ -220,6 +221,18 @@ func ExtractImages(data string, includeBase64 bool) (content string, images []st
 	return content, images
 }
 
+func ExtractImagesFromMarkdown(data string) (images []string) {
+	// extract images like `![image](https://xxx.com/xxx?xxx=xxx&xxx=xxx)` and return urls
+	re := regexp.MustCompile(`!\[.*\]\((https?://\S+)\)`)
+	matches := re.FindAllStringSubmatch(data, -1)
+
+	for _, match := range matches {
+		images = append(images, match[1])
+	}
+
+	return images
+}
+
 func ExtractBase64Images(data string) []string {
 	// get base64 images from data (data:image/png;base64,xxxxxx) (\n \\n [space] \\t \\r \\v \\f break the base64 string)
 	re := regexp.MustCompile(`(data:image/\w+;base64,[\w+/=]+)`)
@@ -229,7 +242,7 @@ func ExtractBase64Images(data string) []string {
 func ExtractExternalImages(data string) []string {
 	// https://platform.openai.com/docs/guides/vision/what-type-of-files-can-i-upload
 
-	re := regexp.MustCompile(`(https?://\S+\.(?:png|jpg|jpeg|gif|webp|heif|heic)(?:\s\S+)?)`)
+	re := regexp.MustCompile(`(https?://\S+\.(?:png|jpg|jpeg|gif|webp|heif|heic|bmp|svg|ico)(?:\?\S+)?)`)
 	return re.FindAllString(data, -1)
 }
 
diff --git a/utils/image.go b/utils/image.go
index 18aaaf03..fc5f5300 100644
--- a/utils/image.go
+++ b/utils/image.go
@@ -3,7 +3,6 @@ package utils
 import (
 	"chat/globals"
 	"fmt"
-	"github.com/chai2010/webp"
 	"image"
 	"image/gif"
 	"image/jpeg"
@@ -13,6 +12,8 @@ import (
 	"os"
 	"path"
 	"strings"
+
+	"github.com/chai2010/webp"
 )
 
 type Image struct {
@@ -226,7 +227,7 @@ func DownloadImage(url string, path string) error {
 
 func StoreImage(url string) string {
 	if globals.AcceptImageStore {
-		hash := Md5Encrypt(url)
+		hash := Md5Encrypt(url) + path.Ext(url)
 
 		if err := DownloadImage(url, fmt.Sprintf("storage/attachments/%s", hash)); err != nil {
 			globals.Warn(fmt.Sprintf("[utils] save image error: %s", err.Error()))