Skip to content

Commit 7fd08b7

Browse files
ayanrajpoot10Mzack9999
ayanrajpoot10
and
Mzack9999
authored
Add support for duplicate custom request headers (#2504)
* refactor: change CustomHeaders to support multiple values per header * fix: clear existing header values before setting custom headers * harden headers * expand tests --------- Co-authored-by: Mzack9999 <mzack9999@protonmail.com>
1 parent 8e7c09a commit 7fd08b7

6 files changed

Lines changed: 204 additions & 25 deletions

File tree

‎common/httputilz/httputilz.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@ func DumpRequest(req *retryablehttp.Request) (string, error) {
2424
}
2525

2626
// ParseRequest from raw string
27-
func ParseRequest(req string, unsafe bool) (method, path string, headers map[string]string, body string, err error) {
28-
headers = make(map[string]string)
27+
func ParseRequest(req string, unsafe bool) (method, path string, headers map[string][]string, body string, err error) {
28+
headers = make(map[string][]string)
2929
reader := bufio.NewReader(strings.NewReader(req))
3030
s, err := reader.ReadString('\n')
3131
if err != nil {
@@ -68,7 +68,7 @@ func ParseRequest(req string, unsafe bool) (method, path string, headers map[str
6868
value = strings.TrimSpace(value)
6969
}
7070

71-
headers[key] = value
71+
headers[key] = append(headers[key], value)
7272
}
7373

7474
// Handle case with the full http url in path. In that case,
@@ -81,7 +81,7 @@ func ParseRequest(req string, unsafe bool) (method, path string, headers map[str
8181
return
8282
}
8383
path = parts[1]
84-
headers["Host"] = parsed.Host
84+
headers["Host"] = []string{parsed.Host}
8585
} else {
8686
path = parts[1]
8787
}

‎common/httputilz/httputilz_test.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package httputilz
2+
3+
import (
4+
"strings"
5+
"testing"
6+
7+
"github.com/stretchr/testify/require"
8+
)
9+
10+
func TestParseRequestPreservesDuplicateHeaders(t *testing.T) {
11+
raw := strings.Join([]string{
12+
"GET /anything HTTP/1.1",
13+
"Host: example.com",
14+
"X-Test: one",
15+
"X-Test: two",
16+
"",
17+
"",
18+
}, "\r\n")
19+
20+
method, path, headers, _, err := ParseRequest(raw, false)
21+
require.NoError(t, err)
22+
require.Equal(t, "GET", method)
23+
require.Equal(t, "/anything", path)
24+
require.Equal(t, []string{"one", "two"}, headers["X-Test"])
25+
}
26+
27+
func TestParseRequestFullURLPathSetsHost(t *testing.T) {
28+
raw := strings.Join([]string{
29+
"GET https://example.com/anything HTTP/1.1",
30+
"Host: ignored.example",
31+
"",
32+
"",
33+
}, "\r\n")
34+
35+
_, path, headers, _, err := ParseRequest(raw, false)
36+
require.NoError(t, err)
37+
require.Equal(t, "https://example.com/anything", path)
38+
require.Equal(t, []string{"example.com"}, headers["Host"])
39+
}
40+
41+
func TestParseRequestParsesBody(t *testing.T) {
42+
raw := strings.Join([]string{
43+
"POST /submit HTTP/1.1",
44+
"Host: example.com",
45+
"Content-Type: application/json",
46+
"",
47+
`{"a":1}`,
48+
}, "\r\n")
49+
50+
method, path, headers, body, err := ParseRequest(raw, false)
51+
require.NoError(t, err)
52+
require.Equal(t, "POST", method)
53+
require.Equal(t, "/submit", path)
54+
require.Equal(t, []string{"application/json"}, headers["Content-Type"])
55+
require.Equal(t, `{"a":1}`, body)
56+
}
57+
58+
func TestParseRequestSafeStripsContentLength(t *testing.T) {
59+
raw := strings.Join([]string{
60+
"GET /anything HTTP/1.1",
61+
"Host: example.com",
62+
"Content-Length: 0",
63+
"",
64+
"",
65+
}, "\r\n")
66+
67+
_, _, headers, _, err := ParseRequest(raw, false)
68+
require.NoError(t, err)
69+
require.NotContains(t, headers, "Content-Length")
70+
}
71+
72+
func TestParseRequestUnsafePreservesRawHeaders(t *testing.T) {
73+
raw := strings.Join([]string{
74+
"GET /anything HTTP/1.1",
75+
"Host: example.com",
76+
"Content-Length: 0",
77+
"X-Test: one",
78+
"X-Test: two",
79+
"",
80+
"",
81+
}, "\r\n")
82+
83+
_, _, headers, _, err := ParseRequest(raw, true)
84+
require.NoError(t, err)
85+
// unsafe mode keeps content-length and does not trim values
86+
require.Contains(t, headers, "Content-Length")
87+
require.Equal(t, []string{" one", " two"}, headers["X-Test"])
88+
}
89+
90+
func TestParseRequestMalformed(t *testing.T) {
91+
_, _, _, _, err := ParseRequest("GET\r\n\r\n", false)
92+
require.Error(t, err)
93+
}

‎common/httpx/httpx.go‎

Lines changed: 26 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"io"
88
"net"
99
"net/http"
10+
"net/textproto"
1011
"net/url"
1112
"os"
1213
"strconv"
@@ -36,7 +37,7 @@ type HTTPX struct {
3637
Filters []Filter
3738
Options *Options
3839
htmlPolicy *bluemonday.Policy
39-
CustomHeaders map[string]string
40+
CustomHeaders map[string][]string
4041
cdn *cdncheck.Client
4142
Dialer *fastdialer.Dialer
4243
NetworkPolicy *networkpolicy.NetworkPolicy
@@ -434,19 +435,31 @@ func (h *HTTPX) NewRequestWithContext(ctx context.Context, method, targetURL str
434435
}
435436

436437
// SetCustomHeaders on the provided request
437-
func (h *HTTPX) SetCustomHeaders(r *retryablehttp.Request, headers map[string]string) {
438-
for name, value := range headers {
439-
switch strings.ToLower(name) {
440-
case "host":
441-
r.Host = value
442-
if h.Options.Unsafe {
443-
r.Header.Set("Host", value)
438+
func (h *HTTPX) SetCustomHeaders(r *retryablehttp.Request, headers map[string][]string) {
439+
// Coalesce values by canonical header key first. net/http canonicalizes keys
440+
// on Del/Add, so case-variant duplicates (e.g. "X-Test" and "x-test") would
441+
// otherwise have the second key's Del wipe the values added for the first.
442+
normalized := make(map[string][]string, len(headers))
443+
for name, values := range headers {
444+
canonical := textproto.CanonicalMIMEHeaderKey(name)
445+
normalized[canonical] = append(normalized[canonical], values...)
446+
}
447+
448+
for name, values := range normalized {
449+
r.Header.Del(name)
450+
for _, value := range values {
451+
switch strings.ToLower(name) {
452+
case "host":
453+
r.Host = value
454+
if h.Options.Unsafe {
455+
r.Header.Add("Host", value)
456+
}
457+
case "cookie":
458+
// cookies are set in the default branch, and reset during the follow redirect flow
459+
fallthrough
460+
default:
461+
r.Header.Add(name, value)
444462
}
445-
case "cookie":
446-
// cookies are set in the default branch, and reset during the follow redirect flow
447-
fallthrough
448-
default:
449-
r.Header.Set(name, value)
450463
}
451464
}
452465
if h.Options.RandomAgent {

‎common/httpx/httpx_test.go‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,79 @@ func TestDo(t *testing.T) {
2929
})
3030
}
3131

32+
func TestSetCustomHeaders(t *testing.T) {
33+
h := &HTTPX{Options: &Options{}}
34+
35+
t.Run("duplicate values preserved in order", func(t *testing.T) {
36+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
37+
require.NoError(t, err)
38+
h.SetCustomHeaders(req, map[string][]string{"X-Test": {"one", "two"}})
39+
require.Equal(t, []string{"one", "two"}, req.Header.Values("X-Test"))
40+
})
41+
42+
t.Run("case-variant duplicates are coalesced", func(t *testing.T) {
43+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
44+
require.NoError(t, err)
45+
h.SetCustomHeaders(req, map[string][]string{"X-Test": {"one"}, "x-test": {"two"}})
46+
require.ElementsMatch(t, []string{"one", "two"}, req.Header.Values("X-Test"))
47+
})
48+
49+
t.Run("custom header replaces existing value", func(t *testing.T) {
50+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
51+
require.NoError(t, err)
52+
req.Header.Set("User-Agent", "default-agent")
53+
h.SetCustomHeaders(req, map[string][]string{"User-Agent": {"custom-agent"}})
54+
require.Equal(t, []string{"custom-agent"}, req.Header.Values("User-Agent"))
55+
})
56+
57+
t.Run("host header sets request host", func(t *testing.T) {
58+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
59+
require.NoError(t, err)
60+
h.SetCustomHeaders(req, map[string][]string{"Host": {"custom.host"}})
61+
require.Equal(t, "custom.host", req.Host)
62+
require.Empty(t, req.Header.Values("Host"))
63+
})
64+
65+
t.Run("multiple distinct headers preserved", func(t *testing.T) {
66+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
67+
require.NoError(t, err)
68+
h.SetCustomHeaders(req, map[string][]string{"X-One": {"1"}, "X-Two": {"2"}})
69+
require.Equal(t, []string{"1"}, req.Header.Values("X-One"))
70+
require.Equal(t, []string{"2"}, req.Header.Values("X-Two"))
71+
})
72+
73+
t.Run("multiple cookie values preserved", func(t *testing.T) {
74+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
75+
require.NoError(t, err)
76+
h.SetCustomHeaders(req, map[string][]string{"Cookie": {"a=1", "b=2"}})
77+
require.Equal(t, []string{"a=1", "b=2"}, req.Header.Values("Cookie"))
78+
})
79+
80+
t.Run("empty value applied as-is", func(t *testing.T) {
81+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
82+
require.NoError(t, err)
83+
h.SetCustomHeaders(req, map[string][]string{"X-Empty": {""}})
84+
require.Equal(t, []string{""}, req.Header.Values("X-Empty"))
85+
})
86+
87+
t.Run("unsafe raw header line stored verbatim as key", func(t *testing.T) {
88+
hu := &HTTPX{Options: &Options{Unsafe: true}}
89+
req, err := retryablehttp.NewRequest(http.MethodGet, "https://example.com", nil)
90+
require.NoError(t, err)
91+
// in unsafe mode the runner stores the whole raw header line as the key
92+
// with an empty value; it must survive canonicalization untouched
93+
hu.SetCustomHeaders(req, map[string][]string{"X-Test: one": {""}})
94+
require.Equal(t, []string{""}, req.Header.Values("X-Test: one"))
95+
})
96+
}
97+
98+
func TestParseCustomCookies(t *testing.T) {
99+
options := &Options{CustomHeaders: map[string][]string{"Cookie": {"a=1", "b=2"}}}
100+
options.parseCustomCookies()
101+
require.True(t, options.hasCustomCookies())
102+
require.Len(t, options.customCookies, 2)
103+
}
104+
32105
func TestHTTP11DisablesRetryableHTTP2FallbackClient(t *testing.T) {
33106
options := DefaultOptions
34107
options.Protocol = HTTP11

‎common/httpx/option.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ type Options struct {
3636
Timeout time.Duration
3737
// RetryMax is the maximum number of retries
3838
RetryMax int
39-
CustomHeaders map[string]string
39+
CustomHeaders map[string][]string
4040
// VHostSimilarityRatio 1 - 100
4141
VHostSimilarityRatio int
4242
FollowRedirects bool
@@ -90,7 +90,7 @@ func (options *Options) parseCustomCookies() {
9090
// parse and fill the custom field
9191
for k, v := range options.CustomHeaders {
9292
if strings.EqualFold(k, "cookie") {
93-
req := http.Request{Header: http.Header{"Cookie": []string{v}}}
93+
req := http.Request{Header: http.Header{"Cookie": v}}
9494
options.customCookies = req.Cookies()
9595
}
9696
}

‎runner/runner.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,14 @@ import (
3030
"github.com/PuerkitoBio/goquery"
3131
"github.com/corona10/goimagehash"
3232
"github.com/gocarina/gocsv"
33+
"github.com/happyhackingspace/dit"
3334
"github.com/mfonda/simhash"
3435
asnmap "github.com/projectdiscovery/asnmap/libs"
3536
"github.com/projectdiscovery/fastdialer/fastdialer"
37+
"github.com/projectdiscovery/httpx/common/authprovider"
3638
"github.com/projectdiscovery/httpx/common/customextract"
3739
"github.com/projectdiscovery/httpx/common/hashes/jarm"
3840
"github.com/projectdiscovery/httpx/common/inputformats"
39-
"github.com/happyhackingspace/dit"
40-
"github.com/projectdiscovery/httpx/common/authprovider"
4141
"github.com/projectdiscovery/httpx/static"
4242
"github.com/projectdiscovery/mapcidr/asn"
4343
"github.com/projectdiscovery/networkpolicy"
@@ -238,12 +238,12 @@ func New(options *Options) (*Runner, error) {
238238
httpxOptions.Protocol = httpx.Proto(options.Protocol)
239239

240240
var key, value string
241-
httpxOptions.CustomHeaders = make(map[string]string)
241+
httpxOptions.CustomHeaders = make(map[string][]string)
242242
for _, customHeader := range options.CustomHeaders {
243243
tokens := strings.SplitN(customHeader, ":", two)
244244
// rawhttp skips all checks
245245
if options.Unsafe {
246-
httpxOptions.CustomHeaders[customHeader] = ""
246+
httpxOptions.CustomHeaders[customHeader] = []string{""}
247247
continue
248248
}
249249

@@ -253,7 +253,7 @@ func New(options *Options) (*Runner, error) {
253253
}
254254
key = strings.TrimSpace(tokens[0])
255255
value = strings.TrimSpace(tokens[1])
256-
httpxOptions.CustomHeaders[key] = value
256+
httpxOptions.CustomHeaders[key] = append(httpxOptions.CustomHeaders[key], value)
257257
}
258258
httpxOptions.SniName = options.SniName
259259

@@ -278,7 +278,7 @@ func New(options *Options) (*Runner, error) {
278278
scanopts.Methods = append(scanopts.Methods, rrMethod)
279279
scanopts.RequestURI = rrPath
280280
for name, value := range rrHeaders {
281-
httpxOptions.CustomHeaders[name] = value
281+
httpxOptions.CustomHeaders[name] = append(httpxOptions.CustomHeaders[name], value...)
282282
}
283283
scanopts.RequestBody = rrBody
284284
options.rawRequest = string(rawRequest)

0 commit comments

Comments
 (0)