From 5bf3f7a2d9957f28cc267d34c005af1274d98d94 Mon Sep 17 00:00:00 2001
From: ZacharyZcR <PayasoNorahC@protonmail.com>
Date: Tue, 2 Sep 2025 12:02:27 +0000
Subject: [PATCH] =?UTF-8?q?refactor:=20=E4=B8=BASMB2=E6=8F=92=E4=BB=B6?=
 =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E5=8C=85=E6=8E=A7=E5=88=B6=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 在testCredential和detectProtocol方法中添加包控制检查
- 统一错误处理和包计数逻辑
- 确保所有TCP连接遵循发包限制
---
 plugins/services/smb2.go | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/plugins/services/smb2.go b/plugins/services/smb2.go
index 9a05eaa..fba2b94 100644
--- a/plugins/services/smb2.go
+++ b/plugins/services/smb2.go
@@ -105,6 +105,11 @@ func (p *Smb2Plugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResul
 
 // testCredential 测试单个凭据 - 简化的SMB2协议检测
 func (p *Smb2Plugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
+	// 检查发包限制
+	if canSend, _ := common.CanSendPacket(); !canSend {
+		return false
+	}
+	
 	// 基于TCP连接的简单SMB2检测
 	timeoutCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
 	defer cancel()
@@ -112,8 +117,10 @@ func (p *Smb2Plugin) testCredential(ctx context.Context, info *common.HostInfo,
 	dialer := &net.Dialer{}
 	conn, err := dialer.DialContext(timeoutCtx, "tcp", fmt.Sprintf("%s:445", info.Host))
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return false
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	// 发送SMB2协商请求
@@ -153,10 +160,17 @@ func (p *Smb2Plugin) testCredential(ctx context.Context, info *common.HostInfo,
 
 // detectProtocol 检测SMB2协议信息
 func (p *Smb2Plugin) detectProtocol(ctx context.Context, info *common.HostInfo) string {
-	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:445", info.Host), 5*time.Second)
-	if err != nil {
+	// 检查发包限制
+	if canSend, _ := common.CanSendPacket(); !canSend {
 		return ""
 	}
+	
+	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:445", info.Host), 5*time.Second)
+	if err != nil {
+		common.IncrementTCPFailedPacketCount()
+		return ""
+	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	// 发送SMB2协商请求获取版本信息