Added CoreDNS plugin setup and replaced forward

config.go

@@ -253,7 +253,7 @@ const coreDNSConfigTemplate = `.:{{.Port}} {
     hosts {
         fallthrough
     }
-    {{if .UpstreamDNS}}forward . {{range .UpstreamDNS}}{{.}} {{end}}{{end}}
+    {{if .UpstreamDNS}}upstream {{range .UpstreamDNS}}{{.}} {{end}} { bootstrap 8.8.8.8:53 }{{end}}
     {{.Cache}}
     {{.Prometheus}}
 }

coredns.go

@@ -8,6 +8,7 @@ import (
 	"sync" // Include all plugins.
 
 	_ "github.com/AdguardTeam/AdGuardHome/coredns_plugin"
+	_ "github.com/AdguardTeam/AdGuardHome/upstream"
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/coremain"
 	_ "github.com/coredns/coredns/plugin/auto"
@@ -79,6 +80,7 @@ var directives = []string{
 	"loop",
 	"forward",
 	"proxy",
+	"upstream",
 	"erratic",
 	"whoami",
 	"on",

upstream/dns_upstream.go

@@ -42,7 +42,20 @@ func NewDnsUpstream(endpoint string, proto string, tlsServerName string) (Upstre
 // Exchange provides an implementation for the Upstream interface
 func (u *DnsUpstream) Exchange(ctx context.Context, query *dns.Msg) (*dns.Msg, error) {
 
-	resp, err := u.exchange(query)
+	resp, err := u.exchange(u.proto, query)
+
+	// Retry over TCP if response is truncated
+	if err == dns.ErrTruncated && u.proto == "udp" {
+		resp, err = u.exchange("tcp", query)
+	} else if err == dns.ErrTruncated && resp != nil {
+		// Reassemble something to be sent to client
+		m := new(dns.Msg)
+		m.SetReply(query)
+		m.Truncated = true
+		m.Authoritative = true
+		m.Rcode = dns.RcodeSuccess
+		return m, nil
+	}
 
 	if err != nil {
 		resp = &dns.Msg{}
@@ -62,10 +75,10 @@ func (u *DnsUpstream) Close() error {
 
 // Performs a synchronous query. It sends the message m via the conn
 // c and waits for a reply. The conn c is not closed.
-func (u *DnsUpstream) exchange(query *dns.Msg) (r *dns.Msg, err error) {
+func (u *DnsUpstream) exchange(proto string, query *dns.Msg) (r *dns.Msg, err error) {
 
 	// Establish a connection if needed (or reuse cached)
-	conn, err := u.transport.Dial(u.proto)
+	conn, err := u.transport.Dial(proto)
 	if err != nil {
 		return nil, err
 	}

upstream/helpers.go

@@ -87,7 +87,7 @@ func IsAlive(u Upstream) (bool, error) {
 	ping := new(dns.Msg)
 	ping.SetQuestion("ipv4only.arpa.", dns.TypeA)
 
-	resp, err := u.Exchange(nil, ping)
+	resp, err := u.Exchange(context.Background(), ping)
 
 	// If we got a header, we're alright, basically only care about I/O errors 'n stuff.
 	if err != nil && resp != nil {

upstream/setup.go

@@ -0,0 +1,83 @@
+package upstream
+
+import (
+	"github.com/coredns/coredns/core/dnsserver"
+	"github.com/coredns/coredns/plugin"
+	"github.com/mholt/caddy"
+	"log"
+)
+
+func init() {
+	caddy.RegisterPlugin("upstream", caddy.Plugin{
+		ServerType: "dns",
+		Action:     setup,
+	})
+}
+
+// Read the configuration and initialize upstreams
+func setup(c *caddy.Controller) error {
+
+	p, err := setupPlugin(c)
+	if err != nil {
+		return err
+	}
+	config := dnsserver.GetConfig(c)
+	config.AddPlugin(func(next plugin.Handler) plugin.Handler {
+		p.Next = next
+		return p
+	})
+
+	c.OnShutdown(p.onShutdown)
+	return nil
+}
+
+// Read the configuration
+func setupPlugin(c *caddy.Controller) (*UpstreamPlugin, error) {
+
+	p := New()
+
+	log.Println("Initializing the Upstream plugin")
+
+	bootstrap := ""
+	upstreamUrls := []string{}
+	for c.Next() {
+		args := c.RemainingArgs()
+		if len(args) > 0 {
+			upstreamUrls = append(upstreamUrls, args...)
+		}
+		for c.NextBlock() {
+			switch c.Val() {
+			case "bootstrap":
+				if !c.NextArg() {
+					return nil, c.ArgErr()
+				}
+				bootstrap = c.Val()
+			}
+		}
+	}
+
+	for _, url := range upstreamUrls {
+		u, err := NewUpstream(url, bootstrap)
+		if err != nil {
+			log.Printf("Cannot initialize upstream %s", url)
+			return nil, err
+		}
+
+		p.Upstreams = append(p.Upstreams, u)
+	}
+
+	return p, nil
+}
+
+func (p *UpstreamPlugin) onShutdown() error {
+	for i := range p.Upstreams {
+
+		u := p.Upstreams[i]
+		err := u.Close()
+		if err != nil {
+			log.Printf("Error while closing the upstream: %s", err)
+		}
+	}
+
+	return nil
+}

upstream/setup_test.go

@@ -0,0 +1,29 @@
+package upstream
+
+import (
+	"github.com/mholt/caddy"
+	"testing"
+)
+
+func TestSetup(t *testing.T) {
+
+	var tests = []struct {
+		config string
+	}{
+		{`upstream 8.8.8.8`},
+		{`upstream 8.8.8.8 {
+	bootstrap 8.8.8.8:53
+}`},
+		{`upstream tls://1.1.1.1 8.8.8.8 {
+	bootstrap 1.1.1.1
+}`},
+	}
+
+	for _, test := range tests {
+		c := caddy.NewTestController("dns", test.config)
+		err := setup(c)
+		if err != nil {
+			t.Fatalf("Test failed")
+		}
+	}
+}

upstream/upstream.go

@@ -5,8 +5,6 @@ import (
 	"github.com/miekg/dns"
 	"github.com/pkg/errors"
 	"golang.org/x/net/context"
-	"log"
-	"runtime"
 	"time"
 )
 
@@ -14,8 +12,6 @@ const (
 	defaultTimeout = 5 * time.Second
 )
 
-// TODO: Add a helper method for health-checking an upstream (see health.go in coredns)
-
 // Upstream is a simplified interface for proxy destination
 type Upstream interface {
 	Exchange(ctx context.Context, query *dns.Msg) (*dns.Msg, error)
@@ -30,10 +26,10 @@ type UpstreamPlugin struct {
 
 // Initialize the upstream plugin
 func New() *UpstreamPlugin {
-	p := &UpstreamPlugin{}
+	p := &UpstreamPlugin{
+		Upstreams: []Upstream{},
+	}
 
-	// Make sure all resources are cleaned up
-	runtime.SetFinalizer(p, (*UpstreamPlugin).finalizer)
 	return p
 }
 
@@ -56,15 +52,3 @@ func (p *UpstreamPlugin) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *
 
 // Name implements interface for CoreDNS plugin
 func (p *UpstreamPlugin) Name() string { return "upstream" }
-
-func (p *UpstreamPlugin) finalizer() {
-
-	for i := range p.Upstreams {
-
-		u := p.Upstreams[i]
-		err := u.Close()
-		if err != nil {
-			log.Printf("Error while closing the upstream: %s", err)
-		}
-	}
-}

upstream/upstream_test.go

@@ -2,6 +2,7 @@ package upstream
 
 import (
 	"github.com/miekg/dns"
+	"golang.org/x/net/context"
 	"net"
 	"testing"
 )
@@ -169,7 +170,7 @@ func testUpstream(t *testing.T, u Upstream) {
 			{Name: test.name, Qtype: dns.TypeA, Qclass: dns.ClassINET},
 		}
 
-		resp, err := u.Exchange(nil, &req)
+		resp, err := u.Exchange(context.Background(), &req)
 
 		if err != nil {
 			t.Errorf("error while making an upstream request: %s", err)