diff --git a/pkg/imds/middleware/cache.go b/pkg/imds/middleware/cache.go
new file mode 100644
index 0000000..a813672
--- /dev/null
+++ b/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()
+ }
+}
diff --git a/pkg/imds/middleware/cache_test.go b/pkg/imds/middleware/cache_test.go
new file mode 100644
index 0000000..a2b3844
--- /dev/null
+++ b/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())
+}
diff --git a/pkg/imds/middleware/json.go b/pkg/imds/middleware/json.go
index 41bb213..de9cae5 100644
--- a/pkg/imds/middleware/json.go
+++ b/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 {
diff --git a/pkg/imds/middleware/writer.go b/pkg/imds/middleware/writer.go
deleted file mode 100644
index 72862ca..0000000
--- a/pkg/imds/middleware/writer.go
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
-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 (
- "bufio"
- "net"
- "net/http"
-
- "github.com/gin-gonic/gin"
-)
-
-// JSONFormatter defines an interface for formatting JSON text
-type JSONFormatter interface {
- Format(in []byte) []byte
- FormatString(in string) string
-}
-
-// A crude wrapper around a gin.ResponseWriter for supporting the custom
-// injection of a JSON formatter
-type jsonRewriter struct {
- ResponseWriter gin.ResponseWriter
- Formatter JSONFormatter
-}
-
-func (w *jsonRewriter) Header() http.Header { return w.ResponseWriter.Header() }
-func (w *jsonRewriter) WriteHeader(code int) { w.ResponseWriter.WriteHeader(code) }
-func (w *jsonRewriter) WriteHeaderNow() { w.ResponseWriter.WriteHeaderNow() }
-
-func (w *jsonRewriter) Write(data []byte) (n int, err error) {
- return w.ResponseWriter.Write(w.Formatter.Format(data))
-}
-
-func (w *jsonRewriter) WriteString(s string) (n int, err error) {
- return w.ResponseWriter.WriteString(w.Formatter.FormatString(s))
-}
-
-func (w *jsonRewriter) Status() int { return w.ResponseWriter.Status() }
-func (w *jsonRewriter) Size() int { return w.ResponseWriter.Size() }
-func (w *jsonRewriter) Written() bool { return w.ResponseWriter.Written() }
-
-func (w *jsonRewriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
- return w.ResponseWriter.Hijack()
-}
-
-func (w *jsonRewriter) CloseNotify() <-chan bool { return w.ResponseWriter.CloseNotify() }
-func (w *jsonRewriter) Flush() { w.ResponseWriter.Flush() }
-func (w *jsonRewriter) Pusher() (pusher http.Pusher) { return w.ResponseWriter.Pusher() }
diff --git a/pkg/imds/server.go b/pkg/imds/server.go
index 8a88bff..6e1a991 100644
--- a/pkg/imds/server.go
+++ b/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
diff --git a/pkg/imds/server_test.go b/pkg/imds/server_test.go
index a4aee2d..261c8b1 100644
--- a/pkg/imds/server_test.go
+++ b/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, `
+
+
+
+ 400 - Bad Request
+
+
+ 400 - Bad Request
+
+`, 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, `
+
+
+
+ 400 - Bad Request
+
+
+ 400 - Bad Request
+
+`, w.Body.String())
+ })
+ }
+}
diff --git a/pkg/imds/token/token.go b/pkg/imds/token/v2.go
similarity index 100%
rename from pkg/imds/token/token.go
rename to pkg/imds/token/v2.go
diff --git a/pkg/imds/token/token_test.go b/pkg/imds/token/v2_test.go
similarity index 100%
rename from pkg/imds/token/token_test.go
rename to pkg/imds/token/v2_test.go