coredns/plugin/rewrite/condition.go

package rewrite

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

	"github.com/coredns/coredns/plugin/pkg/replacer"
	"github.com/coredns/coredns/request"

	"github.com/miekg/dns"
)

// Operators that are defined.
const (
	Is         = "is"
	Not        = "not"
	Has        = "has"
	NotHas     = "not_has"
	StartsWith = "starts_with"
	EndsWith   = "ends_with"
	Match      = "match"
	NotMatch   = "not_match"
)

var repl = replacer.New()

// condition is a rewrite condition.
type condition func(string, string) bool

var conditions = map[string]condition{
	Is:         isFunc,
	Not:        notFunc,
	Has:        hasFunc,
	NotHas:     notHasFunc,
	StartsWith: startsWithFunc,
	EndsWith:   endsWithFunc,
	Match:      matchFunc,
	NotMatch:   notMatchFunc,
}

// isFunc is condition for Is operator. It checks for equality.
func isFunc(a, b string) bool { return a == b }

// notFunc is condition for Not operator. It checks for inequality.
func notFunc(a, b string) bool { return a != b }

// hasFunc is condition for Has operator. It checks if b is a substring of a.
func hasFunc(a, b string) bool { return strings.Contains(a, b) }

// notHasFunc is condition for NotHas operator. It checks if b is not a substring of a.
func notHasFunc(a, b string) bool { return !strings.Contains(a, b) }

// startsWithFunc is condition for StartsWith operator. It checks if b is a prefix of a.
func startsWithFunc(a, b string) bool { return strings.HasPrefix(a, b) }

// endsWithFunc is condition for EndsWith operator. It checks if b is a suffix of a.
func endsWithFunc(a, b string) bool {
	// TODO(miek): IsSubDomain
	return strings.HasSuffix(a, b)
}

// matchFunc is condition for Match operator. It does regexp matching of a against pattern in b
// and returns if they match.
func matchFunc(a, b string) bool {
	matched, _ := regexp.MatchString(b, a)
	return matched
}

// notMatchFunc is condition for NotMatch operator. It does regexp matching of a against pattern in b
// and returns if they do not match.
func notMatchFunc(a, b string) bool {
	matched, _ := regexp.MatchString(b, a)
	return !matched
}

// If is statement for a rewrite condition.
type If struct {
	A        string
	Operator string
	B        string
}

// True returns true if the condition is true and false otherwise.
// If r is not nil, it replaces placeholders before comparison.
func (i If) True(r *dns.Msg) bool {
	if c, ok := conditions[i.Operator]; ok {
		a, b := i.A, i.B
		if r != nil {
			ctx := context.TODO()
			state := request.Request{Req: r, W: nil} // hmm W nil?
			a = repl.Replace(ctx, state, nil, i.A)
			b = repl.Replace(ctx, state, nil, i.B)
		}
		return c(a, b)
	}
	return false
}

// NewIf creates a new If condition.
func NewIf(a, operator, b string) (If, error) {
	if _, ok := conditions[operator]; !ok {
		return If{}, fmt.Errorf("invalid operator %v", operator)
	}
	return If{
		A:        a,
		Operator: operator,
		B:        b,
	}, nil
}