Add support for using the update provider for A data source

Stefan Richter · 02strich · commit 5e162132bdb0 · 2024-09-09T14:38:41.000-07:00
This adds a new option to use the DNS server configured for update requests also with the A and AAAA data source.
diff --git a/docs/data-sources/a_record_set.md b/docs/data-sources/a_record_set.md
@@ -29,6 +29,10 @@ output "google_addrs" {
 
 - `host` (String) Host to look up.
 
+### Optional
+
+- `use_update_server` (Boolean) Whether to use the configured update DNS server
+
 ### Read-Only
 
 - `addrs` (List of String) A list of IP addresses. IP addresses are always sorted to avoid constant changing plans.
diff --git a/internal/provider/data_dns_a_record_set.go b/internal/provider/data_dns_a_record_set.go
@@ -12,17 +12,21 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/datasource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/types"
+	"github.com/miekg/dns"
 )
 
 var (
-	_ datasource.DataSource = (*dnsARecordSetDataSource)(nil)
+	_ datasource.DataSource              = (*dnsARecordSetDataSource)(nil)
+	_ datasource.DataSourceWithConfigure = (*dnsARecordSetDataSource)(nil)
 )
 
 func NewDnsARecordSetDataSource() datasource.DataSource {
 	return &dnsARecordSetDataSource{}
 }
 
-type dnsARecordSetDataSource struct{}
+type dnsARecordSetDataSource struct {
+	client *DNSClient
+}
 
 func (d *dnsARecordSetDataSource) Metadata(ctx context.Context, req datasource.MetadataRequest, resp *datasource.MetadataResponse) {
 	resp.TypeName = req.ProviderTypeName + "_a_record_set"
@@ -36,6 +40,10 @@ func (d *dnsARecordSetDataSource) Schema(ctx context.Context, req datasource.Sch
 				Required:    true,
 				Description: "Host to look up.",
 			},
+			"use_update_server": schema.BoolAttribute{
+				Optional:    true,
+				Description: "Whether to use the configured update DNS server",
+			},
 			"addrs": schema.ListAttribute{
 				ElementType: types.StringType,
 				Computed:    true,
@@ -49,6 +57,24 @@ func (d *dnsARecordSetDataSource) Schema(ctx context.Context, req datasource.Sch
 	}
 }
 
+func (d *dnsARecordSetDataSource) Configure(ctx context.Context, req datasource.ConfigureRequest, resp *datasource.ConfigureResponse) {
+	if req.ProviderData == nil {
+		return
+	}
+
+	client, ok := req.ProviderData.(*DNSClient)
+	if !ok {
+		resp.Diagnostics.AddError(
+			"Unexpected Resource Configure Type",
+			fmt.Sprintf("Expected *DNSClient, got: %T. Please report this issue to the provider developers.", req.ProviderData),
+		)
+
+		return
+	}
+
+	d.client = client
+}
+
 func (d *dnsARecordSetDataSource) Read(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
 	var config aRecordSetConfig
 
@@ -57,16 +83,42 @@ func (d *dnsARecordSetDataSource) Read(ctx context.Context, req datasource.ReadR
 		return
 	}
 
-	host := config.Host.ValueString()
-	a, _, err := lookupIP(host)
-	if err != nil {
-		resp.Diagnostics.AddError(fmt.Sprintf("error looking up A records for %q: ", host), err.Error())
+	if config.UseUpdateServer.ValueBool() && d.client == nil {
+		resp.Diagnostics.AddError("use_update_server enabled, but no update server configured", "If you set use_update_server to true, an update server needs to be configured for the provider")
 		return
 	}
-	sort.Strings(a)
+
+	host := config.Host.ValueString()
+
+	answers := []string{}
+	if !config.UseUpdateServer.ValueBool() || d.client == nil {
+		var err error
+		answers, _, err = lookupIP(host)
+		if err != nil {
+			resp.Diagnostics.AddError(fmt.Sprintf("error looking up A records for %q: ", host), err.Error())
+			return
+		}
+	} else {
+		records, diags := resourceDnsRead_framework_flags(dnsConfig{Name: host}, d.client, dns.TypeA, dns.MsgHdr{RecursionDesired: true})
+		resp.Diagnostics.Append(diags...)
+		if diags.HasError() {
+			return
+		}
+
+		for _, record := range records {
+			addr, _, err := getAVal(record)
+			if err != nil {
+				resp.Diagnostics.AddError("Error querying DNS record:", err.Error())
+				return
+			}
+
+			answers = append(answers, addr)
+		}
+	}
+	sort.Strings(answers)
 
 	var convertDiags diag.Diagnostics
-	config.Addrs, convertDiags = types.ListValueFrom(ctx, config.Addrs.ElementType(ctx), a)
+	config.Addrs, convertDiags = types.ListValueFrom(ctx, config.Addrs.ElementType(ctx), answers)
 	resp.Diagnostics.Append(convertDiags...)
 	if resp.Diagnostics.HasError() {
 		return
@@ -77,7 +129,8 @@ func (d *dnsARecordSetDataSource) Read(ctx context.Context, req datasource.ReadR
 }
 
 type aRecordSetConfig struct {
-	ID    types.String `tfsdk:"id"`
-	Host  types.String `tfsdk:"host"`
-	Addrs types.List   `tfsdk:"addrs"`
+	ID              types.String `tfsdk:"id"`
+	Host            types.String `tfsdk:"host"`
+	UseUpdateServer types.Bool   `tfsdk:"use_update_server"`
+	Addrs           types.List   `tfsdk:"addrs"`
 }
diff --git a/internal/provider/data_dns_a_record_set_test.go b/internal/provider/data_dns_a_record_set_test.go
@@ -30,3 +30,27 @@ data "dns_a_record_set" "test" {
 		},
 	})
 }
+
+func TestAccDataDnsARecordSet_BasicUpdateProvider(t *testing.T) {
+	recordName := "data.dns_a_record_set.test"
+
+	resource.Test(t, resource.TestCase{
+		PreCheck:                 func() { testAccPreCheck(t) },
+		ProtoV5ProviderFactories: testProtoV5ProviderFactories,
+		Steps: []resource.TestStep{
+			{
+				Config: `
+data "dns_a_record_set" "test" {
+  host = "terraform-provider-dns-a.hashicorptest.com"
+  use_update_server = true
+}
+`,
+				Check: resource.ComposeAggregateTestCheckFunc(
+					resource.TestCheckResourceAttr(recordName, "addrs.#", "1"),
+					resource.TestCheckTypeSetElemAttr(recordName, "addrs.*", "127.0.0.1"),
+					resource.TestCheckResourceAttr(recordName, "id", "terraform-provider-dns-a.hashicorptest.com"),
+				),
+			},
+		},
+	})
+}
diff --git a/internal/provider/data_dns_aaaa_record_set.go b/internal/provider/data_dns_aaaa_record_set.go
@@ -12,17 +12,21 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/datasource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/types"
+	"github.com/miekg/dns"
 )
 
 var (
-	_ datasource.DataSource = (*dnsAAAARecordSetDataSource)(nil)
+	_ datasource.DataSource              = (*dnsAAAARecordSetDataSource)(nil)
+	_ datasource.DataSourceWithConfigure = (*dnsAAAARecordSetDataSource)(nil)
 )
 
 func NewDnsAAAARecordSetDataSource() datasource.DataSource {
 	return &dnsAAAARecordSetDataSource{}
 }
 
