dnsforward: imp tests

AdguardTeam · Apr 19, 2024 · ca252e8 · ca252e8
1 parent 9d4de18
commit ca252e8
Showing 1 changed file with 33 additions and 0 deletions.
diff --git a/internal/dnsforward/beforerequest_internal_test.go b/internal/dnsforward/beforerequest_internal_test.go
@@ -25,73 +25,89 @@ func TestServer_HandleBefore(t *testing.T) {
 		clientSrvName     string
 		name              string
 		host              string
+		wantErrMsgUDP     string
 		allowedClients    []string
 		disallowedClients []string
 		blockedHosts      []string
 		wantRCode         int
+		wantRCodeUDP      int
 	}{{
 		want:              assert.NotEmpty,
 		clientSrvName:     tlsServerName,
 		name:              "allow_all",
 		host:              testFQDN,
+		wantErrMsgUDP:     "",
 		allowedClients:    []string{},
 		disallowedClients: []string{},
 		blockedHosts:      []string{},
 		wantRCode:         dns.RcodeSuccess,
+		wantRCodeUDP:      dns.RcodeSuccess,
 	}, {
 		want:              assert.NotEmpty,
 		clientSrvName:     clientID + "." + tlsServerName,
 		name:              "allowed_client_allowed",
 		host:              testFQDN,
+		wantErrMsgUDP:     "i/o timeout",
 		allowedClients:    []string{clientID},
 		disallowedClients: []string{},
 		blockedHosts:      []string{},
 		wantRCode:         dns.RcodeSuccess,
+		wantRCodeUDP:      -1,
 	}, {
 		want:              assert.Empty,
 		clientSrvName:     "client-2." + tlsServerName,
 		name:              "allowed_client_rejected",
 		host:              testFQDN,
+		wantErrMsgUDP:     "i/o timeout",
 		allowedClients:    []string{clientID},
 		disallowedClients: []string{},
 		blockedHosts:      []string{},
 		wantRCode:         dns.RcodeRefused,
+		wantRCodeUDP:      -1,
 	}, {
 		want:              assert.NotEmpty,
 		clientSrvName:     tlsServerName,
 		name:              "disallowed_client_allowed",
 		host:              testFQDN,
+		wantErrMsgUDP:     "",
 		allowedClients:    []string{},
 		disallowedClients: []string{clientID},
 		blockedHosts:      []string{},
 		wantRCode:         dns.RcodeSuccess,
+		wantRCodeUDP:      dns.RcodeSuccess,
 	}, {
 		want:              assert.Empty,
 		clientSrvName:     clientID + "." + tlsServerName,
 		name:              "disallowed_client_rejected",
 		host:              testFQDN,
+		wantErrMsgUDP:     "",
 		allowedClients:    []string{},
 		disallowedClients: []string{clientID},
 		blockedHosts:      []string{},
 		wantRCode:         dns.RcodeRefused,
+		wantRCodeUDP:      dns.RcodeSuccess,
 	}, {
 		want:              assert.NotEmpty,
 		clientSrvName:     tlsServerName,
 		name:              "blocked_hosts_allowed",
 		host:              testFQDN,
+		wantErrMsgUDP:     "",
 		allowedClients:    []string{},
 		disallowedClients: []string{},
 		blockedHosts:      []string{blockedHost},
 		wantRCode:         dns.RcodeSuccess,
+		wantRCodeUDP:      dns.RcodeSuccess,
 	}, {
 		want:              assert.Empty,
 		clientSrvName:     tlsServerName,
 		name:              "blocked_hosts_rejected",
 		host:              dns.Fqdn(blockedHost),
+		wantErrMsgUDP:     "i/o timeout",
 		allowedClients:    []string{},
 		disallowedClients: []string{},
 		blockedHosts:      []string{blockedHost},
 		wantRCode:         dns.RcodeRefused,
+		wantRCodeUDP:      -1,
 	}}
 
 	for _, tc := range testCases {
@@ -130,6 +146,23 @@ func TestServer_HandleBefore(t *testing.T) {
 
 			tc.want(t, reply.Answer)
 			assert.Equal(t, tc.wantRCode, reply.Rcode)
+
+			client = &dns.Client{
+				Net: "udp",
+			}
+
+			req = createTestMessage(tc.host)
+			addr = s.dnsProxy.Addr(proxy.ProtoUDP).String()
+
+			reply, _, err = client.Exchange(req, addr)
+			if tc.wantErrMsgUDP != "" {
+				assert.ErrorContains(t, err, tc.wantErrMsgUDP)
+				assert.Nil(t, reply)
+			} else {
+				require.NotNil(t, reply)
+				assert.Equal(t, tc.wantRCodeUDP, reply.Rcode)
+				assert.NotEmpty(t, reply.Answer)
+			}
 		})
 	}
 }