From 9efc38122446fb2b3b29487e41b8bae72db2d8c8 Mon Sep 17 00:00:00 2001
From: Dimitry Kolyshev <dkolyshev@adguard.com>
Date: Fri, 23 Dec 2022 13:18:04 +0700
Subject: [PATCH] rewrite: test matching cnames

---
 internal/filtering/rewrite/storage_test.go | 68 ++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/internal/filtering/rewrite/storage_test.go b/internal/filtering/rewrite/storage_test.go
index 55dc551c..386a034c 100644
--- a/internal/filtering/rewrite/storage_test.go
+++ b/internal/filtering/rewrite/storage_test.go
@@ -401,6 +401,74 @@ func TestDefaultStorage_MatchRequest_ExceptionCNAME(t *testing.T) {
 	}
 }
 
+func TestDefaultStorage_MatchRequest_CNAMEs(t *testing.T) {
+	// Two cname rules for one subdomain
+	items := []*Item{{
+		Domain: "cname.org",
+		Answer: "1.1.1.1",
+	}, {
+		Domain: "sub_cname.org",
+		Answer: "2.2.2.2",
+	}, {
+		Domain: "*.host.com",
+		Answer: "cname.org",
+	}, {
+		Domain: "*.sub.host.com",
+		Answer: "sub_cname.org",
+	}}
+
+	s, err := NewDefaultStorage(-1, items)
+	require.NoError(t, err)
+
+	testCases := []struct {
+		name            string
+		host            string
+		wantDNSRewrites []*rules.DNSRewrite
+		dtyp            uint16
+	}{{
+		name: "match_my_domain",
+		host: "my.host.com",
+		wantDNSRewrites: []*rules.DNSRewrite{{
+			Value:    nil,
+			NewCNAME: "cname.org",
+			RCode:    dns.RcodeSuccess,
+			RRType:   dns.TypeNone,
+		}, {
+			Value:    net.IP{1, 1, 1, 1}.To16(),
+			NewCNAME: "",
+			RCode:    dns.RcodeSuccess,
+			RRType:   dns.TypeA,
+		}},
+		dtyp: dns.TypeA,
+	}, {
+		name: "match_sub_my_domain",
+		host: "my.sub.host.com",
+		wantDNSRewrites: []*rules.DNSRewrite{{
+			Value:    nil,
+			NewCNAME: "cname.org",
+			RCode:    dns.RcodeSuccess,
+			RRType:   dns.TypeNone,
+		}, {
+			Value:    net.IP{1, 1, 1, 1}.To16(),
+			NewCNAME: "",
+			RCode:    dns.RcodeSuccess,
+			RRType:   dns.TypeA,
+		}},
+		dtyp: dns.TypeA,
+	}}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			dnsRewrites := s.MatchRequest(&urlfilter.DNSRequest{
+				Hostname: tc.host,
+				DNSType:  tc.dtyp,
+			})
+
+			assert.Equal(t, tc.wantDNSRewrites, dnsRewrites)
+		})
+	}
+}
+
 func TestDefaultStorage_MatchRequest_ExceptionIP(t *testing.T) {
 	// Exception for AAAA record.
 	items := []*Item{{