feat: basic spot instance termination simulation (#24)

cmd/root.go

@@ -49,6 +49,7 @@ func Execute(out io.Writer) error {
 	flags.StringToStringVar(&opts.InstanceTags, "instance-tags", imds.DefaultOptions.InstanceTags, "a list of instance tags (key pairs) to expose as metadata")
 	flags.IntVar(&opts.Port, "port", imds.DefaultOptions.Port, "the port to be used at startup")
 	flags.BoolVar(&opts.Pretty, "pretty", imds.DefaultOptions.Pretty, "if instance categories should return pretty printed JSON")
+	flags.BoolVar(&opts.Spot, "spot", imds.DefaultOptions.Spot, "enable simulation of a spot instance and interruption notice")
 
 	rootCmd.AddCommand(newVersionCmd(out))
 	rootCmd.AddCommand(newManPagesCmd(out))

pkg/imds/patch/func.go

@@ -34,10 +34,10 @@ import (
 //		"FirstName": "joe",
 //		"LastName": "bloggs",
 //	}
-// 	out := JSONPairs(pairs)
-// 	fmt.Println(out)
+//	out := JSONPairs(pairs)
+//	fmt.Println(out)
 //
-//  // => "\"FirstName\": \"joe\", \"LastName\": \"bloggs\""
+//	"\"FirstName\": \"joe\", \"LastName\": \"bloggs\""
 func JSONPairs(in map[string]string) string {
 	pairs := make([]string, 0, len(in))
 	for k, v := range in {

pkg/imds/patch/instancetag_test.go

@@ -31,7 +31,7 @@ import (
 	"github.com/tidwall/pretty"
 )
 
-func TestPatch(t *testing.T) {
+func TestInstanceTagPatch(t *testing.T) {
 	tagPatch := patch.InstanceTag{
 		Tags: map[string]string{
 			"Name":        "testing",
@@ -57,7 +57,7 @@ func TestPatch(t *testing.T) {
 `, string(pretty.PrettyOptions(out, opts)))
 }
 
-func TestPatchNoTags(t *testing.T) {
+func TestInstanceTagPatch_NoTags(t *testing.T) {
 	tagPatch := patch.InstanceTag{
 		Tags: map[string]string{},
 	}
@@ -68,7 +68,7 @@ func TestPatchNoTags(t *testing.T) {
 	assert.Equal(t, `{"testing":"123"}`, string(out))
 }
 
-func TestPathInvalidInputJSON(t *testing.T) {
+func TestInstanceTagPatch_InvalidInputJSON(t *testing.T) {
 	tagPatch := patch.InstanceTag{
 		Tags: map[string]string{
 			"Name":        "testing",

pkg/imds/patch/spot.go

@@ -0,0 +1,115 @@
+/*
+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 patch
+
+import (
+	"bytes"
+	"text/template"
+	"time"
+
+	jsonpatch "github.com/evanphx/json-patch/v5"
+)
+
+const spotTemplate = `[
+	{
+		"op": "replace",
+		"path": "/instance-life-cycle",
+		"value": "spot"
+	},
+	{
+		"op": "add", 
+		"path": "/spot", 
+		"value": {
+			"instance-action": {
+				"action": "{{ .Action }}", 
+				"time": "{{ .ActionTime }}"
+			}{{ if .TerminationTime }},
+			"termination-time": "{{ .TerminationTime }}"{{ end }}
+		}
+	},
+	{
+		"op": "add",
+		"path": "/events",
+		"value": {
+			"recommendations": {
+				"rebalance": {
+					"noticeTime": "{{ .RebalanceTime }}"
+				}
+			}
+		}
+	}
+]`
+
+var spotPatch = template.Must(template.New("SpotPatch").
+	Parse(spotTemplate))
+
+// SpotInstanceAction is used to represent the lifecycle event (or action) of a spot instance
+type SpotInstanceAction string
+
+const (
+	TerminateSpotInstanceAction SpotInstanceAction = "terminate"
+	StopSpotInstanceAction      SpotInstanceAction = "stop"
+	HibernateSpotInstanceAction SpotInstanceAction = "hibernate"
+)
+
+// Spot is used to patch a JSON document and replicate the use and lifecycle of a spot EC2 instance.
+// The lifecycle of a spot instance is exposed through the use of an instance action category,
+// see: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/spot-instance-termination-notices.html#instance-action-metadata
+type Spot struct {
+	InstanceAction SpotInstanceAction
+}
+
+// Patch the document based on the provided instance action. The resulting JSON
+// document will conform to the IMDS specification and expose spot details within
+// the IMDS metadata
+func (p Spot) Patch(in []byte) ([]byte, error) {
+	now := time.Now().UTC().Format(time.RFC3339)
+
+	spotDetails := struct {
+		Action          SpotInstanceAction
+		ActionTime      string
+		TerminationTime string
+		RebalanceTime   string
+	}{
+		Action:        p.InstanceAction,
+		ActionTime:    now,
+		RebalanceTime: now,
+	}
+
+	if spotDetails.Action == TerminateSpotInstanceAction {
+		// For backwards compatibility set the termination time
+		spotDetails.TerminationTime = now
+	}
+
+	var buf bytes.Buffer
+	spotPatch.Execute(&buf, spotDetails)
+
+	patch, _ := jsonpatch.DecodePatch(buf.Bytes())
+
+	out, err := patch.Apply(in)
+	if err != nil {
+		return in, err
+	}
+
+	return out, nil
+}

pkg/imds/patch/spot_test.go

@@ -0,0 +1,110 @@
+/*
+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 patch_test
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/purpleclay/imds-mock/pkg/imds/patch"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// Only used for serialising the patch into a struct for assertions
+type spotPatchJSON struct {
+	InstanceLifeCycle string `json:"instance-life-cycle"`
+	Spot              struct {
+		InstanceAction struct {
+			Action     string    `json:"action"`
+			ActionTime time.Time `json:"time"`
+		} `json:"instance-action"`
+
+		TerminationTime *time.Time `json:"termination-time"`
+	} `json:"spot"`
+	Events struct {
+		Recommendations struct {
+			Rebalance struct {
+				NoticeTime time.Time `json:"noticeTime"`
+			} `json:"rebalance"`
+		} `json:"recommendations"`
+	} `json:"events"`
+}
+
+func TestSpotPatch_TerminateAction(t *testing.T) {
+	tests := []struct {
+		name   string
+		action patch.SpotInstanceAction
+	}{
+		{
+			name:   "TerminateAction",
+			action: patch.TerminateSpotInstanceAction,
+		},
+		{
+			name:   "StopAction",
+			action: patch.StopSpotInstanceAction,
+		},
+		{
+			name:   "HibernateAction",
+			action: patch.HibernateSpotInstanceAction,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			spotPatch := patch.Spot{
+				InstanceAction: tt.action,
+			}
+
+			out, err := spotPatch.Patch([]byte(`{"instance-life-cycle":"on-demand"}`))
+			require.NoError(t, err)
+
+			var spotJSON spotPatchJSON
+			json.Unmarshal(out, &spotJSON)
+
+			now := time.Now().UTC()
+
+			assert.Equal(t, "spot", spotJSON.InstanceLifeCycle)
+			assert.Equal(t, string(tt.action), spotJSON.Spot.InstanceAction.Action)
+			assert.WithinDuration(t, now, spotJSON.Spot.InstanceAction.ActionTime, 1*time.Second)
+
+			if tt.action == patch.TerminateSpotInstanceAction {
+				require.NotNil(t, spotJSON.Spot.TerminationTime)
+				assert.WithinDuration(t, now, *spotJSON.Spot.TerminationTime, 1*time.Second)
+			} else {
+				require.Nil(t, spotJSON.Spot.TerminationTime)
+			}
+
+			assert.WithinDuration(t, now, spotJSON.Events.Recommendations.Rebalance.NoticeTime, 1*time.Second)
+		})
+	}
+}
+
+func TestSpotPatch_InvalidInputJSON(t *testing.T) {
+	spotPatch := patch.Spot{
+		InstanceAction: patch.HibernateSpotInstanceAction,
+	}
+
+	_, err := spotPatch.Patch([]byte(`{`))
+	require.Error(t, err)
+}

pkg/imds/server.go

@@ -97,6 +97,11 @@ type Options struct {
 	// Pretty controls if the JSON outputted by any instance category
 	// is pretty printed. By default all JSON will be compacted
 	Pretty bool
+
+	// Spot enables the simulation of a spot instance and interruption notice
+	// through the IMDS mock. By default this will set to false and an on-demand
+	// instance will be simulated
+	Spot bool
 }
 
 // DefaultOptions defines the default set of options that will be applied
@@ -110,13 +115,16 @@ var DefaultOptions = Options{
 	},
 	Port:   1338,
 	Pretty: false,
+	Spot:   false,
 }
 
 // Used as a hashset for quick lookups. Any matched path will just return its value
 // and not be used to perform a key lookup
 var reservedPaths = map[string]struct{}{
-	"iam.info":                 {},
-	"iam.security-credentials": {},
+	"iam.info":                         {},
+	"iam.security-credentials":         {},
+	"spot.instance-action":             {},
+	"events.recommendations.rebalance": {},
 }
 
 // Serve configures the IMDS mock using default options to handle HTTP requests
@@ -255,6 +263,12 @@ func patchResponseJSON(in []byte, opts Options) ([]byte, error) {
 		})
 	}
 
+	if opts.Spot {
+		patches = append(patches, patch.Spot{
+			InstanceAction: patch.TerminateSpotInstanceAction,
+		})
+	}
+
 	var err error
 	for _, p := range patches {
 		in, err = p.Patch(in)

pkg/imds/server_test.go

@@ -41,6 +41,7 @@ var testOptions = imds.Options{
 	ExcludeInstanceTags: imds.DefaultOptions.ExcludeInstanceTags,
 	InstanceTags:        imds.DefaultOptions.InstanceTags,
 	Pretty:              imds.DefaultOptions.Pretty,
+	Spot:                imds.DefaultOptions.Spot,
 }
 
 func TestMain(m *testing.M) {
@@ -332,3 +333,26 @@ func TestAPIToken_BadTTL(t *testing.T) {
 		})
 	}
 }
+
+func TestNoSpotCategoriesByDefault(t *testing.T) {
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/latest/meta-data/spot/instance-action", http.NoBody)
+
+	r, _ := imds.ServeWith(testOptions)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusNotFound, w.Code)
+}
+
+func TestSpotSimulation(t *testing.T) {
+	opts := testOptions
+	opts.Spot = true
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/latest/meta-data/spot/instance-action", http.NoBody)
+
+	r, _ := imds.ServeWith(opts)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+}