From e3fc74e78d264e557e42627cdc2f3207243ee559 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Fri, 6 Jan 2023 15:06:40 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96baseline=E7=9A=84=E6=80=A7?=
 =?UTF-8?q?=E8=83=BD=E4=B8=8E=E4=BF=A1=E6=81=AF=E8=8E=B7=E5=8F=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/pool.go |  2 +-
 pkg/baseline.go  | 22 +++++++++++++---------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/internal/pool.go b/internal/pool.go
index 1e7162a..36c8602 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -49,8 +49,8 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 	}
 
 	p, _ := ants.NewPoolWithFunc(config.Thread, pool.Invoke)
-
 	pool.reqPool = p
+
 	// 挂起一个异步的处理结果线程, 不干扰主线程的请求并发
 	go func() {
 		for bl := range pool.tempCh {
diff --git a/pkg/baseline.go b/pkg/baseline.go
index 8945fba..7760c94 100644
--- a/pkg/baseline.go
+++ b/pkg/baseline.go
@@ -55,14 +55,21 @@ func NewBaseline(u, host string, resp *ihttp.Response) *Baseline {
 	if resp.ClientType == ihttp.STANDARD {
 		bl.Host = host
 	}
+	header := resp.Header()
+	bl.Header = make([]byte, len(header))
+	copy(bl.Header, header)
+	bl.HeaderLength = len(bl.Header)
+
 	body := resp.Body()
 	bl.Body = make([]byte, len(body))
 	copy(bl.Body, body)
 	bl.BodyLength = resp.ContentLength()
-	bl.Header = resp.Header()
-	bl.HeaderLength = len(bl.Header)
-	bl.RedirectURL = resp.GetHeader("Location")
+	if bl.BodyLength == -1 {
+		bl.BodyLength = len(bl.Body)
+	}
+
 	bl.Raw = append(bl.Header, bl.Body...)
+	bl.RedirectURL = resp.GetHeader("Location")
 	return bl
 }
 
@@ -85,12 +92,9 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response, reason string) *Ba
 		bl.Host = host
 	}
 
-	body := resp.Body()
-	bl.Body = make([]byte, len(body))
-	copy(bl.Body, body)
+	// 无效数据也要读取body, 否则keep-alive不生效
+	resp.Body()
 	bl.BodyLength = resp.ContentLength()
-	bl.Header = resp.Header()
-	bl.HeaderLength = len(bl.Header)
 	bl.RedirectURL = string(resp.GetHeader("Location"))
 
 	return bl
@@ -339,8 +343,8 @@ func (bl *Baseline) ColorString() string {
 	}
 	line.WriteString(" - ")
 	line.WriteString(logs.YellowBold(strconv.Itoa(int(bl.Spended)) + "ms"))
+	line.WriteString(logs.YellowBold(" - " + GetSourceName(bl.Source)))
 	line.WriteString(logs.GreenLine(bl.Additional("title")))
-	line.WriteString(logs.GreenLine(bl.Additional("source")))
 	line.WriteString(logs.Cyan(bl.Frameworks.String()))
 	line.WriteString(logs.Cyan(bl.Extracteds.String()))
 	if bl.RedirectURL != "" {