From 5dc8f7107fa9485dc0a1de49ea7d224ede3c07ba Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Wed, 4 Jan 2023 11:31:31 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96content-length=E8=BE=93?=
 =?UTF-8?q?=E5=87=BA,=20=E7=8E=B0=E5=9C=A8=E8=B6=85=E8=BF=87max=E7=9A=84?=
 =?UTF-8?q?=E7=BB=93=E6=9E=9C=E4=B9=9F=E4=BC=9A=E6=AD=A3=E7=A1=AE=E7=9A=84?=
 =?UTF-8?q?=E8=8E=B7=E5=8F=96,=20=E4=BD=86=E6=B7=BB=E5=8A=A0=E4=BA=86`(exc?=
 =?UTF-8?q?eed)`=E6=A0=87=E8=AE=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/pool.go |  3 +++
 pkg/baseline.go  | 65 +++++++++++++++++++++++++++---------------------
 2 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/internal/pool.go b/internal/pool.go
index 4d7106b..51b6189 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -92,6 +92,9 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			}
 		}
 
+		if bl.BodyLength > ihttp.DefaultMaxBodySize {
+			bl.ExceedLength = true
+		}
 		bl.Source = int(unit.source)
 		bl.ReqDepth = unit.depth
 		bl.Spended = time.Since(start).Milliseconds()
diff --git a/pkg/baseline.go b/pkg/baseline.go
index 1f752b2..c5c7639 100644
--- a/pkg/baseline.go
+++ b/pkg/baseline.go
@@ -28,7 +28,7 @@ func NewBaseline(u, host string, resp *ihttp.Response) *Baseline {
 		bl.Host = host
 	}
 	bl.Body = resp.Body()
-	bl.BodyLength = len(bl.Body)
+	bl.BodyLength = resp.ContentLength()
 	bl.Header = resp.Header()
 	bl.HeaderLength = len(bl.Header)
 	bl.RedirectURL = resp.GetHeader("Location")
@@ -55,7 +55,7 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response, reason string) *Ba
 	}
 
 	bl.Body = resp.Body()
-	bl.BodyLength = len(bl.Body)
+	bl.BodyLength = resp.ContentLength()
 	bl.Header = resp.Header()
 	bl.HeaderLength = len(bl.Header)
 	bl.RedirectURL = string(resp.GetHeader("Location"))
@@ -64,33 +64,34 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response, reason string) *Ba
 }
 
 type Baseline struct {
-	Number       int        `json:"number"`
-	Url          *url.URL   `json:"-"`
-	UrlString    string     `json:"url"`
-	Path         string     `json:"path"`
-	Host         string     `json:"host"`
-	Body         []byte     `json:"-"`
-	BodyLength   int        `json:"body_length"`
-	Header       []byte     `json:"-"`
-	Raw          []byte     `json:"-"`
-	HeaderLength int        `json:"header_length"`
-	RedirectURL  string     `json:"redirect_url,omitempty"`
-	FrontURL     string     `json:"front_url,omitempty"`
-	Status       int        `json:"status"`
-	Spended      int64      `json:"spend"` // 耗时, 毫秒
-	Title        string     `json:"title"`
-	Frameworks   Frameworks `json:"frameworks"`
-	Extracteds   Extracteds `json:"extracts"`
-	ErrString    string     `json:"error"`
-	Reason       string     `json:"reason"`
-	IsValid      bool       `json:"valid"`
-	IsFuzzy      bool       `json:"fuzzy"`
-	URLs         []string   `json:"-"`
-	Source       int        `json:"source"`
-	RecuDepth    int        `json:"-"`
-	ReqDepth     int        `json:"depth"`
-	Recu         bool       `json:"-"`
-	*parsers.Hashes
+	Number          int        `json:"number"`
+	Url             *url.URL   `json:"-"`
+	UrlString       string     `json:"url"`
+	Path            string     `json:"path"`
+	Host            string     `json:"host"`
+	Body            []byte     `json:"-"`
+	BodyLength      int        `json:"body_length"`
+	ExceedLength    bool       `json:"-"`
+	Header          []byte     `json:"-"`
+	Raw             []byte     `json:"-"`
+	HeaderLength    int        `json:"header_length"`
+	RedirectURL     string     `json:"redirect_url,omitempty"`
+	FrontURL        string     `json:"front_url,omitempty"`
+	Status          int        `json:"status"`
+	Spended         int64      `json:"spend"` // 耗时, 毫秒
+	Title           string     `json:"title"`
+	Frameworks      Frameworks `json:"frameworks"`
+	Extracteds      Extracteds `json:"extracts"`
+	ErrString       string     `json:"error"`
+	Reason          string     `json:"reason"`
+	IsValid         bool       `json:"valid"`
+	IsFuzzy         bool       `json:"fuzzy"`
+	Source          int        `json:"source"`
+	ReqDepth        int        `json:"depth"`
+	Recu            bool       `json:"-"`
+	RecuDepth       int        `json:"-"`
+	URLs            []string   `json:"-"`
+	*parsers.Hashes `json:"hashes"`
 }
 
 func (bl *Baseline) IsDir() bool {
@@ -292,6 +293,9 @@ func (bl *Baseline) ColorString() string {
 	line.WriteString(logs.GreenBold(strconv.Itoa(bl.Status)))
 	line.WriteString(" - ")
 	line.WriteString(logs.YellowBold(strconv.Itoa(bl.BodyLength)))
+	if bl.ExceedLength {
+		line.WriteString(logs.Red("(exceed)"))
+	}
 	line.WriteString(" - ")
 	line.WriteString(logs.YellowBold(strconv.Itoa(int(bl.Spended)) + "ms"))
 	line.WriteString(logs.GreenLine(bl.Additional("title")))
@@ -339,6 +343,9 @@ func (bl *Baseline) String() string {
 	line.WriteString(strconv.Itoa(bl.Status))
 	line.WriteString(" - ")
 	line.WriteString(strconv.Itoa(bl.BodyLength))
+	if bl.ExceedLength {
+		line.WriteString("(exceed)")
+	}
 	line.WriteString(" - ")
 	line.WriteString(strconv.Itoa(int(bl.Spended)) + "ms")
 	line.WriteString(bl.Additional("title"))