From 7491e753c5eb6df54c9c050b0fa198c81dded2a0 Mon Sep 17 00:00:00 2001
From: Simon Zolin <s.zolin@adguard.com>
Date: Fri, 29 May 2020 18:22:32 +0300
Subject: [PATCH] test

---
 dnsforward/dnsforward_test.go | 59 +++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/dnsforward/dnsforward_test.go b/dnsforward/dnsforward_test.go
index 341f2fad143..cece5682226 100644
--- a/dnsforward/dnsforward_test.go
+++ b/dnsforward/dnsforward_test.go
@@ -627,6 +627,65 @@ func TestBlockedBySafeBrowsing(t *testing.T) {
 	}
 }
 
+func TestRewrite(t *testing.T) {
+	c := dnsfilter.Config{}
+	c.Rewrites = []dnsfilter.RewriteEntry{
+		dnsfilter.RewriteEntry{
+			Domain: "test.com",
+			Answer: "1.2.3.4",
+			Type:   dns.TypeA,
+		},
+		dnsfilter.RewriteEntry{
+			Domain: "alias.test.com",
+			Answer: "test.com",
+			Type:   dns.TypeCNAME,
+		},
+		dnsfilter.RewriteEntry{
+			Domain: "my.alias.example.org",
+			Answer: "example.org",
+			Type:   dns.TypeCNAME,
+		},
+	}
+
+	f := dnsfilter.New(&c, nil)
+	s := NewServer(f, nil, nil)
+	conf := ServerConfig{}
+	conf.UDPListenAddr = &net.UDPAddr{Port: 0}
+	conf.TCPListenAddr = &net.TCPAddr{Port: 0}
+	conf.ProtectionEnabled = true
+	conf.UpstreamDNS = []string{"8.8.8.8:53"}
+
+	err := s.Prepare(&conf)
+	assert.Nil(t, err)
+	err = s.Start()
+	assert.Nil(t, err)
+	addr := s.dnsProxy.Addr(proxy.ProtoUDP)
+
+	req := createTestMessageWithType("test.com.", dns.TypeA)
+	reply, err := dns.Exchange(req, addr.String())
+	assert.Nil(t, err)
+	assert.Equal(t, 1, len(reply.Answer))
+	a, ok := reply.Answer[0].(*dns.A)
+	assert.True(t, ok)
+	assert.Equal(t, "1.2.3.4", a.A.String())
+
+	req = createTestMessageWithType("alias.test.com.", dns.TypeA)
+	reply, err = dns.Exchange(req, addr.String())
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(reply.Answer))
+	assert.Equal(t, "test.com.", reply.Answer[0].(*dns.CNAME).Target)
+	assert.Equal(t, "1.2.3.4", reply.Answer[1].(*dns.A).A.String())
+
+	req = createTestMessageWithType("my.alias.example.org.", dns.TypeA)
+	reply, err = dns.Exchange(req, addr.String())
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(reply.Answer))
+	assert.Equal(t, "example.org.", reply.Answer[0].(*dns.CNAME).Target)
+	assert.Equal(t, dns.TypeA, reply.Answer[1].Header().Rrtype)
+
+	_ = s.Stop()
+}
+
 func createTestServer(t *testing.T) *Server {
 	rules := `||nxdomain.example.org
 ||null.example.org^