From 08ce95b43dc936154863399a67ca2384593188a6 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Tue, 4 Mar 2025 14:58:28 +0800
Subject: [PATCH] fix: check pool net set headers

---
 core/pool/brutepool.go | 8 --------
 core/pool/checkpool.go | 2 +-
 core/runner.go         | 8 ++++++++
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go
index c32701d..7460a0c 100644
--- a/core/pool/brutepool.go
+++ b/core/pool/brutepool.go
@@ -107,14 +107,6 @@ type BrutePool struct {
 }
 
 func (pool *BrutePool) Init() error {
-	if pool.Headers.Get("User-Agent") == "" {
-		pool.Headers.Set("User-Agent", pkg.DefaultUserAgent)
-	}
-
-	if pool.Headers.Get("Accept") == "" {
-		pool.Headers.Set("Accept", "*/*")
-	}
-
 	pool.initwg.Add(2)
 	if pool.Index != "/" {
 		logs.Log.Logf(pkg.LogVerbose, "custom index url: %s", pkg.BaseURL(pool.url)+pkg.FormatURL(pkg.BaseURL(pool.url), pool.Index))
diff --git a/core/pool/checkpool.go b/core/pool/checkpool.go
index b6040e2..136d460 100644
--- a/core/pool/checkpool.go
+++ b/core/pool/checkpool.go
@@ -36,7 +36,7 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) {
 			processCh:  make(chan *baseline.Baseline, config.Thread),
 		},
 	}
-	pool.Headers = map[string]string{"Connection": "close"}
+	pool.Headers.Set("Connection", "close")
 	p, _ := ants.NewPoolWithFunc(config.Thread, pool.Invoke)
 
 	pool.Pool = p
diff --git a/core/runner.go b/core/runner.go
index e36821d..99b9809 100644
--- a/core/runner.go
+++ b/core/runner.go
@@ -114,6 +114,14 @@ func (r *Runner) PrepareConfig() *pool.Config {
 		config.Headers.Set(k, v)
 	}
 
+	if config.Headers.Get("User-Agent") == "" {
+		config.Headers.Set("User-Agent", pkg.DefaultUserAgent)
+	}
+
+	if config.Headers.Get("Accept") == "" {
+		config.Headers.Set("Accept", "*/*")
+	}
+
 	return config
 }