From 7c813f85796eed1373ddf8864b5dc9bb6d36e634 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: Tue, 1 Aug 2023 11:09:04 +0800
Subject: [PATCH] feat: add config sign-server-timeout

---
 cmd/gocq/login.go                 | 14 ++++-----
 cmd/gocq/main.go                  |  4 +--
 internal/base/flag.go             |  4 ++-
 internal/download/download.go     | 52 ++++++++++++++++++++-----------
 modules/config/config.go          |  1 +
 modules/config/default_config.yml |  6 ++--
 6 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/cmd/gocq/login.go b/cmd/gocq/login.go
index a6f6e88..2823ae4 100644
--- a/cmd/gocq/login.go
+++ b/cmd/gocq/login.go
@@ -278,7 +278,7 @@ func energy(uin uint64, id string, _ string, salt []byte) ([]byte, error) {
 		Method: http.MethodGet,
 		URL: signServer + "custom_energy" + fmt.Sprintf("?data=%v&salt=%v&uin=%v&android_id=%v&guid=%v",
 			id, hex.EncodeToString(salt), uin, utils.B2S(device.AndroidId), hex.EncodeToString(device.Guid)),
-	}
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second)
 	if base.IsBelow110 {
 		req.URL = signServer + "custom_energy" + fmt.Sprintf("?data=%v&salt=%v", id, hex.EncodeToString(salt))
 	}
@@ -312,7 +312,7 @@ func signSubmit(uin string, cmd string, callbackID int64, buffer []byte, t strin
 		Method: http.MethodGet,
 		URL: signServer + "submit" + fmt.Sprintf("?uin=%v&cmd=%v&callback_id=%v&buffer=%v",
 			uin, cmd, callbackID, buffStr),
-	}.Bytes()
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second).Bytes()
 	if err != nil {
 		log.Warnf("提交 callback 时出现错误: %v server: %v", err, signServer)
 	}
@@ -343,7 +343,7 @@ func signRequset(seq uint64, uin string, cmd string, qua string, buff []byte) (s
 		Header: map[string]string{"Content-Type": "application/x-www-form-urlencoded"},
 		Body: bytes.NewReader([]byte(fmt.Sprintf("uin=%v&qua=%s&cmd=%s&seq=%v&buffer=%v&android_id=%v&guid=%v",
 			uin, qua, cmd, seq, hex.EncodeToString(buff), utils.B2S(device.AndroidId), hex.EncodeToString(device.Guid)))),
-	}.Bytes()
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second).Bytes()
 	if err != nil {
 		return nil, nil, nil, err
 	}
