From af2ef2347acaef9a204b80136264324ca44a8e74 Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Wed, 21 Oct 2015 14:31:53 -0700
Subject: [PATCH] improvements to middleware and URL resolver

---
 cache/cache_test.go                    |  1 -
 router/middleware/location/location.go | 16 +++++++++++++---
 router/router.go                       |  2 ++
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/cache/cache_test.go b/cache/cache_test.go
index f8cf7998..e984da2c 100644
--- a/cache/cache_test.go
+++ b/cache/cache_test.go
@@ -3,7 +3,6 @@ package cache
 import (
 	"testing"
 
-	// "github.com/drone/drone/model"
 	"github.com/franela/goblin"
 	"github.com/gin-gonic/gin"
 )
diff --git a/router/middleware/location/location.go b/router/middleware/location/location.go
index 6fcd5068..08dc9171 100644
--- a/router/middleware/location/location.go
+++ b/router/middleware/location/location.go
@@ -7,9 +7,9 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
-// Hostname is a middleware function that evaluates the http.Request
-// and adds the real hostname and scheme to the context.
-func Hostname(c *gin.Context) {
+// Resolve is a middleware function that resolves the hostname
+// and scheme for the http.Request and adds to the context.
+func Resolve(c *gin.Context) {
 	c.Set("host", resolveHost(c.Request))
 	c.Set("scheme", resolveScheme(c.Request))
 	c.Next()
@@ -56,3 +56,13 @@ func resolveHost(r *http.Request) string {
 		return "localhost:8000"
 	}
 }
+
+// Hostname returns the hostname associated with
+// the current context.
+func Hostname(c *gin.Context) (host string) {
+	v, ok := c.Get("host")
+	if ok {
+		host = v.(string)
+	}
+	return
+}
diff --git a/router/router.go b/router/router.go
index d231c76b..72dde24c 100644
--- a/router/router.go
+++ b/router/router.go
@@ -9,6 +9,7 @@ import (
 	"github.com/drone/drone/controller"
 	"github.com/drone/drone/router/middleware/cache"
 	"github.com/drone/drone/router/middleware/header"
+	"github.com/drone/drone/router/middleware/location"
 	"github.com/drone/drone/router/middleware/session"
 	"github.com/drone/drone/router/middleware/token"
 	"github.com/drone/drone/static"
@@ -20,6 +21,7 @@ func Load(middleware ...gin.HandlerFunc) http.Handler {
 	e.SetHTMLTemplate(template.Load())
 	e.StaticFS("/static", static.FileSystem())
 
+	e.Use(location.Resolve)
 	e.Use(header.NoCache)
 	e.Use(header.Options)
 	e.Use(header.Secure)