From 67fc00cff383f4beb55970642743d2a9fcb15e65 Mon Sep 17 00:00:00 2001
From: tbxark <tbxark@outlook.com>
Date: Mon, 2 Dec 2024 23:02:44 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=A2=9E=E5=8A=A0=E8=BF=9C=E7=A8=8B?=
 =?UTF-8?q?=E5=9B=BE=E7=89=87=E4=B8=8A=E4=BC=A0=E6=97=B6=E7=9A=84=E5=9F=9F?=
 =?UTF-8?q?=E5=90=8D=E9=AA=8C=E8=AF=81=EF=BC=8C=E9=99=90=E5=88=B6=E5=85=81?=
 =?UTF-8?q?=E8=AE=B8=E7=9A=84=E5=A4=B4=E5=83=8F=E5=9F=9F=E5=90=8D=20#1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/service/api/user.go | 31 ++++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/internal/service/api/user.go b/internal/service/api/user.go
index 8c81774..a13ec7f 100644
--- a/internal/service/api/user.go
+++ b/internal/service/api/user.go
@@ -9,15 +9,21 @@ import (
 	"github.com/tbxark/sphere/internal/pkg/database/ent/user"
 	"github.com/tbxark/sphere/pkg/server/statuserr"
 	"github.com/tbxark/sphere/pkg/storage"
+	"net/url"
 	"strconv"
 	"strings"
 )
 
 var _ apiv1.UserServiceHTTPServer = (*Service)(nil)
 
+var wechatAvatarDomains = map[string]struct{}{
+	"thirdwx.qlogo.cn": {},
+}
+
 const RemoteImageMaxSize = 1024 * 1024 * 2
 
 var ErrImageSizeExceed = fmt.Errorf("image size exceed")
+var ErrImageHostNotAllowed = fmt.Errorf("image host not allowed")
 
 func (s *Service) BindPhoneWxMini(ctx context.Context, req *apiv1.BindPhoneWxMiniRequest) (*apiv1.BindPhoneWxMiniResponse, error) {
 	userId, err := s.GetCurrentID(ctx)
@@ -87,20 +93,31 @@ func (s *Service) Update(ctx context.Context, req *apiv1.UpdateRequest) (*apiv1.
 	}, nil
 }
 
-func (s *Service) uploadRemoteImage(ctx context.Context, url string) (string, error) {
-	key := s.Storage.ExtractKeyFromURL(url)
-	if key == "" {
+func (s *Service) uploadRemoteImage(ctx context.Context, uri string) (string, error) {
+	key, err := s.Storage.ExtractKeyFromURLWithMode(uri, false)
+	if key != "" && err == nil {
 		return key, nil
 	}
-	if !(strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://")) {
-		return key, nil
+	u, err := url.Parse(uri)
+	if err != nil {
+		return "", err
+	}
+	isValidHost := false
+	for domain := range wechatAvatarDomains {
+		if strings.HasSuffix(u.Host, domain) {
+			isValidHost = true
+			break
+		}
+	}
+	if !isValidHost {
+		return "", ErrImageHostNotAllowed
 	}
 	id, err := s.GetCurrentID(ctx)
 	if err != nil {
 		return "", err
 	}
-	key = storage.DefaultKeyBuilder(strconv.Itoa(int(id)))(url, "user")
-	resp, err := s.httpClient.Get(url)
+	key = storage.DefaultKeyBuilder(strconv.Itoa(int(id)))(uri, "user")
+	resp, err := s.httpClient.Get(uri)
 	if err != nil {
 		return "", err
 	}