explicit have attribute for whether to set EDNS0 in the test cases ra…

…thern than infer it in checkDNSService
hashicorp · Aug 15, 2024 · da2c449 · da2c449
1 parent f8c8d0a
commit da2c449
Showing 1 changed file with 28 additions and 26 deletions.
diff --git a/agent/dns_service_lookup_test.go b/agent/dns_service_lookup_test.go
@@ -3017,6 +3017,7 @@ func checkDNSService(
 	qType uint16,
 	expectedResultsCount int,
 	udpSize uint16,
+	setEDNS0 bool,
 ) {
 	a := NewTestAgent(t, `
 		node_name = "test-node"
@@ -3076,7 +3077,7 @@ func checkDNSService(
 			m := new(dns.Msg)
 
 			m.SetQuestion(question, qType)
-			if udpSize > 512 {
+			if setEDNS0 {
 				m.SetEdns0(udpSize, true)
 			}
 			c := &dns.Client{Net: protocol, UDPSize: udpSize}
@@ -3110,33 +3111,34 @@ func TestDNS_ServiceLookup_ARecordLimits(t *testing.T) {
 		expectedSRVResults  int
 		numNodesTotal       int
 		udpSize             uint16
+		setEDNS0            bool
 	}{
 		// UDP + EDNS
-		{"udp-edns-1", UDP, 1, 1, 1, 1, 30, 30, 8192},
-		{"udp-edns-2", UDP, 2, 2, 2, 2, 30, 30, 8192},
-		{"udp-edns-3", UDP, 3, 3, 3, 3, 30, 30, 8192},
-		{"udp-edns-4", UDP, 4, 4, 4, 4, 30, 30, 8192},
-		{"udp-edns-5", UDP, 5, 5, 5, 5, 30, 30, 8192},
-		{"udp-edns-6", UDP, 6, 6, 6, 6, 30, 30, 8192},
-		{"udp-edns-max", UDP, 6, 2, 1, 3, 3, 3, 8192},
+		{"udp-edns-1", UDP, 1, 1, 1, 1, 30, 30, 8192, true},
+		{"udp-edns-2", UDP, 2, 2, 2, 2, 30, 30, 8192, true},
+		{"udp-edns-3", UDP, 3, 3, 3, 3, 30, 30, 8192, true},
+		{"udp-edns-4", UDP, 4, 4, 4, 4, 30, 30, 8192, true},
+		{"udp-edns-5", UDP, 5, 5, 5, 5, 30, 30, 8192, true},
+		{"udp-edns-6", UDP, 6, 6, 6, 6, 30, 30, 8192, true},
+		{"udp-edns-max", UDP, 6, 2, 1, 3, 3, 3, 8192, true},
 		// All UDP without EDNS and no udpAnswerLimit
 		// Size of records is limited by UDP payload
-		{"udp-1", UDP, 1, 1, 0, 1, 1, 1, 512},
-		{"udp-2", UDP, 2, 1, 1, 2, 2, 2, 512},
-		{"udp-3", UDP, 3, 1, 1, 2, 2, 2, 512},
-		{"udp-4", UDP, 4, 1, 1, 2, 2, 2, 512},
-		{"udp-5", UDP, 5, 1, 1, 2, 2, 2, 512},
-		{"udp-6", UDP, 6, 1, 1, 2, 2, 2, 512},
+		{"udp-1", UDP, 1, 1, 0, 1, 1, 1, 512, false},
+		{"udp-2", UDP, 2, 1, 1, 2, 2, 2, 512, false},
+		{"udp-3", UDP, 3, 1, 1, 2, 2, 2, 512, false},
+		{"udp-4", UDP, 4, 1, 1, 2, 2, 2, 512, false},
+		{"udp-5", UDP, 5, 1, 1, 2, 2, 2, 512, false},
+		{"udp-6", UDP, 6, 1, 1, 2, 2, 2, 512, false},
 		// Only 3 A and 3 SRV records on 512 bytes
-		{"udp-max", UDP, 6, 1, 1, 2, 2, 2, 512},
+		{"udp-max", UDP, 6, 1, 1, 2, 2, 2, 512, false},
 
-		{"tcp-1", TCP, 1, 1, 1, 1, 30, 30, 0},
-		{"tcp-2", TCP, 2, 2, 2, 2, 30, 30, 0},
-		{"tcp-3", TCP, 3, 3, 3, 3, 30, 30, 0},
-		{"tcp-4", TCP, 4, 4, 4, 4, 30, 30, 0},
-		{"tcp-5", TCP, 5, 5, 5, 5, 30, 30, 0},
-		{"tcp-6", TCP, 6, 6, 6, 6, 30, 30, 0},
-		{"tcp-max", TCP, 6, 1, 1, 2, 2, 2, 0},
+		{"tcp-1", TCP, 1, 1, 1, 1, 30, 30, 0, false},
+		{"tcp-2", TCP, 2, 2, 2, 2, 30, 30, 0, false},
+		{"tcp-3", TCP, 3, 3, 3, 3, 30, 30, 0, false},
+		{"tcp-4", TCP, 4, 4, 4, 4, 30, 30, 0, false},
+		{"tcp-5", TCP, 5, 5, 5, 5, 30, 30, 0, false},
+		{"tcp-6", TCP, 6, 6, 6, 6, 30, 30, 0, false},
+		{"tcp-max", TCP, 6, 1, 1, 2, 2, 2, 0, false},
 	}
 	for _, test := range tests {
 		test := test // capture loop var
@@ -3146,20 +3148,20 @@ func TestDNS_ServiceLookup_ARecordLimits(t *testing.T) {
 			// All those queries should have at max queriesLimited elements
 
 			t.Run("A", func(t *testing.T) {
-				checkDNSService(t, test.protocol, test.numNodesTotal, test.aRecordLimit, dns.TypeA, test.expectedAResults, test.udpSize)
+				checkDNSService(t, test.protocol, test.numNodesTotal, test.aRecordLimit, dns.TypeA, test.expectedAResults, test.udpSize, test.setEDNS0)
 			})
 
 			t.Run("AAAA", func(t *testing.T) {
-				checkDNSService(t, test.protocol, test.numNodesTotal, test.aRecordLimit, dns.TypeAAAA, test.expectedAAAAResults, test.udpSize)
+				checkDNSService(t, test.protocol, test.numNodesTotal, test.aRecordLimit, dns.TypeAAAA, test.expectedAAAAResults, test.udpSize, test.setEDNS0)
 			})
 
 			t.Run("ANY", func(t *testing.T) {
-				checkDNSService(t, test.protocol, test.numNodesTotal, test.aRecordLimit, dns.TypeANY, test.expectedANYResults, test.udpSize)
+				checkDNSService(t, test.protocol, test.numNodesTotal, test.aRecordLimit, dns.TypeANY, test.expectedANYResults, test.udpSize, test.setEDNS0)
 			})
 
 			// No limits but the size of records for SRV records, since not subject to randomization issues
 			t.Run("SRV", func(t *testing.T) {
-				checkDNSService(t, test.protocol, test.expectedSRVResults, test.aRecordLimit, dns.TypeSRV, test.numNodesTotal, test.udpSize)
+				checkDNSService(t, test.protocol, test.expectedSRVResults, test.aRecordLimit, dns.TypeSRV, test.numNodesTotal, test.udpSize, test.setEDNS0)
 			})
 		})
 	}