-type dnsAAAARecordSetDataSource struct{}
+type dnsAAAARecordSetDataSource struct {
+	client *DNSClient
+}
 
 func (d *dnsAAAARecordSetDataSource) Metadata(ctx context.Context, req datasource.MetadataRequest, resp *datasource.MetadataResponse) {
 	resp.TypeName = req.ProviderTypeName + "_aaaa_record_set"
@@ -36,6 +40,10 @@ func (d *dnsAAAARecordSetDataSource) Schema(ctx context.Context, req datasource.
 				Required:    true,
 				Description: "Host to look up.",
 			},
+			"use_update_server": schema.BoolAttribute{
+				Optional:    true,
+				Description: "Whether to use the configured update DNS server",
+			},
 			"addrs": schema.ListAttribute{
 				ElementType: types.StringType,
 				Computed:    true,
@@ -49,6 +57,24 @@ func (d *dnsAAAARecordSetDataSource) Schema(ctx context.Context, req datasource.
 	}
 }
 
+func (d *dnsAAAARecordSetDataSource) Configure(ctx context.Context, req datasource.ConfigureRequest, resp *datasource.ConfigureResponse) {
+	if req.ProviderData == nil {
+		return
+	}
+
+	client, ok := req.ProviderData.(*DNSClient)
+	if !ok {
+		resp.Diagnostics.AddError(
+			"Unexpected Resource Configure Type",
+			fmt.Sprintf("Expected *DNSClient, got: %T. Please report this issue to the provider developers.", req.ProviderData),
+		)
+
+		return
+	}
+
+	d.client = client
+}
+
 func (d *dnsAAAARecordSetDataSource) Read(ctx context.Context, req datasource.ReadRequest, resp *datasource.ReadResponse) {
 	var config aRecordSetConfig
 
@@ -57,16 +83,42 @@ func (d *dnsAAAARecordSetDataSource) Read(ctx context.Context, req datasource.Re
 		return
 	}
 
-	host := config.Host.ValueString()
-	_, aaaa, err := lookupIP(host)
-	if err != nil {
-		resp.Diagnostics.AddError(fmt.Sprintf("error looking up AAAA records for %q: ", host), err.Error())
+	if config.UseUpdateServer.ValueBool() && d.client == nil {
+		resp.Diagnostics.AddError("use_update_server enabled, but no update server configured", "If you set use_update_server to true, an update server needs to be configured for the provider")
 		return
 	}
-	sort.Strings(aaaa)
+
+	host := config.Host.ValueString()
+
+	answers := []string{}
+	if !config.UseUpdateServer.ValueBool() || d.client == nil {
+		var err error
+		_, answers, err = lookupIP(host)
+		if err != nil {
+			resp.Diagnostics.AddError(fmt.Sprintf("error looking up AAAA records for %q: ", host), err.Error())
+			return
+		}
+	} else {
+		records, diags := resourceDnsRead_framework_flags(dnsConfig{Name: host}, d.client, dns.TypeAAAA, dns.MsgHdr{RecursionDesired: true})
+		resp.Diagnostics.Append(diags...)
+		if diags.HasError() {
+			return
+		}
+
+		for _, record := range records {
+			addr, _, err := getAAAAVal(record)
+			if err != nil {
+				resp.Diagnostics.AddError("Error querying DNS record:", err.Error())
+				return
+			}
+
+			answers = append(answers, addr)
+		}
+	}
+	sort.Strings(answers)
 
 	var convertDiags diag.Diagnostics
-	config.Addrs, convertDiags = types.ListValueFrom(ctx, config.Addrs.ElementType(ctx), aaaa)
+	config.Addrs, convertDiags = types.ListValueFrom(ctx, config.Addrs.ElementType(ctx), answers)
 	resp.Diagnostics.Append(convertDiags...)
 	if resp.Diagnostics.HasError() {
 		return
diff --git a/internal/provider/data_dns_aaaa_record_set_test.go b/internal/provider/data_dns_aaaa_record_set_test.go
@@ -30,3 +30,27 @@ data "dns_aaaa_record_set" "test" {
 		},
 	})
 }
+
+func TestAccDataDnsAAAARecordSet_BasicUpdateServer(t *testing.T) {
+	recordName := "data.dns_aaaa_record_set.test"
+
+	resource.Test(t, resource.TestCase{
+		PreCheck:                 func() { testAccPreCheck(t) },
+		ProtoV5ProviderFactories: testProtoV5ProviderFactories,
+		Steps: []resource.TestStep{
+			{
+				Config: `
+data "dns_aaaa_record_set" "test" {
+  host = "terraform-provider-dns-aaaa.hashicorptest.com"
+  use_update_server = true
+}
+`,
+				Check: resource.ComposeAggregateTestCheckFunc(
+					resource.TestCheckResourceAttr(recordName, "addrs.#", "1"),
+					resource.TestCheckTypeSetElemAttr(recordName, "addrs.*", "::1"),
+					resource.TestCheckResourceAttr(recordName, "id", "terraform-provider-dns-aaaa.hashicorptest.com"),
+				),
+			},
+		},
+	})
+}
diff --git a/internal/provider/provider.go b/internal/provider/provider.go
@@ -489,8 +489,6 @@ func exchange(msg *dns.Msg, tsig bool, client *DNSClient) (*dns.Msg, error) {
 		keyname = k
 	}
 
-	msg.RecursionDesired = false
-
 	if tsig && keyname != "" {
 		msg.SetTsig(keyname, keyalgo, 300, time.Now().Unix())
 	}
diff --git a/internal/provider/provider_framework.go b/internal/provider/provider_framework.go
@@ -325,10 +325,13 @@ func (p *dnsProvider) Configure(ctx context.Context, req provider.ConfigureReque
 		keytab:    keytab,
 	}
 
-	resp.ResourceData, configErr = config.Client(ctx)
+	client, configErr := config.Client(ctx)
 	if configErr != nil {
 		resp.Diagnostics.AddError("Error initializing DNS Client:", configErr.Error())
 	}
+
+	resp.ResourceData = client
+	resp.DataSourceData = client
 }
 
 func (p *dnsProvider) Resources(ctx context.Context) []func() resource.Resource {
@@ -496,11 +499,16 @@ func resourceFQDN_framework(config dnsConfig) string {
 }
 
 func resourceDnsRead_framework(config dnsConfig, client *DNSClient, rrType uint16) ([]dns.RR, diag.Diagnostics) {
+	return resourceDnsRead_framework_flags(config, client, rrType, dns.MsgHdr{RecursionDesired: false})
+}
+
+func resourceDnsRead_framework_flags(config dnsConfig, client *DNSClient, rrType uint16, flags dns.MsgHdr) ([]dns.RR, diag.Diagnostics) {
 	var diags diag.Diagnostics
 	fqdn := resourceFQDN_framework(config)
 
 	msg := new(dns.Msg)
 	msg.SetQuestion(fqdn, rrType)
+	msg.RecursionDesired = flags.RecursionDesired
 
 	r, err := exchange(msg, true, client)
 	if err != nil {
diff --git a/internal/provider/provider_framework_test.go b/internal/provider/provider_framework_test.go

Original file line number	Diff line number	Diff line change
`@@ -489,8 +489,6 @@ func exchange(msg dns.Msg, tsig bool, client DNSClient) (*dns.Msg, error) {`
`489`	`489`	`keyname = k`
`490`	`490`	`}`
`491`	`491`
`492`		`- msg.RecursionDesired = false`
`493`		`-`
`494`	`492`	`if tsig && keyname != "" {`
`495`	`493`	`msg.SetTsig(keyname, keyalgo, 300, time.Now().Unix())`
`496`	`494`	`}`