// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: BUSL-1.1

package agent

import (
	"context"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/miekg/dns"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestDNS_ReverseLookup(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register node
			args := &structs.RegisterRequest{
				Datacenter: "dc1",
				Node:       "foo2",
				Address:    "127.0.0.2",
			}

			var out struct{}
			if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
				t.Fatalf("err: %v", err)
			}

			m := new(dns.Msg)
			m.SetQuestion("2.0.0.127.in-addr.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != "foo2.node.dc1.consul." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_ReverseLookup_CustomDomain(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, `
		domain = "custom"
	`+experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register node
			args := &structs.RegisterRequest{
				Datacenter: "dc1",
				Node:       "foo2",
				Address:    "127.0.0.2",
			}

			var out struct{}
			if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
				t.Fatalf("err: %v", err)
			}

			m := new(dns.Msg)
			m.SetQuestion("2.0.0.127.in-addr.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != "foo2.node.dc1.custom." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_ReverseLookup_IPV6(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register node
			args := &structs.RegisterRequest{
				Datacenter: "dc1",
				Node:       "bar",
				Address:    "::4242:4242",
			}

			var out struct{}
			if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
				t.Fatalf("err: %v", err)
			}

			m := new(dns.Msg)
			m.SetQuestion("2.4.2.4.2.4.2.4.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.ip6.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != "bar.node.dc1.consul." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_Compression_ReverseLookup(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {

			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register node.
			args := &structs.RegisterRequest{
				Datacenter: "dc1",
				Node:       "foo2",
				Address:    "127.0.0.2",
			}
			var out struct{}
			if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
				t.Fatalf("err: %v", err)
			}

			m := new(dns.Msg)
			m.SetQuestion("2.0.0.127.in-addr.arpa.", dns.TypeANY)

			conn, err := dns.Dial("udp", a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// Do a manual exchange with compression on (the default).
			if err := conn.WriteMsg(m); err != nil {
				t.Fatalf("err: %v", err)
			}
			p := make([]byte, dns.MaxMsgSize)
			compressed, err := conn.Read(p)
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// Disable compression and try again.
			a.DNSDisableCompression(true)
			if err := conn.WriteMsg(m); err != nil {
				t.Fatalf("err: %v", err)
			}
			unc, err := conn.Read(p)
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			// We can't see the compressed status given the DNS API, so we just make
			// sure the message is smaller to see if it's respecting the flag.
			if compressed == 0 || unc == 0 || compressed >= unc {
				t.Fatalf("doesn't look compressed: %d vs. %d", compressed, unc)
			}
		})
	}
}

func TestDNS_ServiceReverseLookup(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register a node with a service.
			{
				args := &structs.RegisterRequest{
					Datacenter: "dc1",
					Node:       "foo",
					Address:    "127.0.0.1",
					Service: &structs.NodeService{
						Service: "db",
						Tags:    []string{"primary"},
						Port:    12345,
						Address: "127.0.0.2",
					},
				}

				var out struct{}
				if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
					t.Fatalf("err: %v", err)
				}
			}

			m := new(dns.Msg)
			m.SetQuestion("2.0.0.127.in-addr.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != serviceCanonicalDNSName("db", "service", "dc1", "consul", nil)+"." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_ServiceReverseLookup_IPV6(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register a node with a service.
			{
				args := &structs.RegisterRequest{
					Datacenter: "dc1",
					Node:       "foo",
					Address:    "2001:db8::1",
					Service: &structs.NodeService{
						Service: "db",
						Tags:    []string{"primary"},
						Port:    12345,
						Address: "2001:db8::ff00:42:8329",
					},
				}

				var out struct{}
				if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
					t.Fatalf("err: %v", err)
				}
			}

			m := new(dns.Msg)
			m.SetQuestion("9.2.3.8.2.4.0.0.0.0.f.f.0.0.0.0.0.0.0.0.0.0.0.0.8.b.d.0.1.0.0.2.ip6.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != serviceCanonicalDNSName("db", "service", "dc1", "consul", nil)+"." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_ServiceReverseLookup_CustomDomain(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, `
		domain = "custom"
	`+experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register a node with a service.
			{
				args := &structs.RegisterRequest{
					Datacenter: "dc1",
					Node:       "foo",
					Address:    "127.0.0.1",
					Service: &structs.NodeService{
						Service: "db",
						Tags:    []string{"primary"},
						Port:    12345,
						Address: "127.0.0.2",
					},
				}

				var out struct{}
				if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
					t.Fatalf("err: %v", err)
				}
			}

			m := new(dns.Msg)
			m.SetQuestion("2.0.0.127.in-addr.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != serviceCanonicalDNSName("db", "service", "dc1", "custom", nil)+"." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_ServiceReverseLookupNodeAddress(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Register a node with a service.
			{
				args := &structs.RegisterRequest{
					Datacenter: "dc1",
					Node:       "foo",
					Address:    "127.0.0.1",
					Service: &structs.NodeService{
						Service: "db",
						Tags:    []string{"primary"},
						Port:    12345,
						Address: "127.0.0.1",
					},
				}

				var out struct{}
				if err := a.RPC(context.Background(), "Catalog.Register", args, &out); err != nil {
					t.Fatalf("err: %v", err)
				}
			}

			m := new(dns.Msg)
			m.SetQuestion("1.0.0.127.in-addr.arpa.", dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			if err != nil {
				t.Fatalf("err: %v", err)
			}

			if len(in.Answer) != 1 {
				t.Fatalf("Bad: %#v", in)
			}

			ptrRec, ok := in.Answer[0].(*dns.PTR)
			if !ok {
				t.Fatalf("Bad: %#v", in.Answer[0])
			}
			if ptrRec.Ptr != "foo.node.dc1.consul." {
				t.Fatalf("Bad: %#v", ptrRec)
			}
		})
	}
}

func TestDNS_ReverseLookup_NotFound(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	for name, experimentsHCL := range getVersionHCL(true) {
		t.Run(name, func(t *testing.T) {
			// do not configure recursors
			a := NewTestAgent(t, experimentsHCL)
			defer a.Shutdown()
			testrpc.WaitForLeader(t, a.RPC, "dc1")

			// Do not register any nodes
			m := new(dns.Msg)
			qName := "2.0.0.127.in-addr.arpa."
			m.SetQuestion(qName, dns.TypeANY)

			c := new(dns.Client)
			in, _, err := c.Exchange(m, a.DNSAddr())
			require.NoError(t, err)
			require.Nil(t, in.Answer)
			require.Nil(t, in.Extra)

			require.Equal(t, dns.RcodeNameError, in.Rcode)

			question := in.Question[0]
			require.Equal(t, qName, question.Name)
			require.Equal(t, dns.TypeANY, question.Qtype)
			require.Equal(t, uint16(dns.ClassINET), question.Qclass)

			soa, ok := in.Ns[0].(*dns.SOA)
			require.True(t, ok)
			require.Equal(t, "ns.consul.", soa.Ns)
			require.Equal(t, "hostmaster.consul.", soa.Mbox)
		})
	}
}