diff options
-rw-r--r-- | plugin/rewrite/name.go | 2 | ||||
-rw-r--r-- | plugin/rewrite/name_test.go | 57 |
2 files changed, 37 insertions, 22 deletions
diff --git a/plugin/rewrite/name.go b/plugin/rewrite/name.go index 84ba99824..1c181ec03 100644 --- a/plugin/rewrite/name.go +++ b/plugin/rewrite/name.go @@ -80,7 +80,7 @@ func (rule *prefixNameRule) Rewrite(ctx context.Context, state request.Request) // Rewrite rewrites the current request when the name ends with the matching string. func (rule *suffixNameRule) Rewrite(ctx context.Context, state request.Request) Result { if strings.HasSuffix(state.Name(), rule.Suffix) { - state.Req.Question[0].Name = strings.TrimRight(state.Name(), rule.Suffix) + rule.Replacement + state.Req.Question[0].Name = strings.TrimSuffix(state.Name(), rule.Suffix) + rule.Replacement return RewriteDone } return RewriteIgnored diff --git a/plugin/rewrite/name_test.go b/plugin/rewrite/name_test.go index f2245f448..0d3c3dd28 100644 --- a/plugin/rewrite/name_test.go +++ b/plugin/rewrite/name_test.go @@ -32,31 +32,46 @@ func TestRewriteIllegalName(t *testing.T) { } } -func TestRewriteNamePrefix(t *testing.T) { - r, err := newNameRule("stop", "prefix", "test", "not-a-test") - if err != nil { - t.Fatalf("Expected no error, got %s", err) - } +func TestRewriteNamePrefixSuffix(t *testing.T) { - rw := Rewrite{ - Next: plugin.HandlerFunc(msgPrinter), - Rules: []Rule{r}, - noRevert: true, + ctx, close := context.WithCancel(context.TODO()) + defer close() + + tests := []struct { + next string + args []string + question string + expected string + }{ + {"stop", []string{"prefix", "foo", "bar"}, "foo.example.com.", "bar.example.com."}, + {"stop", []string{"prefix", "foo.", "bar."}, "foo.example.com.", "bar.example.com."}, + {"stop", []string{"suffix", "com", "org"}, "foo.example.com.", "foo.example.org."}, + {"stop", []string{"suffix", ".com", ".org"}, "foo.example.com.", "foo.example.org."}, } + for _, tc := range tests { + r, err := newNameRule(tc.next, tc.args...) + if err != nil { + t.Fatalf("Expected no error, got %s", err) + } - ctx := context.TODO() - m := new(dns.Msg) - m.SetQuestion("test.example.org.", dns.TypeA) + rw := Rewrite{ + Next: plugin.HandlerFunc(msgPrinter), + Rules: []Rule{r}, + noRevert: true, + } - rec := dnstest.NewRecorder(&test.ResponseWriter{}) - _, err = rw.ServeDNS(ctx, rec, m) - if err != nil { - t.Fatalf("Expected no error, got %s", err) - } - expected := "not-a-test.example.org." - actual := rec.Msg.Question[0].Name - if actual != expected { - t.Fatalf("Expected rewrite to %v, got %v", expected, actual) + m := new(dns.Msg) + m.SetQuestion(tc.question, dns.TypeA) + + rec := dnstest.NewRecorder(&test.ResponseWriter{}) + _, err = rw.ServeDNS(ctx, rec, m) + if err != nil { + t.Fatalf("Expected no error, got %s", err) + } + actual := rec.Msg.Question[0].Name + if actual != tc.expected { + t.Fatalf("Expected rewrite to %v, got %v", tc.expected, actual) + } } } |