From ea77f2a2caea7973615c93b787f98e67ec6f31a3 Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Thu, 10 Aug 2017 21:31:36 +0100
Subject: [PATCH] core: replace GetMiddleware (#885)

* core: replace GetMiddleware

See the discussion in #881. GetMiddleware would add a `nil` middleware
to the callstack thereby breaking functionality.

This PR drops it in favor of RegisterHandler which is a completely
standalone registry for middleware that want to let it self know to
other middleware.

Currenly *autopath* uses this to call *kubernetes*'s AutoPath method
for dynamic autopathing.

* Drop GetMiddleware

* Register metrics

* drop the panic
---
 core/dnsserver/config.go       | 20 +++++---------------
 core/dnsserver/register.go     | 26 ++++++++++++++++++++++++++
 middleware/auto/setup.go       |  4 +++-
 middleware/autopath/setup.go   | 23 ++++++++++++-----------
 middleware/kubernetes/setup.go |  3 +++
 middleware/metrics/setup.go    |  3 +++
 middleware/proxy/setup.go      |  2 +-
 7 files changed, 53 insertions(+), 28 deletions(-)

diff --git a/core/dnsserver/config.go b/core/dnsserver/config.go
index c0f879a31..32751b326 100644
--- a/core/dnsserver/config.go
+++ b/core/dnsserver/config.go
@@ -38,6 +38,11 @@ type Config struct {
 
 	// Compiled middleware stack.
 	middlewareChain middleware.Handler
+
+	// Middleware interested in announcing that they exist, so other middleware can call methods
+	// on them should register themselves here. The name should be the name as return by the
+	// Handler's Name method.
+	Registry map[string]middleware.Handler
 }
 
 // GetConfig gets the Config that corresponds to c.
@@ -53,18 +58,3 @@ func GetConfig(c *caddy.Controller) *Config {
 	ctx.saveConfig(c.Key, &Config{})
 	return GetConfig(c)
 }
-
-// GetMiddleware returns the middleware handler that has been added to the config under name.
-// This is useful to inspect if a certain middleware is active in this server.
-// Note that this is order dependent and the order is defined in directives.go, i.e. if your middleware
-// comes before the middleware you are checking; it will not be there (yet).
-func GetMiddleware(c *caddy.Controller, name string) middleware.Handler {
-	conf := GetConfig(c)
-	for _, h := range conf.Middleware {
-		x := h(nil)
-		if name == x.Name() {
-			return x
-		}
-	}
-	return nil
-}
diff --git a/core/dnsserver/register.go b/core/dnsserver/register.go
index 4b0ec5945..b5589e165 100644
--- a/core/dnsserver/register.go
+++ b/core/dnsserver/register.go
@@ -124,6 +124,32 @@ func (c *Config) AddMiddleware(m middleware.Middleware) {
 	c.Middleware = append(c.Middleware, m)
 }
 
+// RegisterHandler adds a handler to a site's handler registration. Handlers
+// should use this if the want to announce that they exist to other middleware.
+func (c *Config) RegisterHandler(h middleware.Handler) {
+	if c.Registry == nil {
+		c.Registry = make(map[string]middleware.Handler)
+	}
+
+	// Just overwrite...
+	c.Registry[h.Name()] = h
+}
+
+// GetHandler returns the middleware handler that has been added to the config under its name.
+// This is useful to inspect if a certain middleware is active in this server.
+// Note that this is order dependent and the order is defined in directives.go, i.e. if your middleware
+// comes before the middleware you are checking; it will not be there (yet).
+// See RegisterHandler on how to register the middleware with this server.
+func (c *Config) GetHandler(name string) middleware.Handler {
+	if c.Registry == nil {
+		return nil
+	}
+	if h, ok := c.Registry[name]; ok {
+		return h
+	}
+	return nil
+}
+
 // groupSiteConfigsByListenAddr groups site configs by their listen
 // (bind) address, so sites that use the same listener can be served
 // on the same server instance. The return value maps the listen
diff --git a/middleware/auto/setup.go b/middleware/auto/setup.go
index 98eefca4a..bd5949d7e 100644
--- a/middleware/auto/setup.go
+++ b/middleware/auto/setup.go
@@ -32,7 +32,9 @@ func setup(c *caddy.Controller) error {
 	}
 
 	// If we have enabled prometheus we should add newly discovered zones to it.
-	met := dnsserver.GetMiddleware(c, "prometheus")
+	// This does not have to happen in a on.Startup because monitoring is one of the first
+	// to be initialized.
+	met := dnsserver.GetConfig(c).GetHandler("prometheus")
 	if met != nil {
 		a.metrics = met.(*metrics.Metrics)
 	}
diff --git a/middleware/autopath/setup.go b/middleware/autopath/setup.go
index 804b61855..1855dd440 100644
--- a/middleware/autopath/setup.go
+++ b/middleware/autopath/setup.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/middleware"
+	"github.com/coredns/coredns/middleware/kubernetes"
 
 	"github.com/mholt/caddy"
 	"github.com/miekg/dns"
@@ -19,23 +20,21 @@ func init() {
 }
 
 func setup(c *caddy.Controller) error {
-	ap, _, err := autoPathParse(c)
+	ap, mw, err := autoPathParse(c)
 	if err != nil {
 		return middleware.Error("autopath", err)
 	}
 
+	// Do this in OnStartup, so all middleware has been initialized.
 	c.OnStartup(func() error {
-		// Do this in OnStartup, so all middleware has been initialized.
 		// TODO(miek): fabricate test to proof this is not thread safe.
-		// TODO(miek): disable this for now: See https://github.com/coredns/coredns/issues/881
-		/*
-			switch mw {
-			case "kubernetes":
-				if k, ok := m.(kubernetes.Kubernetes); ok {
-					ap.searchFunc = k.AutoPath
-				}
-			}
-		*/
+		m := dnsserver.GetConfig(c).GetHandler(mw)
+		if m == nil {
+			return nil
+		}
+		if k, ok := m.(kubernetes.Kubernetes); ok {
+			ap.searchFunc = k.AutoPath
+		}
 		return nil
 	})
 
@@ -47,6 +46,8 @@ func setup(c *caddy.Controller) error {
 	return nil
 }
 
+// allowedMiddleware has a list of middleware that can be used by autopath. For this to work, they
+// need to register themselves with dnsserver.RegisterHandler.
 var allowedMiddleware = map[string]bool{
 	"@kubernetes": true,
 }
diff --git a/middleware/kubernetes/setup.go b/middleware/kubernetes/setup.go
index 1c77c11ec..6b8cbba6e 100644
--- a/middleware/kubernetes/setup.go
+++ b/middleware/kubernetes/setup.go
@@ -51,6 +51,9 @@ func setup(c *caddy.Controller) error {
 		return kubernetes
 	})
 
+	// Also register kubernetes for use in autopath.
+	dnsserver.GetConfig(c).RegisterHandler(kubernetes)
+
 	return nil
 }
 
diff --git a/middleware/metrics/setup.go b/middleware/metrics/setup.go
index 1ed455ffc..b6b8e5291 100644
--- a/middleware/metrics/setup.go
+++ b/middleware/metrics/setup.go
@@ -38,6 +38,9 @@ func setup(c *caddy.Controller) error {
 	}
 	c.OnFinalShutdown(m.OnShutdown)
 
+	// Also register metrics for use in other middleware.
+	dnsserver.GetConfig(c).RegisterHandler(m)
+
 	return nil
 }
 
diff --git a/middleware/proxy/setup.go b/middleware/proxy/setup.go
index de979a5df..bd036d4cb 100644
--- a/middleware/proxy/setup.go
+++ b/middleware/proxy/setup.go
@@ -20,7 +20,7 @@ func setup(c *caddy.Controller) error {
 		return middleware.Error("proxy", err)
 	}
 
-	t := dnsserver.GetMiddleware(c, "trace")
+	t := dnsserver.GetConfig(c).GetHandler("trace")
 	P := &Proxy{Trace: t}
 	dnsserver.GetConfig(c).AddMiddleware(func(next middleware.Handler) middleware.Handler {
 		P.Next = next