Use Trim(Prefix/Suffix) instead of Trim(Left/Right) in rewrite prefix plugin (#2364) (#2372)

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)
+		}
 	}
 }