1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
// Package rules implements match-and-replace: user-defined rules that
// rewrite request/response headers as they pass through the proxy.
// Deliberately headers-only for now - see ApplyHeaders for why bodies
// are a separate, harder problem (noted as a follow-up in PLAN.md).
package rules
import (
"bufio"
"net/http"
"net/textproto"
"regexp"
"sort"
"strings"
)
// Rule is one match-and-replace rule.
type Rule struct {
ID int64
Enabled bool
Name string
Scope string // "request" or "response"
Part string // "header" (only part supported so far)
Match string
Replace string
IsRegex bool
// Position orders rule application (ascending) when several rules
// could touch the same text.
Position int
}
// ApplyHeaders rewrites h in place by serializing it to a raw
// "Name: value\r\n" block, running every enabled rule with Part=="header"
// over that text (in Position order), and reparsing the result. Working
// on the raw text rather than per-value substitution is what lets a rule
// add or remove a header entirely, not just rewrite an existing value -
// matching how Burp's header match/replace works. If a rule's output
// doesn't parse back as valid headers, ApplyHeaders returns h unchanged
// rather than risk sending something corrupted.
func ApplyHeaders(h http.Header, rs []Rule) http.Header {
keys := make([]string, 0, len(h))
for k := range h {
keys = append(keys, k)
}
sort.Strings(keys)
var block strings.Builder
for _, k := range keys {
for _, v := range h[k] {
block.WriteString(k)
block.WriteString(": ")
block.WriteString(v)
block.WriteString("\r\n")
}
}
text := block.String()
changed := false
for _, r := range sortedByPosition(rs) {
if !r.Enabled || r.Part != "header" {
continue
}
if next, ok := apply(text, r); ok {
text, changed = next, true
}
}
if !changed {
return h
}
tp := textproto.NewReader(bufio.NewReader(strings.NewReader(text + "\r\n")))
mh, err := tp.ReadMIMEHeader()
if err != nil {
return h
}
return http.Header(mh)
}
func sortedByPosition(rs []Rule) []Rule {
out := make([]Rule, len(rs))
copy(out, rs)
sort.SliceStable(out, func(i, j int) bool { return out[i].Position < out[j].Position })
return out
}
func apply(text string, r Rule) (string, bool) {
if r.IsRegex {
re, err := regexp.Compile(r.Match)
if err != nil {
return text, false
}
return re.ReplaceAllString(text, r.Replace), true
}
if r.Match == "" {
return text, false
}
return strings.ReplaceAll(text, r.Match, r.Replace), true
}
|