From c746c26ff9eba3a4e83b8693fe88557d2a101521 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Tue, 4 Mar 2025 14:56:48 +0800
Subject: [PATCH] fix: brute init not set header,
 https://github.com/chainreactors/spray/issues/94

---
 core/ihttp/request.go  | 10 +++++-----
 core/option.go         |  2 +-
 core/pool/brutepool.go |  8 ++++----
 core/pool/config.go    |  3 ++-
 core/runner.go         |  7 ++++++-
 5 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/core/ihttp/request.go b/core/ihttp/request.go
index 4e52912..1302e17 100644
--- a/core/ihttp/request.go
+++ b/core/ihttp/request.go
@@ -31,18 +31,18 @@ type Request struct {
 	ClientType      int
 }
 
-func (r *Request) SetHeaders(header map[string]string, RandomUA bool) {
+func (r *Request) SetHeaders(header http.Header, RandomUA bool) {
 	if RandomUA {
 		r.SetHeader("User-Agent", pkg.RandomUA())
 	}
 
 	if r.StandardRequest != nil {
-		for k, v := range header {
-			r.StandardRequest.Header.Set(k, v)
-		}
+		r.StandardRequest.Header = header
 	} else if r.FastRequest != nil {
 		for k, v := range header {
-			r.FastRequest.Header.Set(k, v)
+			for _, i := range v {
+				r.FastRequest.Header.Set(k, i)
+			}
 		}
 	}
 }
diff --git a/core/option.go b/core/option.go
index 7e790bb..14a38d1 100644
--- a/core/option.go
+++ b/core/option.go
@@ -99,7 +99,7 @@ type OutputOptions struct {
 
 type RequestOptions struct {
 	Method          string   `short:"X" long:"method" default:"GET" description:"String, request method, e.g.: --method POST" config:"method"`
-	Headers         []string `long:"header" description:"Strings, custom headers, e.g.: --header 'Auth: example_auth'" config:"headers"`
+	Headers         []string `short:"H" long:"header" description:"Strings, custom headers, e.g.: --header 'Auth: example_auth'" config:"headers"`
 	UserAgent       string   `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom" config:"useragent"`
 	RandomUserAgent bool     `long:"random-agent" description:"Bool, use random with default user-agent" config:"random-useragent"`
 	Cookie          []string `long:"cookie" description:"Strings, custom cookie" config:"cookies"`
diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go
index b051c80..c32701d 100644
--- a/core/pool/brutepool.go
+++ b/core/pool/brutepool.go
@@ -107,12 +107,12 @@ type BrutePool struct {
 }
 
 func (pool *BrutePool) Init() error {
-	if pool.Headers["User-Agent"] == "" {
-		pool.Headers["User-Agent"] = pkg.DefaultUserAgent
+	if pool.Headers.Get("User-Agent") == "" {
+		pool.Headers.Set("User-Agent", pkg.DefaultUserAgent)
 	}
 
-	if pool.Headers["Accept"] == "" {
-		pool.Headers["Accept"] = "*/*"
+	if pool.Headers.Get("Accept") == "" {
+		pool.Headers.Set("Accept", "*/*")
 	}
 
 	pool.initwg.Add(2)
diff --git a/core/pool/config.go b/core/pool/config.go
index 6a0c1f2..5178dfe 100644
--- a/core/pool/config.go
+++ b/core/pool/config.go
@@ -7,6 +7,7 @@ import (
 	"github.com/chainreactors/words"
 	"github.com/chainreactors/words/rule"
 	"github.com/expr-lang/expr/vm"
+	"net/http"
 	"sync"
 	"time"
 )
@@ -27,7 +28,7 @@ type Config struct {
 	BreakThreshold    int32
 	Method            string
 	Mod               SprayMod
-	Headers           map[string]string
+	Headers           http.Header
 	ClientType        int
 	MatchExpr         *vm.Program
 	FilterExpr        *vm.Program
diff --git a/core/runner.go b/core/runner.go
index 168c052..e36821d 100644
--- a/core/runner.go
+++ b/core/runner.go
@@ -15,6 +15,7 @@ import (
 	"github.com/panjf2000/ants/v2"
 	"github.com/vbauerster/mpb/v8"
 	"github.com/vbauerster/mpb/v8/decor"
+	"net/http"
 	"strings"
 	"sync"
 	"time"
@@ -67,7 +68,7 @@ func (r *Runner) PrepareConfig() *pool.Config {
 		Thread:         r.Threads,
 		Timeout:        time.Duration(r.Timeout) * time.Second,
 		RateLimit:      r.RateLimit,
-		Headers:        r.Headers,
+		Headers:        make(http.Header),
 		Method:         r.Method,
 		Mod:            pool.ModMap[r.Mod],
 		OutputCh:       r.outputCh,
@@ -109,6 +110,10 @@ func (r *Runner) PrepareConfig() *pool.Config {
 		}
 	}
 
+	for k, v := range r.Headers {
+		config.Headers.Set(k, v)
+	}
+
 	return config
 }