Fix a crash in rewrite plugin when rule type is missing (#5459)

2 files changed, 35 insertions, 47 deletions

diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index c315fa73a..b676a1f2d 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -104,6 +104,9 @@ func newRule(args ...string) (Rule, error) {
 		expectNumArgs = len(args) - 1
 		startArg = 2
 	case Stop:
+		if len(args) < 2 {
+			return nil, fmt.Errorf("stop rule must begin with a rule type")
+		}
 		ruleType = strings.ToLower(args[1])
 		expectNumArgs = len(args) - 1
 		startArg = 2
diff --git a/plugin/rewrite/setup_test.go b/plugin/rewrite/setup_test.go
index 00808bd9c..128f3d11c 100644
--- a/plugin/rewrite/setup_test.go
+++ b/plugin/rewrite/setup_test.go
@@ -8,60 +8,45 @@ import (
 )
 
 func TestParse(t *testing.T) {
-	c := caddy.NewTestController("dns", `rewrite`)
-	_, err := rewriteParse(c)
-	if err == nil {
-		t.Errorf("Expected error but found nil for `rewrite`")
-	}
-	c = caddy.NewTestController("dns", `rewrite name`)
-	_, err = rewriteParse(c)
-	if err == nil {
-		t.Errorf("Expected error but found nil for `rewrite name`")
-	}
-	c = caddy.NewTestController("dns", `rewrite name a.com b.com`)
-	_, err = rewriteParse(c)
-	if err != nil {
-		t.Errorf("Expected success but found %s for `rewrite name a.com b.com`", err)
-	}
 
-	c = caddy.NewTestController("dns",
-		`rewrite stop {
+	tests := []struct {
+		inputFileRules string
+		shouldErr      bool
+		errContains    string
+	}{
+		// parse errors
+		{`rewrite`, true, ""},
+		{`rewrite name`, true, ""},
+		{`rewrite name a.com b.com`, false, ""},
+		{`rewrite stop {
     name regex foo bar
     answer name bar foo
-}`)
-	_, err = rewriteParse(c)
-	if err != nil {
-		t.Errorf("Expected success but found %s for valid response rewrite", err)
-	}
-
-	c = caddy.NewTestController("dns", `rewrite stop name regex foo bar answer name bar foo`)
-	_, err = rewriteParse(c)
-	if err != nil {
-		t.Errorf("Expected success but found %s for valid response rewrite", err)
-	}
-
-	c = caddy.NewTestController("dns",
-		`rewrite stop {
+}`, false, ""},
+		{`rewrite stop name regex foo bar answer name bar foo`, false, ""},
+		{`rewrite stop {
     name regex foo bar
     answer name bar foo
     name baz
-}`)
-	_, err = rewriteParse(c)
-	if err == nil {
-		t.Errorf("Expected error but got success for invalid response rewrite")
-	} else if !strings.Contains(err.Error(), "2 arguments required") {
-		t.Errorf("Got wrong error for invalid response rewrite: %v", err.Error())
-	}
-
-	c = caddy.NewTestController("dns",
-		`rewrite stop {
+}`, true, "2 arguments required"},
+		{`rewrite stop {
     answer name bar foo
     name regex foo bar
-}`)
-	_, err = rewriteParse(c)
-	if err == nil {
-		t.Errorf("Expected error but got success for invalid response rewrite")
-	} else if !strings.Contains(err.Error(), "must begin with a name rule") {
-		t.Errorf("Got wrong error for invalid response rewrite: %v", err.Error())
+}`, true, "must begin with a name rule"},
+		{`rewrite stop`, true, ""},
+	}
+
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.inputFileRules)
+		_, err := rewriteParse(c)
+		if err == nil && test.shouldErr {
+			t.Fatalf("Test %d expected errors, but got no error\n---\n%s", i, test.inputFileRules)
+		} else if err != nil && !test.shouldErr {
+			t.Fatalf("Test %d expected no errors, but got '%v'\n---\n%s", i, err, test.inputFileRules)
+		}
+
+		if err != nil && test.errContains != "" && !strings.Contains(err.Error(), test.errContains) {
+			t.Errorf("Test %d got wrong error for invalid response rewrite: '%v'\n---\n%s", i, err.Error(), test.inputFileRules)
+		}
 	}
+
 }