cleaning up the middleware and adding caching with TTL

Makefile

@@ -36,7 +36,7 @@ build:
 	go build
 
 build_static:
-	go build --ldflags '-extldflags "-static"' -o drone_static
+	go build --ldflags '-extldflags "-static" -X main.version=$(BUILD_NUMBER)' -o drone_static
 
 test:
 	go test -cover $(PACKAGES)

controller/user.go

@@ -9,19 +9,8 @@ import (
 	"github.com/drone/drone/router/middleware/context"
 	"github.com/drone/drone/router/middleware/session"
 	"github.com/drone/drone/shared/token"
-	"github.com/hashicorp/golang-lru"
 )
 
-var cache *lru.Cache
-
-func init() {
-	var err error
-	cache, err = lru.New(1028)
-	if err != nil {
-		panic(err)
-	}
-}
-
 func GetSelf(c *gin.Context) {
 	c.IndentedJSON(200, session.User(c))
 }
@@ -52,11 +41,9 @@ func GetRemoteRepos(c *gin.Context) {
 	user := session.User(c)
 	remote := context.Remote(c)
 
-	// attempt to get the repository list from the
-	// cache since the operation is expensive
-	v, ok := cache.Get(user.Login)
+	reposv, ok := c.Get("repos")
 	if ok {
-		c.IndentedJSON(http.StatusOK, v)
+		c.IndentedJSON(http.StatusOK, reposv)
 		return
 	}
 
@@ -65,7 +52,8 @@ func GetRemoteRepos(c *gin.Context) {
 		c.AbortWithStatus(http.StatusInternalServerError)
 		return
 	}
-	cache.Add(user.Login, repos)
+
+	c.Set("repos", repos)
 	c.IndentedJSON(http.StatusOK, repos)
 }
 

drone.go

