Add private DNS support for azure plugin (#3516)

* plugin/azure: fix bug in setting up plugin

Signed-off-by: darshanime <deathbullet@gmail.com>

* plugin/azure: add support for private zones

Signed-off-by: darshanime <deathbullet@gmail.com>

* plugin/azure: change syntax for access level

Signed-off-by: darshanime <deathbullet@gmail.com>

* plugin/azure: change import alias for azure dns

Signed-off-by: darshanime <deathbullet@gmail.com>

* plugin/azure: reword readme, var names

Signed-off-by: darshanime <deathbullet@gmail.com>

* plugin/azure: remove newline in imports

Signed-off-by: darshanime <deathbullet@gmail.com>

* fix import grouping

Co-authored-by: Chris O'Haver <cohaver@infoblox.com>

4 files changed, 188 insertions, 60 deletions

diff --git a/plugin/azure/README.md b/plugin/azure/README.md
index 70cf9a15d..00462290d 100644
--- a/plugin/azure/README.md
+++ b/plugin/azure/README.md
@@ -20,11 +20,12 @@ azure RESOURCE_GROUP:ZONE... {
     subscription SUBSCRIPTION_ID
     environment ENVIRONMENT
     fallthrough [ZONES...]
+    access private
 }
 ~~~
 
 *   **RESOURCE_GROUP:ZONE** is the resource group to which the hosted zones belongs on Azure,
-    and  **ZONE** the zone that contains data.
+    and **ZONE** the zone that contains data.
 
 *   **CLIENT_ID** and **CLIENT_SECRET** are the credentials for Azure, and `tenant` specifies the
     **TENANT_ID** to be used. **SUBSCRIPTION_ID** is the subscription ID. All of these are needed
@@ -36,17 +37,20 @@ azure RESOURCE_GROUP:ZONE... {
     If **ZONES** is omitted, then fallthrough happens for all zones for which the plugin is
     authoritative.
 
+*   `access`  specifies if the zone is `public` or `private`. Default is `public`.
+
 ## Examples
 
-Enable the *azure* plugin with Azure credentials for the zone `example.org`:
+Enable the *azure* plugin with Azure credentials for private zones `example.org`, `example.private`:
 
 ~~~ txt
 example.org {
-    azure resource_group_foo:example.org {
+    azure resource_group_foo:example.org resource_group_foo:example.private {
       tenant 123abc-123abc-123abc-123abc
       client 123abc-123abc-123abc-234xyz
       subscription 123abc-123abc-123abc-563abc
       secret mysecret
+      access private
     }
 }
 ~~~
diff --git a/plugin/azure/azure.go b/plugin/azure/azure.go
index 8432af316..f65503947 100644
--- a/plugin/azure/azure.go
+++ b/plugin/azure/azure.go
@@ -13,53 +13,68 @@ import (
 	"github.com/coredns/coredns/plugin/pkg/upstream"
 	"github.com/coredns/coredns/request"
 
-	azuredns "github.com/Azure/azure-sdk-for-go/profiles/latest/dns/mgmt/dns"
+	publicdns "github.com/Azure/azure-sdk-for-go/profiles/latest/dns/mgmt/dns"
+	privatedns "github.com/Azure/azure-sdk-for-go/profiles/latest/privatedns/mgmt/privatedns"
 	"github.com/miekg/dns"
 )
 
 type zone struct {
-	id   string
-	z    *file.Zone
-	zone string
+	id      string
+	z       *file.Zone
+	zone    string
+	private bool
 }
 
 type zones map[string][]*zone
 
 // Azure is the core struct of the azure plugin.
 type Azure struct {
-	zoneNames []string
-	client    azuredns.RecordSetsClient
-	upstream  *upstream.Upstream
-	zMu       sync.RWMutex
-	zones     zones
+	zoneNames     []string
+	publicClient  publicdns.RecordSetsClient
+	privateClient privatedns.RecordSetsClient
+	upstream      *upstream.Upstream
+	zMu           sync.RWMutex
+	zones         zones
 
 	Next plugin.Handler
 	Fall fall.F
 }
 
 // New validates the input DNS zones and initializes the Azure struct.
-func New(ctx context.Context, dnsClient azuredns.RecordSetsClient, keys map[string][]string) (*Azure, error) {
+func New(ctx context.Context, publicClient publicdns.RecordSetsClient, privateClient privatedns.RecordSetsClient, keys map[string][]string, accessMap map[string]string) (*Azure, error) {
 	zones := make(map[string][]*zone, len(keys))
 	names := make([]string, len(keys))
+	var private bool
 
 	for resourceGroup, znames := range keys {
 		for _, name := range znames {
-			if _, err := dnsClient.ListAllByDNSZone(context.Background(), resourceGroup, name, nil, ""); err != nil {
-				return nil, err
+			switch accessMap[resourceGroup+name] {
+			case "public":
+				if _, err := publicClient.ListAllByDNSZone(context.Background(), resourceGroup, name, nil, ""); err != nil {
+					return nil, err
+				}
+				private = false
+			case "private":
+				if _, err := privateClient.ListComplete(context.Background(), resourceGroup, name, nil, ""); err != nil {
+					return nil, err
+				}
+				private = true
 			}
 
 			fqdn := dns.Fqdn(name)
 			if _, ok := zones[fqdn]; !ok {
 				names = append(names, fqdn)
 			}
-			zones[fqdn] = append(zones[fqdn], &zone{id: resourceGroup, zone: fqdn, z: file.NewZone(fqdn, "")})
+			zones[fqdn] = append(zones[fqdn], &zone{id: resourceGroup, zone: name, private: private, z: file.NewZone(fqdn, "")})
 		}
 	}
+
 	return &Azure{
-		client:    dnsClient,
-		zones:     zones,
-		zoneNames: names,
-		upstream:  upstream.New(),
+		publicClient:  publicClient,
+		privateClient: privateClient,
+		zones:         zones,
+		zoneNames:     names,
+		upstream:      upstream.New(),
 	}, nil
 }
 
@@ -85,14 +100,23 @@ func (h *Azure) Run(ctx context.Context) error {
 }
 
 func (h *Azure) updateZones(ctx context.Context) error {
+	var err error
+	var publicSet publicdns.RecordSetListResultPage
+	var privateSet privatedns.RecordSetListResultPage
+	var newZ *file.Zone
 	errs := make([]string, 0)
 	for zName, z := range h.zones {
 		for i, hostedZone := range z {
-			recordSet, err := h.client.ListByDNSZone(ctx, hostedZone.id, hostedZone.zone, nil, "")
+			if hostedZone.private {
+				privateSet, err = h.privateClient.List(ctx, hostedZone.id, hostedZone.zone, nil, "")
+				newZ = updateZoneFromPrivateResourceSet(privateSet, zName)
+			} else {
+				publicSet, err = h.publicClient.ListByDNSZone(ctx, hostedZone.id, hostedZone.zone, nil, "")
+				newZ = updateZoneFromPublicResourceSet(publicSet, zName)
+			}
 			if err != nil {
 				errs = append(errs, fmt.Sprintf("failed to list resource records for %v from azure: %v", hostedZone.zone, err))
 			}
-			newZ := updateZoneFromResourceSet(recordSet, zName)
 			newZ.Upstream = h.upstream
 			h.zMu.Lock()
 			(*z[i]).z = newZ
@@ -107,7 +131,7 @@ func (h *Azure) updateZones(ctx context.Context) error {
 
 }
 
-func updateZoneFromResourceSet(recordSet azuredns.RecordSetListResultPage, zName string) *file.Zone {
+func updateZoneFromPublicResourceSet(recordSet publicdns.RecordSetListResultPage, zName string) *file.Zone {
 	newZ := file.NewZone(zName, "")
 
 	for _, result := range *(recordSet.Response().Value) {
@@ -196,6 +220,87 @@ func updateZoneFromResourceSet(recordSet azuredns.RecordSetListResultPage, zName
 	return newZ
 }
 
+func updateZoneFromPrivateResourceSet(recordSet privatedns.RecordSetListResultPage, zName string) *file.Zone {
+	newZ := file.NewZone(zName, "")
+
+	for _, result := range *(recordSet.Response().Value) {
+		resultFqdn := *(result.RecordSetProperties.Fqdn)
+		resultTTL := uint32(*(result.RecordSetProperties.TTL))
+		if result.RecordSetProperties.ARecords != nil {
+			for _, A := range *(result.RecordSetProperties.ARecords) {
+				a := &dns.A{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: resultTTL},
+					A: net.ParseIP(*(A.Ipv4Address))}
+				newZ.Insert(a)
+			}
+		}
+		if result.RecordSetProperties.AaaaRecords != nil {
+			for _, AAAA := range *(result.RecordSetProperties.AaaaRecords) {
+				aaaa := &dns.AAAA{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: resultTTL},
+					AAAA: net.ParseIP(*(AAAA.Ipv6Address))}
+				newZ.Insert(aaaa)
+			}
+		}
+
+		if result.RecordSetProperties.MxRecords != nil {
+			for _, MX := range *(result.RecordSetProperties.MxRecords) {
+				mx := &dns.MX{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeMX, Class: dns.ClassINET, Ttl: resultTTL},
+					Preference: uint16(*(MX.Preference)),
+					Mx:         dns.Fqdn(*(MX.Exchange))}
+				newZ.Insert(mx)
+			}
+		}
+
+		if result.RecordSetProperties.PtrRecords != nil {
+			for _, PTR := range *(result.RecordSetProperties.PtrRecords) {
+				ptr := &dns.PTR{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypePTR, Class: dns.ClassINET, Ttl: resultTTL},
+					Ptr: dns.Fqdn(*(PTR.Ptrdname))}
+				newZ.Insert(ptr)
+			}
+		}
+
+		if result.RecordSetProperties.SrvRecords != nil {
+			for _, SRV := range *(result.RecordSetProperties.SrvRecords) {
+				srv := &dns.SRV{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeSRV, Class: dns.ClassINET, Ttl: resultTTL},
+					Priority: uint16(*(SRV.Priority)),
+					Weight:   uint16(*(SRV.Weight)),
+					Port:     uint16(*(SRV.Port)),
+					Target:   dns.Fqdn(*(SRV.Target))}
+				newZ.Insert(srv)
+			}
+		}
+
+		if result.RecordSetProperties.TxtRecords != nil {
+			for _, TXT := range *(result.RecordSetProperties.TxtRecords) {
+				txt := &dns.TXT{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: resultTTL},
+					Txt: *(TXT.Value)}
+				newZ.Insert(txt)
+			}
+		}
+
+		if result.RecordSetProperties.SoaRecord != nil {
+			SOA := result.RecordSetProperties.SoaRecord
+			soa := &dns.SOA{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: resultTTL},
+				Minttl:  uint32(*(SOA.MinimumTTL)),
+				Expire:  uint32(*(SOA.ExpireTime)),
+				Retry:   uint32(*(SOA.RetryTime)),
+				Refresh: uint32(*(SOA.RefreshTime)),
+				Serial:  uint32(*(SOA.SerialNumber)),
+				Mbox:    dns.Fqdn(*(SOA.Email)),
+				Ns:      *(SOA.Host)}
+			newZ.Insert(soa)
+		}
+
+		if result.RecordSetProperties.CnameRecord != nil {
+			CNAME := result.RecordSetProperties.CnameRecord.Cname
+			cname := &dns.CNAME{Hdr: dns.RR_Header{Name: resultFqdn, Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: resultTTL},
+				Target: dns.Fqdn(*CNAME)}
+			newZ.Insert(cname)
+		}
+
+	}
+	return newZ
+}
+
 // ServeDNS implements the plugin.Handler interface.
 func (h *Azure) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
 	state := request.Request{W: w, Req: r}
diff --git a/plugin/azure/setup.go b/plugin/azure/setup.go
index 15ebb7d6f..4f0ec3cbe 100644
--- a/plugin/azure/setup.go
+++ b/plugin/azure/setup.go
@@ -9,7 +9,8 @@ import (
 	"github.com/coredns/coredns/plugin/pkg/fall"
 	clog "github.com/coredns/coredns/plugin/pkg/log"
 
-	azuredns "github.com/Azure/azure-sdk-for-go/profiles/latest/dns/mgmt/dns"
+	publicAzureDNS "github.com/Azure/azure-sdk-for-go/profiles/latest/dns/mgmt/dns"
+	privateAzureDNS "github.com/Azure/azure-sdk-for-go/profiles/latest/privatedns/mgmt/privatedns"
 	azurerest "github.com/Azure/go-autorest/autorest/azure"
 	"github.com/Azure/go-autorest/autorest/azure/auth"
 	"github.com/caddyserver/caddy"
@@ -20,18 +21,23 @@ var log = clog.NewWithPlugin("azure")
 func init() { plugin.Register("azure", setup) }
 
 func setup(c *caddy.Controller) error {
-	env, keys, fall, err := parse(c)
+	env, keys, accessMap, fall, err := parse(c)
 	if err != nil {
 		return plugin.Error("azure", err)
 	}
 	ctx := context.Background()
 
-	dnsClient := azuredns.NewRecordSetsClient(env.Values[auth.SubscriptionID])
-	if dnsClient.Authorizer, err = env.GetAuthorizer(); err != nil {
+	publicDNSClient := publicAzureDNS.NewRecordSetsClient(env.Values[auth.SubscriptionID])
+	if publicDNSClient.Authorizer, err = env.GetAuthorizer(); err != nil {
 		return plugin.Error("azure", err)
 	}
 
-	h, err := New(ctx, dnsClient, keys)
+	privateDNSClient := privateAzureDNS.NewRecordSetsClient(env.Values[auth.SubscriptionID])
+	if privateDNSClient.Authorizer, err = env.GetAuthorizer(); err != nil {
+		return plugin.Error("azure", err)
+	}
+
+	h, err := New(ctx, publicDNSClient, privateDNSClient, keys, accessMap)
 	if err != nil {
 		return plugin.Error("azure", err)
 	}
@@ -47,13 +53,17 @@ func setup(c *caddy.Controller) error {
 	return nil
 }
 
-func parse(c *caddy.Controller) (auth.EnvironmentSettings, map[string][]string, fall.F, error) {
+func parse(c *caddy.Controller) (auth.EnvironmentSettings, map[string][]string, map[string]string, fall.F, error) {
 	resourceGroupMapping := map[string][]string{}
+	accessMap := map[string]string{}
 	resourceGroupSet := map[string]struct{}{}
 	azureEnv := azurerest.PublicCloud
 	env := auth.EnvironmentSettings{Values: map[string]string{}}
 
 	var fall fall.F
+	var access string
+	var resourceGroup string
+	var zoneName string
 
 	for c.Next() {
 		args := c.RemainingArgs()
@@ -61,60 +71,70 @@ func parse(c *caddy.Controller) (auth.EnvironmentSettings, map[string][]string,
 		for i := 0; i < len(args); i++ {
 			parts := strings.SplitN(args[i], ":", 2)
 			if len(parts) != 2 {
-				return env, resourceGroupMapping, fall, c.Errf("invalid resource group/zone: %q", args[i])
+				return env, resourceGroupMapping, accessMap, fall, c.Errf("invalid resource group/zone: %q", args[i])
 			}
-			resourceGroup, zoneName := parts[0], parts[1]
+			resourceGroup, zoneName = parts[0], parts[1]
 			if resourceGroup == "" || zoneName == "" {
-				return env, resourceGroupMapping, fall, c.Errf("invalid resource group/zone: %q", args[i])
+				return env, resourceGroupMapping, accessMap, fall, c.Errf("invalid resource group/zone: %q", args[i])
 			}
-			if _, ok := resourceGroupSet[args[i]]; ok {
-				return env, resourceGroupMapping, fall, c.Errf("conflicting zone: %q", args[i])
+			if _, ok := resourceGroupSet[resourceGroup+zoneName]; ok {
+				return env, resourceGroupMapping, accessMap, fall, c.Errf("conflicting zone: %q", args[i])
 			}
 
-			resourceGroupSet[args[i]] = struct{}{}
+			resourceGroupSet[resourceGroup+zoneName] = struct{}{}
+			accessMap[resourceGroup+zoneName] = "public"
 			resourceGroupMapping[resourceGroup] = append(resourceGroupMapping[resourceGroup], zoneName)
 		}
+
 		for c.NextBlock() {
 			switch c.Val() {
 			case "subscription":
 				if !c.NextArg() {
-					return env, resourceGroupMapping, fall, c.ArgErr()
+					return env, resourceGroupMapping, accessMap, fall, c.ArgErr()
 				}
 				env.Values[auth.SubscriptionID] = c.Val()
 			case "tenant":
 				if !c.NextArg() {
-					return env, resourceGroupMapping, fall, c.ArgErr()
+					return env, resourceGroupMapping, accessMap, fall, c.ArgErr()
 				}
 				env.Values[auth.TenantID] = c.Val()
 			case "client":
 				if !c.NextArg() {
-					return env, resourceGroupMapping, fall, c.ArgErr()
+					return env, resourceGroupMapping, accessMap, fall, c.ArgErr()
 				}
 				env.Values[auth.ClientID] = c.Val()
 			case "secret":
 				if !c.NextArg() {
-					return env, resourceGroupMapping, fall, c.ArgErr()
+					return env, resourceGroupMapping, accessMap, fall, c.ArgErr()
 				}
 				env.Values[auth.ClientSecret] = c.Val()
 			case "environment":
 				if !c.NextArg() {
-					return env, resourceGroupMapping, fall, c.ArgErr()
+					return env, resourceGroupMapping, accessMap, fall, c.ArgErr()
 				}
 				env.Values[auth.ClientSecret] = c.Val()
 				var err error
 				if azureEnv, err = azurerest.EnvironmentFromName(c.Val()); err != nil {
-					return env, resourceGroupMapping, fall, c.Errf("cannot set azure environment: %q", err.Error())
+					return env, resourceGroupMapping, accessMap, fall, c.Errf("cannot set azure environment: %q", err.Error())
 				}
 			case "fallthrough":
 				fall.SetZonesFromArgs(c.RemainingArgs())
+			case "access":
+				if !c.NextArg() {
+					return env, resourceGroupMapping, accessMap, fall, c.ArgErr()
+				}
+				access = c.Val()
+				if access != "public" && access != "private" {
+					return env, resourceGroupMapping, accessMap, fall, c.Errf("invalid access value: can be public/private, found: %s", access)
+				}
+				accessMap[resourceGroup+zoneName] = access
 			default:
-				return env, resourceGroupMapping, fall, c.Errf("unknown property: %q", c.Val())
+				return env, resourceGroupMapping, accessMap, fall, c.Errf("unknown property: %q", c.Val())
 			}
 		}
 	}
 
 	env.Values[auth.Resource] = azureEnv.ResourceManagerEndpoint
 	env.Environment = azureEnv
-
-	return env, resourceGroupMapping, fall, nil
+	return env, resourceGroupMapping, accessMap, fall, nil
 }
diff --git a/plugin/azure/setup_test.go b/plugin/azure/setup_test.go
index c0b22d581..a8df18cd7 100644
--- a/plugin/azure/setup_test.go
+++ b/plugin/azure/setup_test.go
@@ -18,34 +18,30 @@ func TestSetup(t *testing.T) {
     tenant
 }`, true},
 		{`azure resource_set:zone {
-    tenant
-}`, true},
+    tenant abc
+}`, false},
 		{`azure resource_set:zone {
     client
 }`, true},
 		{`azure resource_set:zone {
-    secret
-}`, true},
+    client abc
+}`, false},
 		{`azure resource_set:zone {
     subscription
 }`, true},
 		{`azure resource_set:zone {
-    upstream 10.0.0.1
-}`, true},
-
-		{`azure resource_set:zone {
-    upstream
-}`, true},
+    subscription abc
+}`, false},
 		{`azure resource_set:zone {
-    foobar
+    foo
 }`, true},
 		{`azure resource_set:zone {
     tenant tenant_id
     client client_id
     secret client_secret
     subscription subscription_id
+    access public
 }`, false},
-
 		{`azure resource_set:zone {
     fallthrough
 }`, false},
@@ -56,16 +52,19 @@ func TestSetup(t *testing.T) {
 			fallthrough
 		}`, true},
 		{`azure resource_set:zone,zone2 {
-			fallthrough
+			access private
 		}`, false},
-		{`azure resource-set {
-			fallthrough
+		{`azure resource-set:zone {
+			access public
+		}`, false},
+		{`azure resource-set:zone {
+			access foo
 		}`, true},
 	}
 
 	for i, test := range tests {
 		c := caddy.NewTestController("dns", test.body)
-		if _, _, _, err := parse(c); (err == nil) == test.expectedError {
+		if _, _, _, _, err := parse(c); (err == nil) == test.expectedError {
 			t.Fatalf("Unexpected errors: %v in test: %d\n\t%s", err, i, test.body)
 		}
 	}