From 2cad04ec10584a19dae5b34434d34dedccaabd26 Mon Sep 17 00:00:00 2001
From: Chris O'Haver <cohaver@infoblox.com>
Date: Fri, 16 Feb 2018 03:45:25 -0500
Subject: [PATCH] plugin/template: add upstream option (#1529)

* add upstream

* docs

* tests
---
 plugin/template/README.md     |  4 ++++
 plugin/template/setup.go      |  8 ++++++++
 plugin/template/setup_test.go | 14 ++++++++++++++
 plugin/template/template.go   |  8 +++++++-
 4 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/plugin/template/README.md b/plugin/template/README.md
index 34a429cbb..1ec699511 100644
--- a/plugin/template/README.md
+++ b/plugin/template/README.md
@@ -18,6 +18,7 @@ template CLASS TYPE [ZONE...] {
     [authority RR]
     [...]
     [rcode CODE]
+    [upstream [ADDRESS...]]
     [fallthrough [ZONE...]]
 }
 ~~~
@@ -29,6 +30,9 @@ template CLASS TYPE [ZONE...] {
 * `answer|additional|authority` **RR** A [RFC 1035](https://tools.ietf.org/html/rfc1035#section-5) style resource record fragment
   built by a [Go template](https://golang.org/pkg/text/template/) that contains the reply.
 * `rcode` **CODE** A response code (`NXDOMAIN, SERVFAIL, ...`). The default is `SUCCESS`.
+* `upstream` [**ADDRESS**...] defines the upstream resolvers used for resolving CNAME.
+  If no **ADDRESS** is given, CoreDNS will resolve CNAMEs against itself. **ADDRESS**
+  can be an IP, an IP:port, or a path to a file structured like resolv.conf.
 * `fallthrough` Continue with the next plugin if the zone matched but no regex matched.
   If specific zones are listed (for example `in-addr.arpa` and `ip6.arpa`), then only queries for
   those zones will be subject to fallthrough.
diff --git a/plugin/template/setup.go b/plugin/template/setup.go
index 616531e49..a1190d481 100644
--- a/plugin/template/setup.go
+++ b/plugin/template/setup.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/pkg/upstream"
 
 	"github.com/mholt/caddy"
 	"github.com/miekg/dns"
@@ -142,6 +143,13 @@ func templateParse(c *caddy.Controller) (handler Handler, err error) {
 			case "fallthrough":
 				t.fall.SetZonesFromArgs(c.RemainingArgs())
 
+			case "upstream":
+				args := c.RemainingArgs()
+				u, err := upstream.NewUpstream(args)
+				if err != nil {
+					return handler, err
+				}
+				t.upstream = u
 			default:
 				return handler, c.ArgErr()
 			}
diff --git a/plugin/template/setup_test.go b/plugin/template/setup_test.go
index 8233aed66..e4d4ff804 100644
--- a/plugin/template/setup_test.go
+++ b/plugin/template/setup_test.go
@@ -133,6 +133,20 @@ func TestSetupParse(t *testing.T) {
 				}`,
 			false,
 		},
+		{
+			`template ANY ANY up.stream.local {
+					answer "up.stream.local 5 IN CNAME up.river.local"
+					upstream
+				}`,
+			false,
+		},
+		{
+			`template ANY ANY up.stream.local {
+					answer "up.stream.local 5 IN CNAME up.river.local"
+					upstream invalid-upstream-argument
+				}`,
+			true,
+		},
 	}
 	for i, test := range tests {
 		c := caddy.NewTestController("dns", test.inputFileRules)
diff --git a/plugin/template/template.go b/plugin/template/template.go
index 0d1d0c251..0320d4c60 100644
--- a/plugin/template/template.go
+++ b/plugin/template/template.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/plugin/pkg/fall"
+	"github.com/coredns/coredns/plugin/pkg/upstream"
 	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
@@ -32,6 +33,7 @@ type template struct {
 	qclass     uint16
 	qtype      uint16
 	fall       fall.F
+	upstream   upstream.Upstream
 }
 
 type templateData struct {
@@ -48,7 +50,7 @@ type templateData struct {
 
 // ServeDNS implements the plugin.Handler interface.
 func (h Handler) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
-	state := request.Request{W: w, Req: r}
+	state := request.Request{W: w, Req: r, Context: ctx}
 
 	zone := plugin.Zones(h.Zones).Matches(state.Name())
 	if zone == "" {
@@ -81,6 +83,10 @@ func (h Handler) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 				return dns.RcodeServerFailure, err
 			}
 			msg.Answer = append(msg.Answer, rr)
+			if rr.Header().Rrtype == dns.TypeCNAME {
+				up, _ := template.upstream.Lookup(state, rr.(*dns.CNAME).Target, dns.TypeA)
+				msg.Answer = append(msg.Answer, up.Answer...)
+			}
 		}
 		for _, additional := range template.additional {
 			rr, err := executeRRTemplate("additional", additional, data)