From b243d6c92e2e2f7e5f2357d21235f6a10fe9c252 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Sat, 26 Mar 2022 12:16:25 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=BB=E5=87=BB=E6=B5=81=E9=87=8F=E7=BB=9F?=
 =?UTF-8?q?=E8=AE=A1=E4=BD=BF=E7=94=A8=E4=B8=8A=E8=A1=8C=E6=B5=81=E9=87=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/http_request.go | 44 +++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/internal/nodes/http_request.go b/internal/nodes/http_request.go
index 6fcd203..26c0d0b 100644
--- a/internal/nodes/http_request.go
+++ b/internal/nodes/http_request.go
@@ -312,17 +312,23 @@ func (this *HTTPRequest) doEnd() {
 
 	// 流量统计
 	// TODO 增加是否开启开关
-	// TODO 增加Header统计，考虑从Conn中读取
 	if this.ReqServer != nil {
+		var countCached int64 = 0
+		var cachedBytes int64 = 0
+
+		var countAttacks int64 = 0
+		var attackBytes int64 = 0
+
 		if this.isCached {
-			stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), this.writer.SentBodyBytes(), 1, 1, 0, 0, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
-		} else {
-			if this.isAttack {
-				stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), 0, 1, 0, 1, this.writer.SentBodyBytes(), this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
-			} else {
-				stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), 0, 1, 0, 0, 0, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
-			}
+			countCached = 1
+			cachedBytes = this.writer.SentBodyBytes()
 		}
+		if this.isAttack {
+			countAttacks = 1
+			attackBytes = this.CalculateSize()
+		}
+
+		stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes(), cachedBytes, 1, countCached, countAttacks, attackBytes, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
 	}
 
 	// 指标
@@ -1278,6 +1284,28 @@ func (this *HTTPRequest) ContentLength() int64 {
 	return this.RawReq.ContentLength
 }
 
+// CalculateSize 计算当前请求的尺寸（预估）
+func (this *HTTPRequest) CalculateSize() (size int64) {
+	// Get /xxx HTTP/1.1
+	size += int64(len(this.RawReq.Method)) + 1
+	size += int64(len(this.RawReq.URL.String())) + 1
+	size += int64(len(this.RawReq.Proto)) + 1
+	for k, v := range this.RawReq.Header {
+		for _, v1 := range v {
+			size += int64(len(k) + 2 /** : **/ + len(v1))
+		}
+	}
+
+	size += 1 /** \r\n **/
+
+	if this.RawReq.ContentLength > 0 {
+		size += this.RawReq.ContentLength
+	} else if len(this.requestBodyData) > 0 {
+		size += int64(len(this.requestBodyData))
+	}
+	return size
+}
+
 // Method 请求方法
 func (this *HTTPRequest) Method() string {
 	return this.RawReq.Method