From f43ccf4da7b99863e52c52e1f05b007111f09f62 Mon Sep 17 00:00:00 2001 From: Purple Clay <106762954+purpleclay@users.noreply.github.com> Date: Sun, 21 Aug 2022 14:28:31 +0100 Subject: [PATCH] feat: basic spot instance termination simulation (#24) --- cmd/root.go | 1 + pkg/imds/patch/func.go | 6 +- pkg/imds/patch/instancetag_test.go | 6 +- pkg/imds/patch/spot.go | 115 +++++++++++++++++++++++++++++ pkg/imds/patch/spot_test.go | 110 +++++++++++++++++++++++++++ pkg/imds/server.go | 18 ++++- pkg/imds/server_test.go | 24 ++++++ 7 files changed, 272 insertions(+), 8 deletions(-) create mode 100644 pkg/imds/patch/spot.go create mode 100644 pkg/imds/patch/spot_test.go diff --git a/cmd/root.go b/cmd/root.go index 522ba7b..74a6fe3 100644 --- a/cmd/root.go +++ b/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)) diff --git a/pkg/imds/patch/func.go b/pkg/imds/patch/func.go index a5f908a..507cd46 100644 --- a/pkg/imds/patch/func.go +++ b/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 { diff --git a/pkg/imds/patch/instancetag_test.go b/pkg/imds/patch/instancetag_test.go index 0e35edf..b767856 100644 --- a/pkg/imds/patch/instancetag_test.go +++ b/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", diff --git a/pkg/imds/patch/spot.go b/pkg/imds/patch/spot.go new file mode 100644 index 0000000..4e9e9be --- /dev/null +++ b/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 +} diff --git a/pkg/imds/patch/spot_test.go b/pkg/imds/patch/spot_test.go new file mode 100644 index 0000000..f7eb36d --- /dev/null +++ b/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) +} diff --git a/pkg/imds/server.go b/pkg/imds/server.go index 6e1a991..69f4553 100644 --- a/pkg/imds/server.go +++ b/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) diff --git a/pkg/imds/server_test.go b/pkg/imds/server_test.go index 261c8b1..d556bb2 100644 --- a/pkg/imds/server_test.go +++ b/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) +}