From e9aa97b2ab6ac2b1f42341734cb9814c912d5fcf Mon Sep 17 00:00:00 2001
From: Emmanuel Bretelle <chantra@fb.com>
Date: Tue, 22 Oct 2019 18:33:35 -0700
Subject: [PATCH] [zone parser] disallow nested $GENERATE directive

While the range number of GENERATE is now limited, one can pass
a line with 2 $GENERATE directive that will exponentially increase the
time spent generating RRs.
Limit to only one per line.
Fixes #1020
---
 generate.go      | 1 +
 generate_test.go | 5 +++++
 scan.go          | 4 ++++
 3 files changed, 10 insertions(+)

diff --git a/generate.go b/generate.go
index 26edc4f40..f7e91a23f 100644
--- a/generate.go
+++ b/generate.go
@@ -85,6 +85,7 @@ func (zp *ZoneParser) generate(l lex) (RR, bool) {
 	}
 	zp.sub = NewZoneParser(r, zp.origin, zp.file)
 	zp.sub.includeDepth, zp.sub.includeAllowed = zp.includeDepth, zp.includeAllowed
+	zp.sub.generateDisallowed = true
 	zp.sub.SetDefaultTTL(defaultTtl)
 	return zp.subNext()
 }
diff --git a/generate_test.go b/generate_test.go
index 15ab9c989..481c9b281 100644
--- a/generate_test.go
+++ b/generate_test.go
@@ -61,6 +61,10 @@ $GENERATE 0-1/0 dhcp-${0,4,d} A 10.0.0.$
 `, true},
 		{`@ IN SOA ns.test. hostmaster.test. ( 1 8h 2h 7d 1d )
 $GENERATE 0-1 $$INCLUDE ` + tmpdir + string(filepath.Separator) + `${0,4,d}.conf
+`, false},
+{`@ IN SOA ns.test. hostmaster.test. ( 1 8h 2h 7d 1d )
+$GENERATE 0-1 dhcp-${0,4,d} A 10.0.0.$
+$GENERATE 0-2 dhcp-${0,4,d} A 10.1.0.$
 `, false},
 	}
 Outer:
@@ -214,6 +218,7 @@ func TestCrasherString(t *testing.T) {
 		{"$GENERATE 0-5414137360", "bad range in $GENERATE"},
 		{"$GENERATE       11522-3668518066406258", "bad range in $GENERATE"},
 		{"$GENERATE 0-200\"(;00000000000000\n$$GENERATE 0-0", "dns: garbage after $GENERATE range: \"\\\"\" at line: 1:16"},
+		{"$GENERATE 6-2048 $$GENERATE 6-036160 $$$$ORIGIN \\$", `dns: nested $GENERATE directive not allowed: "6-036160" at line: 1:19`},
 	}
 	for _, tc := range tests {
 		t.Run(tc.in, func(t *testing.T) {
diff --git a/scan.go b/scan.go
index 0f9361af0..671018b1f 100644
--- a/scan.go
+++ b/scan.go
@@ -248,6 +248,7 @@ type ZoneParser struct {
 	includeDepth uint8
 
 	includeAllowed bool
+	generateDisallowed bool
 }
 
 // NewZoneParser returns an RFC 1035 style zonefile parser that reads
@@ -547,6 +548,9 @@ func (zp *ZoneParser) Next() (RR, bool) {
 
 			st = zExpectDirGenerate
 		case zExpectDirGenerate:
+			if zp.generateDisallowed {
+				return zp.setParseError("nested $GENERATE directive not allowed", l)
+			}
 			if l.value != zString {
 				return zp.setParseError("expecting $GENERATE value, not this...", l)
 			}