feat: use an in-memory cache to cache responses by category paths (#17)

pkg/imds/middleware/cache.go

@@ -0,0 +1,64 @@
+/*
+Copyright (c) 2022 Purple Clay
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+package middleware
+
+import (
+	"bytes"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+)
+
+// Capture any write to the response with an in memory buffer
+type cacheWriter struct {
+	gin.ResponseWriter
+	body bytes.Buffer
+}
+
+func (w *cacheWriter) Write(data []byte) (n int, err error) {
+	w.body.Write(data)
+	return w.ResponseWriter.Write(data)
+}
+
+// Cache provides middleware caching any request to the IMDS mock using
+// an in memory map. The IMDS response is cached using a lookup query based
+// on the IMDS instance category path
+func Cache() gin.HandlerFunc {
+	cache := map[string]string{}
+
+	return func(c *gin.Context) {
+		if res, hit := cache[c.Request.URL.Path]; hit {
+			c.Writer.Header().Add("Content-Type", "text/plain")
+			c.String(http.StatusOK, res)
+
+			c.Abort()
+			return
+		}
+
+		c.Writer = &cacheWriter{ResponseWriter: c.Writer}
+		c.Next()
+
+		// Grab the contents of the cache
+		cache[c.Request.URL.Path] = c.Writer.(*cacheWriter).body.String()
+	}
+}

pkg/imds/middleware/cache_test.go

@@ -0,0 +1,70 @@
+/*
+Copyright (c) 2022 Purple Clay
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+package middleware_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/purpleclay/imds-mock/pkg/imds/middleware"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCache_Miss(t *testing.T) {
+	count := 0
+
+	r := gin.Default()
+	r.GET("/cache", middleware.Cache(), func(c *gin.Context) {
+		count++
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/cache", http.NoBody)
+
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, 1, count)
+}
+
+func TestCache_Hit(t *testing.T) {
+	count := 0
+
+	r := gin.Default()
+	r.GET("/cache", middleware.Cache(), func(c *gin.Context) {
+		count++
+		c.String(http.StatusOK, "ok")
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/cache", http.NoBody)
+
+	r.ServeHTTP(w, req)
+	// Its important to reset the buffer here
+	w.Body.Reset()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, 1, count)
+	assert.Equal(t, "ok", w.Body.String())
+}

pkg/imds/middleware/json.go

@@ -27,6 +27,22 @@ import (
 	"github.com/tidwall/pretty"
 )
 
+// JSONFormatter defines an interface for formatting JSON text
+type JSONFormatter interface {
+	Format(in []byte) []byte
+}
+
+// A crude wrapper around a gin.ResponseWriter for supporting the custom
+// injection of a JSON formatter
+type jsonRewriter struct {
+	gin.ResponseWriter
+	Formatter JSONFormatter
+}
+
+func (w *jsonRewriter) Write(data []byte) (n int, err error) {
+	return w.ResponseWriter.Write(w.Formatter.Format(data))
+}
+
 type compactJSONFormatter struct{}
 
 func (f compactJSONFormatter) Format(in []byte) []byte {
@@ -41,10 +57,6 @@ func (f compactJSONFormatter) Format(in []byte) []byte {
 	return pretty.Ugly(in)
 }
 
-func (f compactJSONFormatter) FormatString(in string) string {
-	return string(f.Format([]byte(in)))
-}
-
 // CompactJSON provides middleware capable of rewriting all JSON responses
 // into a compact format
 func CompactJSON() gin.HandlerFunc {
@@ -79,10 +91,6 @@ func (f prettyJSONFormatter) Format(in []byte) []byte {
 	return pretty.PrettyOptions(in, prettyOptions)
 }
 
-func (f prettyJSONFormatter) FormatString(in string) string {
-	return string(f.Format([]byte(in)))
-}
-
 // PrettyJSON provides middleware capable of rewriting all JSON responses
 // into a pretty format
 func PrettyJSON() gin.HandlerFunc {

pkg/imds/server.go

@@ -141,11 +141,11 @@ func ServeWith(opts Options) (*gin.Engine, error) {
 		return nil, err
 	}
 
-	r.GET("/latest/meta-data", func(c *gin.Context) {
+	r.GET("/latest/meta-data", middleware.Cache(), func(c *gin.Context) {
 		c.String(http.StatusOK, keys(mockResponse, ""))
 	})
 
-	r.GET("/latest/meta-data/*category", func(c *gin.Context) {
+	r.GET("/latest/meta-data/*category", middleware.Cache(), func(c *gin.Context) {
 		categoryPath := c.Param("category")
 		if categoryPath == "/" {
 			// Exact same behaviour as /latest/meta-data

pkg/imds/server_test.go

@@ -30,6 +30,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/purpleclay/imds-mock/pkg/imds"
+	"github.com/purpleclay/imds-mock/pkg/imds/token"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tidwall/pretty"
@@ -260,3 +261,74 @@ func TestExcludeInstanceTags(t *testing.T) {
 
 	assert.Equal(t, http.StatusNotFound, w.Code)
 }
+
+func TestAPIToken(t *testing.T) {
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodPut, "/latest/api/token", http.NoBody)
+	req.Header.Add(imds.V2TokenTTLHeader, "10")
+
+	r, _ := imds.ServeWith(testOptions)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	assert.NotEmpty(t, w.Body.String())
+}
+
+func TestAPIToken_MissingHeader(t *testing.T) {
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodPut, "/latest/api/token", http.NoBody)
+
+	r, _ := imds.ServeWith(testOptions)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusBadRequest, w.Code)
+	assert.Equal(t, `<?xml version="1.0" encoding="iso-8859-1"?>
+<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
+	"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
+<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="en" lang="en">
+ <head>
+  <title>400 - Bad Request</title>
+ </head>
+ <body>
+  <h1>400 - Bad Request</h1>
+ </body>
+</html>`, w.Body.String())
+}
+
+func TestAPIToken_BadTTL(t *testing.T) {
+	tests := []struct {
+		name string
+		ttl  int
+	}{
+		{
+			name: "LessThan1",
+			ttl:  0,
+		},
+		{
+			name: "GreaterThanMax",
+			ttl:  token.MaxTTLInSeconds,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := httptest.NewRecorder()
+			req, _ := http.NewRequest(http.MethodPut, "/latest/api/token", http.NoBody)
+
+			r, _ := imds.ServeWith(testOptions)
+			r.ServeHTTP(w, req)
+
+			assert.Equal(t, http.StatusBadRequest, w.Code)
+			assert.Equal(t, `<?xml version="1.0" encoding="iso-8859-1"?>
+<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
+	"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
+<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="en" lang="en">
+ <head>
+  <title>400 - Bad Request</title>
+ </head>
+ <body>
+  <h1>400 - Bad Request</h1>
+ </body>
+</html>`, w.Body.String())
+		})
+	}
+}