From 0ffd8b27cbeb8b6ef135dcac372bd6a67a098612 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 5 May 2024 03:32:09 +0900
Subject: [PATCH] fix(pool): image pooling logic

---
 img/pool/img.go    | 39 +++++++++++++++++----------------------
 img/pool/sender.go | 10 ++++++++--
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/img/pool/img.go b/img/pool/img.go
index 6a680ca..d4aa16e 100644
--- a/img/pool/img.go
+++ b/img/pool/img.go
@@ -13,6 +13,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/wdvxdr1123/ZeroBot/message"
 
+	"github.com/FloatTech/floatbox/web"
 	"github.com/FloatTech/zbputils/ctxext"
 )
 
@@ -41,20 +42,15 @@ func GetImage(name string) (m *Image, err error) {
 	m.n = name
 	m.item, err = getItem(name)
 	if err == nil && m.u != "" {
-		var resp *http.Response
-		resp, err = http.Head(m.String())
+		_, err = web.RequestDataWithHeaders(web.NewDefaultClient(), m.String(), "HEAD", func(r *http.Request) error {
+			r.Header.Set("user-agent", web.RandUA())
+			return nil
+		}, nil)
 		if err == nil {
-			_ = resp.Body.Close()
-		} else {
-			goto OUTDATE
-		}
-		if resp.StatusCode != http.StatusOK {
-			goto OUTDATE
+			return
 		}
-		return
-	OUTDATE:
+		logrus.Debugln("[imgpool] image", name, m, "outdated:", err)
 		err = ErrImgFileOutdated
-		logrus.Debugln("[imgpool] image", name, m, "outdated")
 		return
 	}
 	err = ErrNoSuchImg
@@ -63,24 +59,22 @@ func GetImage(name string) (m *Image, err error) {
 }
 
 // NewImage context name file
-func NewImage(send ctxext.NoCtxSendMsg, get ctxext.NoCtxGetMsg, name, f string) (m *Image, err error) {
+func NewImage(send ctxext.NoCtxSendMsg, get ctxext.NoCtxGetMsg, name, f string) (m *Image, hassent bool, err error) {
 	m = new(Image)
 	m.n = name
 	m.SetFile(f)
 	m.item, err = getItem(name)
 	if err == nil && m.item.u != "" {
-		var resp *http.Response
-		resp, err = http.Head(m.String())
+		_, err = web.RequestDataWithHeaders(web.NewDefaultClient(), m.String(), "HEAD", func(r *http.Request) error {
+			r.Header.Set("user-agent", web.RandUA())
+			return nil
+		}, nil)
 		if err == nil {
-			_ = resp.Body.Close()
-			if resp.StatusCode == http.StatusOK {
-				return
-			}
+			return
 		}
-		logrus.Debugln("[imgpool] image", name, m, "outdated, updating...")
-		get = nil
+		logrus.Debugln("[imgpool] image", name, m, "outdated:", err, "updating...")
 	}
-	err = m.Push(send, get)
+	hassent, err = m.Push(send, get)
 	return
 }
 
@@ -109,12 +103,13 @@ func (m *Image) SetFile(f string) {
 }
 
 // Push context
-func (m *Image) Push(send ctxext.NoCtxSendMsg, get ctxext.NoCtxGetMsg) (err error) {
+func (m *Image) Push(send ctxext.NoCtxSendMsg, get ctxext.NoCtxGetMsg) (hassent bool, err error) {
 	id := send(message.Message{message.Image(m.f)})
 	if id == 0 {
 		err = ErrSendImg
 		return
 	}
+	hassent = true
 	if get == nil {
 		return
 	}
diff --git a/img/pool/sender.go b/img/pool/sender.go
index 0378901..7799714 100644
--- a/img/pool/sender.go
+++ b/img/pool/sender.go
@@ -23,7 +23,10 @@ func SendImageFromPool(imgname, imgpath string, genimg func() error, send ctxext
 			}
 		}
 		m.SetFile(file.BOTPATH + "/" + imgpath)
-		return m.Push(send, get)
+		hassent, err := m.Push(send, get)
+		if hassent {
+			return err
+		}
 	}
 	// 发送图片
 	img := message.Image(m.String())
@@ -46,7 +49,10 @@ func SendRemoteImageFromPool(imgname, imgurl string, send ctxext.NoCtxSendMsg, g
 		if err == ErrImgFileOutdated {
 			get = nil
 		}
-		return m.Push(send, get)
+		hassent, err := m.Push(send, get)
+		if hassent {
+			return err
+		}
 	}
 	// 发送图片
 	img := message.Image(m.String())