@@ -7,6 +7,7 @@ import (
 	"github.com/drone/drone/remote"
 	"github.com/drone/drone/router"
 	"github.com/drone/drone/router/middleware/context"
+	"github.com/drone/drone/router/middleware/header"
 	"github.com/drone/drone/shared/database"
 	"github.com/drone/drone/shared/envconfig"
 	"github.com/drone/drone/shared/server"
@@ -14,6 +15,10 @@ import (
 	"github.com/Sirupsen/logrus"
 )
 
+// build revision number populated by the continuous
+// integration server at compile time.
+var build string = "custom"
+
 var (
 	dotenv = flag.String("config", ".env", "")
 	debug  = flag.Bool("debug", false, "")
@@ -43,6 +48,7 @@ func main() {
 	server_ := server.Load(env)
 	server_.Run(
 		router.Load(
+			header.Version(build),
 			context.SetDatabase(database_),
 			context.SetRemote(remote_),
 			context.SetEngine(engine_),

router/middleware/cache/cache.go

@@ -0,0 +1,49 @@
+package cache
+
+import (
+	"time"
+
+	"github.com/hashicorp/golang-lru"
+)
+
+// single instance of a thread-safe lru cache
+var cache *lru.Cache
+
+func init() {
+	var err error
+	cache, err = lru.New(2048)
+	if err != nil {
+		panic(err)
+	}
+}
+
+// item is a simple wrapper around a cacheable object
+// that tracks the ttl for item expiration in the cache.
+type item struct {
+	value interface{}
+	ttl   time.Time
+}
+
+// set adds the key value pair to the cache with the
+// specified ttl expiration.
+func set(key string, value interface{}, ttl int64) {
+	ttlv := time.Now().Add(time.Duration(ttl) * time.Second)
+	cache.Add(key, &item{value, ttlv})
+}
+
+// get gets the value from the cache for the given key.
+// if the value does not exist, a nil value is returned.
+// if the value exists, but is expired, the value is returned
+// with a bool flag set to true.
+func get(key string) (interface{}, bool) {
+	v, ok := cache.Get(key)
+	if !ok {
+		return nil, false
+	}
+	vv := v.(*item)
+	expired := vv.ttl.Before(time.Now())
+	if expired {
+		cache.Remove(key)
+	}
+	return vv.value, expired
+}

router/middleware/cache/cache_test.go

@@ -0,0 +1,40 @@
+package cache
+
+import (
+	"testing"
+
+	"github.com/franela/goblin"
+)
+
+func TestCache(t *testing.T) {
+
+	g := goblin.Goblin(t)
+	g.Describe("Cache", func() {
+
+		g.BeforeEach(func() {
+			cache.Purge()
+		})
+
+		g.It("should set and get item", func() {
+			set("foo", "bar", 1000)
+			val, expired := get("foo")
+			g.Assert(val).Equal("bar")
+			g.Assert(expired).Equal(false)
+		})
+
+		g.It("should return nil when item not found", func() {
+			val, expired := get("foo")
+			g.Assert(val == nil).IsTrue()
+			g.Assert(expired).Equal(false)
+		})
+
+		g.It("should get expired item and purge", func() {
+			set("foo", "bar", -900)
+			val, expired := get("foo")
+			g.Assert(val).Equal("bar")
+			g.Assert(expired).Equal(true)
+			val, _ = get("foo")
+			g.Assert(val == nil).IsTrue()
+		})
+	})
+}

router/middleware/cache/perms.go

@@ -0,0 +1,52 @@
+package cache
+
+import (
+	"fmt"
+
+	"github.com/drone/drone/model"
+	"github.com/gin-gonic/gin"
+)
+
+const permKey = "perm"
+
+// Perms is a middleware function that attempts to cache the
+// user's remote rempository permissions (ie in GitHub) to minimize
+// remote calls that might be expensive, slow or rate-limited.
+func Perms(c *gin.Context) {
+	var (
+		owner   = c.Param("owner")
+		name    = c.Param("name")
+		user, _ = c.Get("user")
+	)
+
+	if user == nil {
+		c.Next()
+		return
+	}
+
+	key := fmt.Sprintf("perm/%s/%s/%s",
+		user.(*model.User).Login,
+		owner,
+		name,
+	)
+
+	// if the item already exists in the cache
+	// we can continue the middleware chain and
+	// exit afterwards.
+	v, _ := get(key)
+	if v != nil {
+		c.Set("perm", v)
+		c.Next()
+		return
+	}
+
+	// otherwise, if the item isn't cached we execute
+	// the middleware chain and then cache the permissions
+	// after the request is processed.
+	c.Next()
+
+	perm, ok := c.Get("perm")
+	if ok {
+		set(key, perm, 86400) // 24 hours
+	}
+}

router/middleware/cache/perms_test.go

@@ -0,0 +1,60 @@
+package cache
+
+import (
+	"testing"
+
+	"github.com/drone/drone/model"
+	"github.com/franela/goblin"
+	"github.com/gin-gonic/gin"
+)
+
+func TestPermCache(t *testing.T) {
+
+	g := goblin.Goblin(t)
+	g.Describe("Perm Cache", func() {
+
+		g.BeforeEach(func() {
+			cache.Purge()
+		})
+
+		g.It("should skip when no user session", func() {
+			c := &gin.Context{}
+			c.Params = gin.Params{
+				gin.Param{Key: "owner", Value: "octocat"},
+				gin.Param{Key: "name", Value: "hello-world"},
+			}
+
+			Perms(c)
+
+			_, ok := c.Get("perm")
+			g.Assert(ok).IsFalse()
+		})
+
+		g.It("should get perms from cache", func() {
+			c := &gin.Context{}
+			c.Params = gin.Params{
+				gin.Param{Key: "owner", Value: "octocat"},
+				gin.Param{Key: "name", Value: "hello-world"},
+			}
+			c.Set("user", fakeUser)
+			set("perm/octocat/octocat/hello-world", fakePerm, 999)
+
+			Perms(c)
+
+			perm, ok := c.Get("perm")
+			g.Assert(ok).IsTrue()
+			g.Assert(perm).Equal(fakePerm)
+		})
+
+	})
+}
+
+var fakePerm = &model.Perm{
+	Pull:  true,
+	Push:  true,
+	Admin: true,
+}
+
+var fakeUser = &model.User{
+	Login: "octocat",
+}

router/middleware/cache/repos.go

@@ -0,0 +1,44 @@
+package cache
+
+import (
+	"fmt"
+
+	"github.com/drone/drone/model"
+	"github.com/gin-gonic/gin"
+)
+
+// Repos is a middleware function that attempts to cache the
+// user's list of remote repositories (ie in GitHub) to minimize
+// remote calls that might be expensive, slow or rate-limited.
+func Repos(c *gin.Context) {
+	var user, _ = c.Get("user")
+
+	if user == nil {
+		c.Next()
+		return
+	}
+
+	key := fmt.Sprintf("repos/%s",
+		user.(*model.User).Login,
+	)
+
+	// if the item already exists in the cache
+	// we can continue the middleware chain and
+	// exit afterwards.
+	v, _ := get(key)
+	if v != nil {
+		c.Set("repos", v)
+		c.Next()
+		return
+	}
+
+	// otherwise, if the item isn't cached we execute
+	// the middleware chain and then cache the permissions
+	// after the request is processed.
+	c.Next()
+
+	repos, ok := c.Get("repos")
+	if ok {
+		set(key, repos, 86400) // 24 hours
+	}
+}

router/middleware/cache/repos_test.go

@@ -0,0 +1,46 @@
+package cache
+
+import (
+	"testing"
+
+	"github.com/drone/drone/model"
+	"github.com/franela/goblin"
+	"github.com/gin-gonic/gin"
+)
+
+func TestReposCache(t *testing.T) {
+
+	g := goblin.Goblin(t)
+	g.Describe("Repo List Cache", func() {
+
+		g.BeforeEach(func() {
+			cache.Purge()
+		})
+
+		g.It("should skip when no user session", func() {
+			c := &gin.Context{}
+
+			Perms(c)
+
+			_, ok := c.Get("perm")
+			g.Assert(ok).IsFalse()
+		})
+
+		g.It("should get repos from cache", func() {
+			c := &gin.Context{}
+			c.Set("user", fakeUser)
+			set("repos/octocat", fakeRepos, 999)
+
+			Repos(c)
+
+			repos, ok := c.Get("repos")
+			g.Assert(ok).IsTrue()
+			g.Assert(repos).Equal(fakeRepos)
+		})
+
+	})
+}
+
+var fakeRepos = []*model.RepoLite{
+	{Owner: "octocat", Name: "hello-world"},
+}

router/middleware/header/header.go

@@ -7,34 +7,53 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
-func SetHeaders() gin.HandlerFunc {
-	return func(c *gin.Context) {
+var version string
 
-		c.Writer.Header().Add("Access-Control-Allow-Origin", "*")
-		c.Writer.Header().Add("X-Frame-Options", "DENY")
-		c.Writer.Header().Add("X-Content-Type-Options", "nosniff")
-		c.Writer.Header().Add("X-XSS-Protection", "1; mode=block")
-		c.Writer.Header().Add("Cache-Control", "no-cache")
-		c.Writer.Header().Add("Cache-Control", "no-store")
-		c.Writer.Header().Add("Cache-Control", "max-age=0")
-		c.Writer.Header().Add("Cache-Control", "must-revalidate")
-		c.Writer.Header().Add("Cache-Control", "value")
-		c.Writer.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
-		c.Writer.Header().Set("Expires", "Thu, 01 Jan 1970 00:00:00 GMT")
-		//c.Writer.Header().Set("Content-Security-Policy", "script-src 'self' https://cdnjs.cloudflare.com")
+// NoCache is a middleware function that appends headers
+// to prevent the client from caching the HTTP response.
+func NoCache(c *gin.Context) {
+	c.Header("Cache-Control", "no-cache, no-store, max-age=0, must-revalidate, value")
+	c.Header("Expires", "Thu, 01 Jan 1970 00:00:00 GMT")
+	c.Header("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
+	c.Next()
+}
+
+// Options is a middleware function that appends headers
+// for options requests and aborts then exits the middleware
+// chain and ends the request.
+func Options(c *gin.Context) {
+	if c.Request.Method != "OPTIONS" {
+		c.Next()
+	} else {
+		c.Header("Access-Control-Allow-Methods", "GET,POST,PUT,PATCH,DELETE,OPTIONS")
+		c.Header("Access-Control-Allow-Headers", "Authorization")
+		c.Header("Allow", "HEAD,GET,POST,PUT,PATCH,DELETE,OPTIONS")
+		c.Header("Content-Type", "application/json")
+		c.AbortWithStatus(200)
+	}
+}
+
+// Secure is a middleware function that appends security
+// and resource access headers.
+func Secure(c *gin.Context) {
+	c.Header("Access-Control-Allow-Origin", "*")
+	c.Header("X-Frame-Options", "DENY")
+	c.Header("X-Content-Type-Options", "nosniff")
+	c.Header("X-XSS-Protection", "1; mode=block")
 	if c.Request.TLS != nil {
-			c.Writer.Header().Add("Strict-Transport-Security", "max-age=31536000")
+		c.Header("Strict-Transport-Security", "max-age=31536000")
 	}
 
-		if c.Request.Method == "OPTIONS" {
-			c.Writer.Header().Set("Access-Control-Allow-Methods", "GET,POST,PUT,PATCH,DELETE,OPTIONS")
-			c.Writer.Header().Set("Access-Control-Allow-Headers", "Authorization")
-			c.Writer.Header().Set("Allow", "HEAD,GET,POST,PUT,PATCH,DELETE,OPTIONS")
-			c.Writer.Header().Set("Content-Type", "application/json")
-			c.Writer.WriteHeader(200)
-			return
+	// Also consider adding Content-Security-Policy headers
+	// c.Header("Content-Security-Policy", "script-src 'self' https://cdnjs.cloudflare.com")
 }
 
+// Version is a middleware function that appends the Drone
+// version information to the HTTP response. This is intended
+// for debugging and troubleshooting.
+func Version(version string) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		c.Header("X-DRONE-VERSION", "0.4.0-beta+"+version)
 		c.Next()
 	}
 }

router/middleware/location/location.go

@@ -0,0 +1,58 @@
+package location
+
+import (
+	"net/http"
+	"strings"
+
+	"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) {
+	c.Set("host", resolveHost(c.Request))
+	c.Set("scheme", resolveScheme(c.Request))
+	c.Next()
+}
+
+// resolveScheme is a helper function that evaluates the http.Request
+// and returns the scheme, HTTP or HTTPS. It is able to detect,
+// using the X-Forwarded-Proto, if the original request was HTTPS
+// and routed through a reverse proxy with SSL termination.
+func resolveScheme(r *http.Request) string {
+	switch {
+	case r.URL.Scheme == "https":
+		return "https"
+	case r.TLS != nil:
+		return "https"
+	case strings.HasPrefix(r.Proto, "HTTPS"):
+		return "https"
+	case r.Header.Get("X-Forwarded-Proto") == "https":
+		return "https"
+	default:
+		return "http"
+	}
+}
+
+// resolveHost is a helper function that evaluates the http.Request
+// and returns the hostname. It is able to detect, using the
+// X-Forarded-For header, the original hostname when routed
+// through a reverse proxy.
+func resolveHost(r *http.Request) string {
+	switch {
+	case len(r.Host) != 0:
+		return r.Host
+	case len(r.URL.Host) != 0:
+		return r.URL.Host
+	case len(r.Header.Get("X-Forwarded-For")) != 0:
+		return r.Header.Get("X-Forwarded-For")
+	case len(r.Header.Get("X-Host")) != 0:
+		return r.Header.Get("X-Host")
+	case len(r.Header.Get("XFF")) != 0:
+		return r.Header.Get("XFF")
+	case len(r.Header.Get("X-Real-IP")) != 0:
+		return r.Header.Get("X-Real-IP")
+	default:
+		return "localhost:8000"
+	}
+}

router/middleware/session/repo.go

@@ -1,7 +1,6 @@
 package session
 
 import (
-	"fmt"
 	"net/http"
 
 	"github.com/drone/drone/model"
@@ -10,29 +9,30 @@ import (
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/gin-gonic/gin"
-	"github.com/hashicorp/golang-lru"
 )
 
-var cache *lru.Cache
-
-func init() {
-	var err error
-	cache, err = lru.New(1028)
-	if err != nil {
-		panic(err)
-	}
-}
-
 func Repo(c *gin.Context) *model.Repo {
 	v, ok := c.Get("repo")
 	if !ok {
 		return nil
 	}
-	u, ok := v.(*model.Repo)
+	r, ok := v.(*model.Repo)
 	if !ok {
 		return nil
 	}
-	return u
+	return r
+}
+
+func Repos(c *gin.Context) []*model.RepoLite {
+	v, ok := c.Get("repos")
+	if !ok {
+		return nil
+	}
+	r, ok := v.([]*model.RepoLite)
+	if !ok {
+		return nil
+	}
+	return r
 }
 
 func SetRepo() gin.HandlerFunc {
@@ -106,10 +106,8 @@ func SetPerm() gin.HandlerFunc {
 		if user != nil {
 			// attempt to get the permissions from a local cache
 			// just to avoid excess API calls to GitHub
-			key := fmt.Sprintf("%d.%d", user.ID, repo.ID)
-			val, ok := cache.Get(key)
+			val, ok := c.Get("perm")
 			if ok {
-				c.Set("perm", val.(*model.Perm))
 				c.Next()
 
 				log.Debugf("%s using cached %+v permission to %s",
@@ -161,13 +159,6 @@ func SetPerm() gin.HandlerFunc {
 		}
 
 		if user != nil {
-
-			// cache the updated repository permissions to
-			// prevent un-necessary GitHub API requests.
-			key := fmt.Sprintf("%d.%d", user.ID, repo.ID)
-			cache.Add(key, perm)
-
-			// debug
 			log.Debugf("%s granted %+v permission to %s",
 				user.Login, perm, repo.FullName)
 

router/middleware/token/token.go

@@ -1,4 +1,4 @@
-package refresh
+package token
 
 import (
 	"time"

router/router.go

@@ -7,9 +7,10 @@ import (
 	"github.com/gin-gonic/gin"
 
 	"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/refresh"
 	"github.com/drone/drone/router/middleware/session"
+	"github.com/drone/drone/router/middleware/token"
 	"github.com/drone/drone/static"
 	"github.com/drone/drone/template"
 )
@@ -19,10 +20,13 @@ func Load(middleware ...gin.HandlerFunc) http.Handler {
 	e.SetHTMLTemplate(template.Load())
 	e.StaticFS("/static", static.FileSystem())
 
-	e.Use(header.SetHeaders())
+	e.Use(header.NoCache)
+	e.Use(header.Options)
+	e.Use(header.Secure)
 	e.Use(middleware...)
 	e.Use(session.SetUser())
-	e.Use(refresh.Refresh)
+	e.Use(cache.Perms)
+	e.Use(token.Refresh)
 
 	e.GET("/", controller.ShowIndex)
 	e.GET("/login", controller.ShowLogin)
@@ -58,7 +62,7 @@ func Load(middleware ...gin.HandlerFunc) http.Handler {
 		user.GET("", controller.GetSelf)
 		user.GET("/builds", controller.GetFeed)
 		user.GET("/repos", controller.GetRepos)
-		user.GET("/repos/remote", controller.GetRemoteRepos)
+		user.GET("/repos/remote", cache.Repos, controller.GetRemoteRepos)
 		user.POST("/token", controller.PostToken)
 	}