From 59bf13467bd15b68ccd9934586d6e8a3d788d493 Mon Sep 17 00:00:00 2001
From: Mmx233 <admin@mmxblog.com>
Date: Sun, 3 Dec 2023 20:34:57 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E8=87=AA=E5=AE=9A?=
 =?UTF-8?q?=E4=B9=89=E8=AF=B7=E6=B1=82=20header?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/srun/api.go  | 35 +++++++++++++++++++++++++++--------
 pkg/srun/srun.go | 12 +++++++++---
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/pkg/srun/api.go b/pkg/srun/api.go
index da8c793..641f5ad 100644
--- a/pkg/srun/api.go
+++ b/pkg/srun/api.go
@@ -17,18 +17,29 @@ type Api struct {
 	Client  *http.Client
 	// 禁用自动重定向
 	NoDirect *http.Client
+
+	CustomHeader map[string]interface{}
 }
 
-func (a *Api) Init(https bool, domain string, client *http.Client) {
+type ApiConfig struct {
+	Https        bool
+	Domain       string
+	Client       *http.Client
+	CustomHeader map[string]interface{}
+}
+
+func (a *Api) Init(conf *ApiConfig) {
 	a.BaseUrl = "http"
-	if https {
+	if conf.Https {
 		a.BaseUrl += "s"
 	}
-	a.BaseUrl = a.BaseUrl + "://" + domain + "/"
+	a.BaseUrl = a.BaseUrl + "://" + conf.Domain + "/"
+
+	a.CustomHeader = conf.CustomHeader
 
 	// 初始化 http client
-	a.Client = client
-	copyClient := *client
+	a.Client = conf.Client
+	copyClient := *conf.Client
 	a.NoDirect = &copyClient
 	a.NoDirect.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
 		return http.ErrUseLastResponse
@@ -46,8 +57,9 @@ func (a *Api) request(path string, query map[string]interface{}) (map[string]int
 	query["_"] = timestamp
 	httpTool := tool.NewHttpTool(a.Client)
 	req, err := httpTool.GenReq("GET", &tool.DoHttpReq{
-		Url:   a.BaseUrl + path,
-		Query: query,
+		Url:    a.BaseUrl + path,
+		Query:  query,
+		Header: a.CustomHeader,
 	})
 	if err != nil {
 		log.Debugln(err)
@@ -84,7 +96,14 @@ func (a *Api) DetectAcid() (string, error) {
 	addr := a.BaseUrl
 	for {
 		log.Debugln("HTTP GET ", addr)
-		res, err := a.NoDirect.Get(addr)
+		req, err := http.NewRequest("GET", addr, nil)
+		if err != nil {
+			return "", err
+		}
+		for k, v := range a.CustomHeader {
+			req.Header.Set(k, fmt.Sprint(v))
+		}
+		res, err := a.NoDirect.Do(req)
 		if err != nil {
 			return "", err
 		}
diff --git a/pkg/srun/srun.go b/pkg/srun/srun.go
index d8b6e19..8bd2c8d 100644
--- a/pkg/srun/srun.go
+++ b/pkg/srun/srun.go
@@ -12,15 +12,21 @@ type Conf struct {
 	//调用 API 时直接访问 https URL
 	Https bool
 	//登录参数，不可缺省
-	LoginInfo LoginInfo
-	Client    *http.Client
+	LoginInfo    LoginInfo
+	Client       *http.Client
+	CustomHeader map[string]interface{}
 }
 
 func New(conf *Conf) *Srun {
 	srun := &Srun{
 		LoginInfo: conf.LoginInfo,
 	}
-	srun.api.Init(conf.Https, conf.LoginInfo.Form.Domain, conf.Client)
+	srun.api.Init(&ApiConfig{
+		Https:        conf.Https,
+		Domain:       conf.LoginInfo.Form.Domain,
+		Client:       conf.Client,
+		CustomHeader: conf.CustomHeader,
+	})
 	return srun
 }