coredns/plugin/rewrite/cname_target.go

package rewrite

import (
	"context"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"github.com/coredns/coredns/core/dnsserver"
	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/pkg/log"
	"github.com/coredns/coredns/plugin/pkg/upstream"
	"github.com/coredns/coredns/request"

	"github.com/miekg/dns"
)

const maxCNAMERewriteDepth = 8

// UpstreamInt wraps the Upstream API for dependency injection during testing
type UpstreamInt interface {
	Lookup(ctx context.Context, state request.Request, name string, typ uint16) (*dns.Msg, error)
}

// cnameTargetRule is cname target rewrite rule.
type cnameTargetRule struct {
	rewriteType     string
	paramFromTarget string
	paramToTarget   string
	nextAction      string
	pattern         *regexp.Regexp // Compiled paramFromTarget regex for RegexMatch
	Upstream        UpstreamInt    // Upstream for looking up external names during the resolution process.
}

// cnameTargetRuleWithReqState is cname target rewrite rule state
type cnameTargetRuleWithReqState struct {
	rule  cnameTargetRule
	state request.Request
	ctx   context.Context
}

func (r *cnameTargetRule) getFromAndToTarget(inputCName string) (from string, to string) {
	switch r.rewriteType {
	case ExactMatch:
		return r.paramFromTarget, r.paramToTarget
	case PrefixMatch:
		if after, ok := strings.CutPrefix(inputCName, r.paramFromTarget); ok {
			return inputCName, r.paramToTarget + after
		}
	case SuffixMatch:
		if before, ok := strings.CutSuffix(inputCName, r.paramFromTarget); ok {
			return inputCName, before + r.paramToTarget
		}
	case SubstringMatch:
		if strings.Contains(inputCName, r.paramFromTarget) {
			return inputCName, strings.ReplaceAll(inputCName, r.paramFromTarget, r.paramToTarget)
		}
	case RegexMatch:
		regexGroups := r.pattern.FindStringSubmatch(inputCName)
		if len(regexGroups) == 0 {
			return "", ""
		}
		substitution := r.paramToTarget
		for groupIndex, groupValue := range regexGroups {
			groupIndexStr := "{" + strconv.Itoa(groupIndex) + "}"
			substitution = strings.ReplaceAll(substitution, groupIndexStr, groupValue)
		}
		return inputCName, substitution
	}
	return "", ""
}

func (r *cnameTargetRuleWithReqState) RewriteResponse(res *dns.Msg, rr dns.RR) {
	// logic to rewrite the cname target of dns response
	if rr.Header().Rrtype != dns.TypeCNAME {
		return
	}
	// rename the target of the cname response
	cname, ok := rr.(*dns.CNAME)
	if !ok {
		return
	}
	fromTarget, toTarget := r.rule.getFromAndToTarget(cname.Target)
	if cname.Target != fromTarget {
		return
	}

	// Limit internal lookups that re-enter the server.
	loop, _ := r.ctx.Value(dnsserver.LoopKey{}).(int)
	if loop > maxCNAMERewriteDepth {
		return
	}
	ctx := context.WithValue(r.ctx, dnsserver.LoopKey{}, loop+1)

	// create upstream request with the new target with the same qtype
	r.state.Req.Question[0].Name = toTarget
	// upRes can be nil if the internal query path didn't write a response
	// (e.g. a plugin returned a success rcode without writing, dropped the query,
	// or the context was canceled). Guard upRes before dereferencing.
	upRes, err := r.rule.Upstream.Lookup(ctx, r.state, toTarget, r.state.Req.Question[0].Qtype)
	if err != nil {
		log.Errorf("upstream lookup failed: %v", err)
		return
	}
	if upRes == nil {
		log.Errorf("upstream lookup returned nil")
		return
	}

	var newAnswer []dns.RR
	// iterate over first upstream response
	// add the cname record to the new answer
	for _, rr := range res.Answer {
		if cname, ok := rr.(*dns.CNAME); ok {
			// preserve CNAME records until the rewrite target
			newAnswer = append(newAnswer, rr)
			if cname.Target == fromTarget {
				// change the target name in the response
				cname.Target = toTarget
				break
			}
		}
	}
	// add the upstream response to the new answer
	newAnswer = append(newAnswer, upRes.Answer...)
	res.Answer = newAnswer
	// if not propagated, the truncated response might get cached,
	// and it will be impossible to resolve the full response
	res.Truncated = upRes.Truncated
}

func newCNAMERule(nextAction string, args ...string) (Rule, error) {
	var rewriteType string
	var paramFromTarget, paramToTarget string
	if len(args) == 3 {
		rewriteType = (strings.ToLower(args[0]))
		switch rewriteType {
		case ExactMatch:
		case PrefixMatch:
		case SuffixMatch:
		case SubstringMatch:
		case RegexMatch:
		default:
			return nil, fmt.Errorf("unknown cname rewrite type: %s", rewriteType)
		}
		paramFromTarget, paramToTarget = strings.ToLower(args[1]), strings.ToLower(args[2])
	} else if len(args) == 2 {
		rewriteType = ExactMatch
		paramFromTarget, paramToTarget = strings.ToLower(args[0]), strings.ToLower(args[1])
	} else {
		return nil, fmt.Errorf("too few (%d) arguments for a cname rule", len(args))
	}
	if rewriteType == ExactMatch {
		paramFromTarget = plugin.Name(paramFromTarget).Normalize()
		paramToTarget = plugin.Name(paramToTarget).Normalize()
	}
	rule := cnameTargetRule{
		rewriteType:     rewriteType,
		paramFromTarget: paramFromTarget,
		paramToTarget:   paramToTarget,
		nextAction:      nextAction,
		Upstream:        upstream.New(),
	}
	if rewriteType == RegexMatch {
		if len(paramFromTarget) > maxRegexpLen {
			return nil, fmt.Errorf("regex pattern too long in a cname rule: %d > %d", len(paramFromTarget), maxRegexpLen)
		}
		re, err := regexp.Compile(paramFromTarget)
		if err != nil {
			return nil, fmt.Errorf("invalid cname rewrite regex pattern: %w", err)
		}
		rule.pattern = re
	}
	return &rule, nil
}

// Rewrite rewrites the current request.
func (r *cnameTargetRule) Rewrite(ctx context.Context, state request.Request) (ResponseRules, Result) {
	if r != nil && len(r.rewriteType) > 0 && len(r.paramFromTarget) > 0 && len(r.paramToTarget) > 0 {
		return ResponseRules{&cnameTargetRuleWithReqState{
			rule:  *r,
			state: state,
			ctx:   ctx,
		}}, RewriteDone
	}
	return nil, RewriteIgnored
}

// Mode returns the processing mode.
func (r *cnameTargetRule) Mode() string { return r.nextAction }