1 files changed, 112 insertions, 99 deletions

diff --git a/plugin/pkg/replacer/replacer.go b/plugin/pkg/replacer/replacer.go
index 8c8032c76..f72e1ec9c 100644
--- a/plugin/pkg/replacer/replacer.go
+++ b/plugin/pkg/replacer/replacer.go
@@ -13,123 +13,134 @@ import (
 	"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)
+// Replacer replaces labels for values in strings.
+type Replacer struct {
+	valueFunc func(request.Request, *dnstest.Recorder, string) string
+	labels    []string
 }
 
-type replacer struct {
-	ctx          context.Context
-	replacements map[string]string
-	emptyValue   string
+// labels are all supported labels that can be used in the default Replacer.
+var labels = []string{
+	"{type}",
+	"{name}",
+	"{class}",
+	"{proto}",
+	"{size}",
+	"{remote}",
+	"{port}",
+	"{local}",
+	// Header values.
+	headerReplacer + "id}",
+	headerReplacer + "opcode}",
+	headerReplacer + "do}",
+	headerReplacer + "bufsize}",
+	// Recorded replacements.
+	"{rcode}",
+	"{rsize}",
+	"{duration}",
+	headerReplacer + "rrflags}",
 }
 
-// New 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 New(ctx context.Context, r *dns.Msg, rr *dnstest.Recorder, emptyValue string) Replacer {
-	req := request.Request{W: rr, Req: r}
-	rep := replacer{
-		ctx: ctx,
-		replacements: map[string]string{
-			"{type}":   req.Type(),
-			"{name}":   req.Name(),
-			"{class}":  req.Class(),
-			"{proto}":  req.Proto(),
-			"{when}":   "", // made a noop
-			"{size}":   strconv.Itoa(req.Len()),
-			"{remote}": addrToRFC3986(req.IP()),
-			"{port}":   req.Port(),
-			"{local}":  addrToRFC3986(req.LocalIP()),
-		},
-		emptyValue: emptyValue,
-	}
-	if rr != nil {
+// value returns the current value of label.
+func value(state request.Request, rr *dnstest.Recorder, label string) string {
+	switch label {
+	case "{type}":
+		return state.Type()
+	case "{name}":
+		return state.Name()
+	case "{class}":
+		return state.Class()
+	case "{proto}":
+		return state.Proto()
+	case "{size}":
+		return strconv.Itoa(state.Req.Len())
+	case "{remote}":
+		return addrToRFC3986(state.IP())
+	case "{port}":
+		return state.Port()
+	case "{local}":
+		return addrToRFC3986(state.LocalIP())
+	// Header placeholders (case-insensitive).
+	case headerReplacer + "id}":
+		return strconv.Itoa(int(state.Req.Id))
+	case headerReplacer + "opcode}":
+		return strconv.Itoa(state.Req.Opcode)
+	case headerReplacer + "do}":
+		return boolToString(state.Do())
+	case headerReplacer + "bufsize}":
+		return strconv.Itoa(state.Size())
+	// Recorded replacements.
+	case "{rcode}":
+		if rr == nil {
+			return EmptyValue
+		}
 		rcode := dns.RcodeToString[rr.Rcode]
 		if rcode == "" {
 			rcode = strconv.Itoa(rr.Rcode)
 		}
-		rep.replacements["{rcode}"] = rcode
-		rep.replacements["{rsize}"] = strconv.Itoa(rr.Len)
-		rep.replacements["{duration}"] = strconv.FormatFloat(time.Since(rr.Start).Seconds(), 'f', -1, 64) + "s"
-		if rr.Msg != nil {
-			rep.replacements[headerReplacer+"rflags}"] = flagsToString(rr.Msg.MsgHdr)
+		return rcode
+	case "{rsize}":
+		if rr == nil {
+			return EmptyValue
+		}
+		return strconv.Itoa(rr.Len)
+	case "{duration}":
+		if rr == nil {
+			return EmptyValue
+		}
+		return strconv.FormatFloat(time.Since(rr.Start).Seconds(), 'f', -1, 64) + "s"
+	case headerReplacer + "rrflags}":
+		if rr != nil && rr.Msg != nil {
+			return flagsToString(rr.Msg.MsgHdr)
 		}
+		return EmptyValue
 	}
+	return EmptyValue
+}
 
-	// Header placeholders (case-insensitive)
-	rep.replacements[headerReplacer+"id}"] = strconv.Itoa(int(r.Id))
-	rep.replacements[headerReplacer+"opcode}"] = strconv.Itoa(r.Opcode)
-	rep.replacements[headerReplacer+"do}"] = boolToString(req.Do())
-	rep.replacements[headerReplacer+"bufsize}"] = strconv.Itoa(req.Size())
-
-	return rep
+// New makes a new replacer. This only needs to be called once in the setup and then call Replace for each incoming message.
+// A replacer is safe for concurrent use.
+func New() Replacer {
+	return Replacer{
+		valueFunc: value,
+		labels:    labels,
+	}
 }
 
-// Replace performs a replacement of values on s and returns
-// the string with the replaced values.
-func (r replacer) Replace(s string) string {
-
-	// declare a function that replace based on header matching
-	fscanAndReplace := func(s string, header string, replace func(string) string) string {
-		b := strings.Builder{}
-		for strings.Contains(s, header) {
-			idxStart := strings.Index(s, header)
-			endOffset := idxStart + len(header)
-			idxEnd := strings.Index(s[endOffset:], "}")
-			if idxEnd > -1 {
-				placeholder := strings.ToLower(s[idxStart : endOffset+idxEnd+1])
-				replacement := replace(placeholder)
-				if replacement == "" {
-					replacement = r.emptyValue
-				}
-				b.WriteString(s[:idxStart])
-				b.WriteString(replacement)
-				s = s[endOffset+idxEnd+1:]
-			} else {
-				break
-			}
+// Replace performs a replacement of values on s and returns the string with the replaced values.
+func (r Replacer) Replace(ctx context.Context, state request.Request, rr *dnstest.Recorder, s string) string {
+	for _, placeholder := range r.labels {
+		if strings.Contains(s, placeholder) {
+			s = strings.Replace(s, placeholder, r.valueFunc(state, rr, placeholder), -1)
 		}
-		b.WriteString(s)
-		return b.String()
 	}
 
-	// Header replacements - these are case-insensitive, so we can't just use strings.Replace()
-	s = fscanAndReplace(s, headerReplacer, func(placeholder string) string {
-		return r.replacements[placeholder]
-	})
+	// Metadata label replacements. Scan for {/ and search for next }, replace that metadata label with
+	// any meta data that is available.
+	b := strings.Builder{}
+	for strings.Contains(s, labelReplacer) {
+		idxStart := strings.Index(s, labelReplacer)
+		endOffset := idxStart + len(labelReplacer)
+		idxEnd := strings.Index(s[endOffset:], "}")
+		if idxEnd > -1 {
+			label := s[idxStart+2 : endOffset+idxEnd]
+
+			fm := metadata.ValueFunc(ctx, label)
+			replacement := EmptyValue
+			if fm != nil {
+				replacement = fm()
+			}
 
-	// Regular replacements - these are easier because they're case-sensitive
-	for placeholder, replacement := range r.replacements {
-		if replacement == "" {
-			replacement = r.emptyValue
+			b.WriteString(s[:idxStart])
+			b.WriteString(replacement)
+			s = s[endOffset+idxEnd+1:]
+		} else {
+			break
 		}
-		s = strings.Replace(s, placeholder, replacement, -1)
 	}
 
-	// Metadata label replacements
-	s = fscanAndReplace(s, headerLabelReplacer, func(placeholder string) string {
-		// label place holder has the format {/<label>}
-		fm := metadata.ValueFunc(r.ctx, placeholder[len(headerLabelReplacer):len(placeholder)-1])
-		if fm != nil {
-			return fm()
-		}
-		return ""
-	})
-
-	return s
-}
-
-// Set sets key to value in the replacements map.
-func (r replacer) Set(key, value string) {
-	r.replacements["{"+key+"}"] = value
+	b.WriteString(s)
+	return b.String()
 }
 
 func boolToString(b bool) string {
@@ -190,6 +201,8 @@ func addrToRFC3986(addr string) string {
 }
 
 const (
-	headerReplacer      = "{>"
-	headerLabelReplacer = "{/"
+	headerReplacer = "{>"
+	labelReplacer  = "{/"
+	// EmptyValue is the default empty value.
+	EmptyValue = "-"
 )