@@ -371,7 +371,7 @@ func signRegister(uin int64, androidID, guid []byte, qimei36, key string) {
 		Method: http.MethodGet,
 		URL: signServer + "register" + fmt.Sprintf("?uin=%v&android_id=%v&guid=%v&qimei36=%v&key=%s",
 			uin, utils.B2S(androidID), hex.EncodeToString(guid), qimei36, key),
-	}.Bytes()
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second).Bytes()
 	if err != nil {
 		log.Warnf("注册QQ实例时出现错误: %v server: %v", err, signServer)
 		return
@@ -393,7 +393,7 @@ func signRefreshToken(uin string) error {
 	resp, err := download.Request{
 		Method: http.MethodGet,
 		URL:    signServer + "request_token" + fmt.Sprintf("?uin=%v", uin),
-	}.Bytes()
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second).Bytes()
 	if err != nil {
 		return err
 	}
@@ -463,7 +463,7 @@ func signServerDestroy(uin string) error {
 	resp, err := download.Request{
 		Method: http.MethodGet,
 		URL:    signServer + "destroy" + fmt.Sprintf("?uin=%v&key=%v", uin, base.Key),
-	}.Bytes()
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second).Bytes()
 	if err != nil || gjson.GetBytes(resp, "code").Int() != 0 {
 		return errors.Wrapf(err, "destroy 实例出现错误, server: %v", signServer)
 	}
@@ -475,7 +475,7 @@ func signVersion() (version string, err error) {
 	resp, err := download.Request{
 		Method: http.MethodGet,
 		URL:    signServer,
-	}.Bytes()
+	}.WithTimeout(time.Duration(base.SignServerTimeout) * time.Second).Bytes()
 	if err != nil {
 		return "", err
 	}
diff --git a/cmd/gocq/main.go b/cmd/gocq/main.go
index 9204e15..7dbaff1 100644
--- a/cmd/gocq/main.go
+++ b/cmd/gocq/main.go
@@ -166,7 +166,6 @@ func LoginInteract() {
 
 	if base.SignServer != "-" && base.SignServer != "" {
 		log.Infof("使用服务器 %s 进行数据包签名", base.SignServer)
-		download.SetTimeout(time.Duration(base.HTTPTimeout) * time.Second) // 设置签名超时时间
 		// 等待签名服务器直到连接成功
 		if !signWaitServer() {
 			log.Fatalf("连接签名服务器失败")
@@ -304,9 +303,9 @@ func LoginInteract() {
 		cli.Uin = base.Account.Uin
 		cli.PasswordMd5 = base.PasswordHash
 	}
+	download.SetTimeout(time.Duration(base.HTTPTimeout) * time.Second)
 	if !base.FastStart {
 		log.Infof("正在检查协议更新...")
-		download.SetTimeout(time.Second * 5) // 防止协议更新堵塞过久
 		currentVersionName := device.Protocol.Version().SortVersionName
 		remoteVersion, err := getRemoteLatestProtocolVersion(int(device.Protocol.Version().Protocol))
 		if err == nil {
@@ -391,7 +390,6 @@ func LoginInteract() {
 	})
 	saveToken()
 	cli.AllowSlider = true
-	download.SetTimeout(time.Duration(base.HTTPTimeout) * time.Second) // 登陆完成进行最终超时设置
 	log.Infof("登录成功 欢迎使用: %v", cli.Nickname)
 	log.Info("开始加载好友列表...")
 	global.Check(cli.ReloadFriendList(), true)
diff --git a/internal/base/flag.go b/internal/base/flag.go
index 33ec62d..97e5128 100644
--- a/internal/base/flag.go
+++ b/internal/base/flag.go
@@ -41,7 +41,8 @@ var (
 	SignServer          string // 使用特定的服务器进行签名
 	Key                 string // 签名服务器密钥
 	IsBelow110          bool   // 签名服务器版本是否低于1.1.0及以下
-	HTTPTimeout         int
+	HTTPTimeout         int    // download 超时时间
+	SignServerTimeout   int    // 签名服务器超时时间
 
 	PostFormat        string                 // 上报格式 string or array
 	Proxy             string                 // 存储 proxy_rewrite,用于设置代理
@@ -94,6 +95,7 @@ func Init() {
 		Key = conf.Account.Key
 		IsBelow110 = conf.Account.IsBelow110
 		HTTPTimeout = conf.Message.HTTPTimeout
+		SignServerTimeout = conf.Message.SignServerTimeout
 	}
 	{ // others
 		Proxy = conf.Message.ProxyRewrite
diff --git a/internal/download/download.go b/internal/download/download.go
index ac5ae03..786658f 100644
--- a/internal/download/download.go
+++ b/internal/download/download.go
@@ -21,20 +21,7 @@ import (
 	"github.com/Mrs4s/go-cqhttp/internal/base"
 )
 
-var client = &http.Client{
-	Transport: &http.Transport{
-		Proxy: func(request *http.Request) (*url.URL, error) {
-			if base.Proxy == "" {
-				return http.ProxyFromEnvironment(request)
-			}
-			return url.Parse(base.Proxy)
-		},
-		// Disable http2
-		TLSNextProto:        map[string]func(authority string, c *tls.Conn) http.RoundTripper{},
-		MaxIdleConnsPerHost: 999,
-	},
-	Timeout: time.Minute,
-}
+var client = newcli(time.Second * 15)
 
 var clienth2 = &http.Client{
 	Transport: &http.Transport{
@@ -50,12 +37,35 @@ var clienth2 = &http.Client{
 	Timeout: time.Second * 15,
 }
 
+func newcli(t time.Duration) *http.Client {
+	return &http.Client{
+		Transport: &http.Transport{
+			Proxy: func(request *http.Request) (*url.URL, error) {
+				if base.Proxy == "" {
+					return http.ProxyFromEnvironment(request)
+				}
+				return url.Parse(base.Proxy)
+			},
+			// Disable http2
+			TLSNextProto:        map[string]func(authority string, c *tls.Conn) http.RoundTripper{},
+			MaxIdleConnsPerHost: 999,
+		},
+		Timeout: t,
+	}
+}
+
 // ErrOverSize 响应主体过大时返回此错误
 var ErrOverSize = errors.New("oversize")
 
 // UserAgent HTTP请求时使用的UA
 const UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/87.0.4280.88 Safari/537.36 Edg/87.0.664.66"
 
+// WithTimeout get a download instance with timeout t
+func (r Request) WithTimeout(t time.Duration) *Request {
+	r.custcli = newcli(t)
+	return &r
+}
+
 // SetTimeout set internal/download client timeout
 func SetTimeout(t time.Duration) {
 	if t == 0 {
@@ -67,14 +77,18 @@ func SetTimeout(t time.Duration) {
 
 // Request is a file download request
 type Request struct {
-	Method string
-	URL    string
-	Header map[string]string
-	Limit  int64
-	Body   io.Reader
+	Method  string
+	URL     string
+	Header  map[string]string
+	Limit   int64
+	Body    io.Reader
+	custcli *http.Client
 }
 
 func (r Request) client() *http.Client {
+	if r.custcli != nil {
+		return r.custcli
+	}
 	if strings.Contains(r.URL, "go-cqhttp.org") {
 		return clienth2
 	}
diff --git a/modules/config/config.go b/modules/config/config.go
index f424054..5a02a7d 100644
--- a/modules/config/config.go
+++ b/modules/config/config.go
@@ -63,6 +63,7 @@ type Config struct {
 		SkipMimeScan        bool   `yaml:"skip-mime-scan"`
 		ConvertWebpImage    bool   `yaml:"convert-webp-image"`
 		HTTPTimeout         int    `yaml:"http-timeout"`
+		SignServerTimeout   int    `yaml:"sign-server-timeout"`
 	} `yaml:"message"`
 
 	Output struct {
diff --git a/modules/config/default_config.yml b/modules/config/default_config.yml
index df2244b..b8f9b6d 100644
--- a/modules/config/default_config.yml
+++ b/modules/config/default_config.yml
@@ -70,8 +70,10 @@ message:
   skip-mime-scan: false
   # 是否自动转换 WebP 图片
   convert-webp-image: false
-  # http超时时间，同时作为签名服务超时时间
-  http-timeout: 0
+  # download 超时时间(s)
+  http-timeout: 15
+  # 签名服务超时时间(s)
+  sign-server-timeout: 60
 
 output:
   # 日志等级 trace,debug,info,warn,error