2 files changed, 31 insertions, 15 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