forked from coredns/coredns
-
Notifications
You must be signed in to change notification settings - Fork 0
/
replacer.go
115 lines (102 loc) · 3.06 KB
/
replacer.go
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
package middleware
import (
"strconv"
"strings"
"time"
"github.com/miekg/dns"
)
// Replacer is a type which can replace placeholder
// substrings in a string with actual values from a
// dns.Msg and responseRecorder. Always use
// NewReplacer to get one of these.
type Replacer interface {
Replace(string) string
Set(key, value string)
}
type replacer struct {
replacements map[string]string
emptyValue string
}
// NewReplacer makes a new replacer based on r and rr.
// Do not create a new replacer until r and rr have all
// the needed values, because this function copies those
// values into the replacer. rr may be nil if it is not
// available. emptyValue should be the string that is used
// in place of empty string (can still be empty string).
func NewReplacer(r *dns.Msg, rr *ResponseRecorder, emptyValue string) Replacer {
state := State{W: rr, Req: r}
rep := replacer{
replacements: map[string]string{
"{type}": state.Type(),
"{name}": state.Name(),
"{class}": state.Class(),
"{proto}": state.Proto(),
"{when}": func() string {
return time.Now().Format(timeFormat)
}(),
"{remote}": state.IP(),
"{port}": func() string {
p, _ := state.Port()
return p
}(),
},
emptyValue: emptyValue,
}
if rr != nil {
rcode := dns.RcodeToString[rr.rcode]
if rcode == "" {
rcode = strconv.Itoa(rr.rcode)
}
rep.replacements["{rcode}"] = rcode
rep.replacements["{size}"] = strconv.Itoa(rr.size)
rep.replacements["{duration}"] = time.Since(rr.start).String()
}
// Header placeholders (case-insensitive)
rep.replacements[headerReplacer+"id}"] = strconv.Itoa(int(r.Id))
rep.replacements[headerReplacer+"opcode}"] = strconv.Itoa(int(r.Opcode))
rep.replacements[headerReplacer+"do}"] = boolToString(state.Do())
rep.replacements[headerReplacer+"bufsize}"] = strconv.Itoa(state.Size())
return rep
}
// Replace performs a replacement of values on s and returns
// the string with the replaced values.
func (r replacer) Replace(s string) string {
// Header replacements - these are case-insensitive, so we can't just use strings.Replace()
for strings.Contains(s, headerReplacer) {
idxStart := strings.Index(s, headerReplacer)
endOffset := idxStart + len(headerReplacer)
idxEnd := strings.Index(s[endOffset:], "}")
if idxEnd > -1 {
placeholder := strings.ToLower(s[idxStart : endOffset+idxEnd+1])
replacement := r.replacements[placeholder]
if replacement == "" {
replacement = r.emptyValue
}
s = s[:idxStart] + replacement + s[endOffset+idxEnd+1:]
} else {
break
}
}
// Regular replacements - these are easier because they're case-sensitive
for placeholder, replacement := range r.replacements {
if replacement == "" {
replacement = r.emptyValue
}
s = strings.Replace(s, placeholder, replacement, -1)
}
return s
}
// Set sets key to value in the replacements map.
func (r replacer) Set(key, value string) {
r.replacements["{"+key+"}"] = value
}
func boolToString(b bool) string {
if b {
return "true"
}
return "false"
}
const (
timeFormat = "02/Jan/2006:15:04:05 -0700"
headerReplacer = "{>"
)