-
Notifications
You must be signed in to change notification settings - Fork 62
/
replace.go
96 lines (79 loc) · 2.29 KB
/
replace.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
package rdns
import (
"errors"
"regexp"
"github.com/miekg/dns"
)
// Replace is a resolver that modifies queries according to regular expressions
// and forwards the modified queries to another resolver. Responses are then
// mapped back to the original query string.
type Replace struct {
id string
resolver Resolver
exp replaceExpressions
}
var _ Resolver = &Replace{}
type replaceExp struct {
from *regexp.Regexp
to string
}
type replaceExpressions []replaceExp
func (r replaceExpressions) apply(name string) string {
for _, e := range r {
name = e.from.ReplaceAllString(name, e.to)
}
return name
}
type ReplaceOperation struct {
From string
To string
}
// NewReplace returns a new instance of a Replace resolver.
func NewReplace(id string, resolver Resolver, list ...ReplaceOperation) (*Replace, error) {
var exp replaceExpressions
for _, o := range list {
re, err := regexp.Compile(o.From)
if err != nil {
return nil, err
}
exp = append(exp, replaceExp{re, o.To})
}
return &Replace{id: id, resolver: resolver, exp: exp}, nil
}
// Resolve a DNS query by first replacing the query string with another
// sending the query upstream and replace the name in the response with
// the original query string again.
func (r *Replace) Resolve(q *dns.Msg, ci ClientInfo) (*dns.Msg, error) {
if len(q.Question) < 1 {
return nil, errors.New("no question in query")
}
oldName := q.Question[0].Name
newName := r.exp.apply(oldName)
log := logger(r.id, q, ci)
// if nothing needs modifying, we can stop here and use the original query
if newName == oldName {
log.Debug("forwarding unmodified query to resolver")
return r.resolver.Resolve(q, ci)
}
// Modify the query string
q.Question[0].Name = newName
// Send the query upstream
log.WithField("new-qname", newName).WithField("resolver", r.resolver).Debug("forwarding modified query to resolver")
a, err := r.resolver.Resolve(q, ci)
if err != nil || a == nil {
return nil, err
}
// Set the question back to the original name
a.Question[0].Name = oldName
// Now put the original name in all answer records that have the
// new name
for _, answer := range a.Answer {
if answer.Header().Name == newName {
answer.Header().Name = oldName
}
}
return a, nil
}
func (r *Replace) String() string {
return r.id
}