Gracefully handle creation failures of custom resources

pkg/cloudformation/customresources/customresources.go

@@ -0,0 +1,252 @@
+// Package customresources provides a Go library for building CloudFormation
+// custom resource handlers.
+package customresources
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"hash/fnv"
+	"io/ioutil"
+	"net/http"
+	"time"
+
+	"github.com/remind101/empire/pkg/base62"
+	"golang.org/x/net/context"
+)
+
+// Possible request types.
+const (
+	Create = "Create"
+	Update = "Update"
+	Delete = "Delete"
+)
+
+// Possible response statuses.
+const (
+	StatusSuccess = "SUCCESS"
+	StatusFailed  = "FAILED"
+)
+
+// Provisioner is something that can provision custom resources.
+type Provisioner interface {
+	// Provision should do the appropriate provisioning, then return:
+	//
+	// 1. The physical id that was created, if any.
+	// 2. The data to return.
+	Provision(context.Context, Request) (string, interface{}, error)
+
+	// Properties should return an instance of a type that the properties
+	// can be json.Unmarshalled into.
+	Properties() interface{}
+}
+
+// Request represents a Custom Resource request.
+//
+// See http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/crpg-ref-requests.html
+type Request struct {
+	// The request type is set by the AWS CloudFormation stack operation
+	// (create-stack, update-stack, or delete-stack) that was initiated by
+	// the template developer for the stack that contains the custom
+	// resource.
+	//
+	// Must be one of: Create, Update, or Delete.
+	RequestType string `json:"RequestType"`
+
+	// The response URL identifies a pre-signed Amazon S3 bucket that
+	// receives responses from the custom resource provider to AWS
+	// CloudFormation.
+	ResponseURL string `json:"ResponseURL"`
+
+	// The Amazon Resource Name (ARN) that identifies the stack containing
+	// the custom resource.
+	//
+	// Combining the StackId with the RequestId forms a value that can be
+	// used to uniquely identify a request on a particular custom resource.
+	StackId string `json:"StackId"`
+
+	// A unique ID for the request.
+	//
+	// Combining the StackId with the RequestId forms a value that can be
+	// used to uniquely identify a request on a particular custom resource.
+	RequestId string `json:"RequestId"`
+
+	// The template developer-chosen resource type of the custom resource in
+	// the AWS CloudFormation template. Custom resource type names can be up
+	// to 60 characters long and can include alphanumeric and the following
+	// characters: _@-.
+	ResourceType string `json:"ResourceType"`
+
+	// The template developer-chosen name (logical ID) of the custom
+	// resource in the AWS CloudFormation template. This is provided to
+	// facilitate communication between the custom resource provider and the
+	// template developer.
+	LogicalResourceId string `json:"LogicalResourceId"`
+
+	// A required custom resource provider-defined physical ID that is
+	// unique for that provider.
+	//
+	// Always sent with Update and Delete requests; never sent with Create.
+	PhysicalResourceId string `json:"PhysicalResourceId"`
+
+	// This field contains the contents of the Properties object sent by the
+	// template developer. Its contents are defined by the custom resource
+	// provider.
+	ResourceProperties interface{} `json:"ResourceProperties"`
+
+	// Used only for Update requests. Contains the resource properties that
+	// were declared previous to the update request.
+	OldResourceProperties interface{} `json:"OldResourceProperties"`
+}
+
+// Hash returns a compact unique identifier for the request.
+func (r *Request) Hash() string {
+	h := fnv.New64()
+	h.Write([]byte(fmt.Sprintf("%s.%s", r.StackId, r.RequestId)))
+	return base62.Encode(h.Sum64())
+}
+
+// Response represents the response body we send back to CloudFormation when
+// provisioning is complete.
+//
+// See http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/crpg-ref-responses.html
+type Response struct {
+	// The status value sent by the custom resource provider in response to
+	// an AWS CloudFormation-generated request.
+	//
+	// Must be either SUCCESS or FAILED.
+	Status string `json:"Status"`
+
+	// Describes the reason for a failure response.
+	//
+	// Required if Status is FAILED; optional otherwise.
+	Reason string `json:"Reason,omitempty"`
+
+	// This value should be an identifier unique to the custom resource
+	// vendor, and can be up to 1Kb in size. The value must be a non-empty
+	// string.
+	PhysicalResourceId string `json:"PhysicalResourceId"`
+
+	// The Amazon Resource Name (ARN) that identifies the stack containing
+	// the custom resource. This response value should be copied verbatim
+	// from the request.
+	StackId string `json:"StackId"`
+
+	// A unique ID for the request. This response value should be copied
+	// verbatim from the request.
+	RequestId string `json:"RequestId"`
+
+	// The template developer-chosen name (logical ID) of the custom
+	// resource in the AWS CloudFormation template. This response value
+	// should be copied verbatim from the request.
+	LogicalResourceId string `json:"LogicalResourceId"`
+
+	// Optional, custom resource provider-defined name-value pairs to send
+	// with the response. The values provided here can be accessed by name
+	// in the template with Fn::GetAtt.
+	Data interface{} `json:"Data,omitempty"`
+}
+
+// NewResponseFromRequest initializes a new Response from a Request, filling in
+// the required verbatim fields.
+func NewResponseFromRequest(req Request) Response {
+	return Response{
+		StackId:           req.StackId,
+		RequestId:         req.RequestId,
+		LogicalResourceId: req.LogicalResourceId,
+	}
+}
+
+// SendResponse sends the response the Response to the requests response url
+func SendResponse(req Request, response Response) error {
+	return SendResponseWithClient(http.DefaultClient, req, response)
+}
+
+// SendResponseWithClient uploads the response to the requests signed
+// ResponseURL.
+func SendResponseWithClient(client interface {
+	Do(*http.Request) (*http.Response, error)
+}, req Request, response Response) error {
+	c := ResponseClient{client}
+	return c.SendResponse(req, response)
+}
+
+// ResponseClient is a client that can send responses to a requests ResponseURL.
+type ResponseClient struct {
+	client interface {
+		Do(*http.Request) (*http.Response, error)
+	}
+}
+
+// SendResponse sends the response to the request's ResponseURL.
+func (c *ResponseClient) SendResponse(req Request, response Response) error {
+	raw, err := json.Marshal(response)
+	if err != nil {
+		return err
+	}
+
+	r, err := http.NewRequest("PUT", req.ResponseURL, bytes.NewReader(raw))
+	if err != nil {
+		return err
+	}
+
+	resp, err := c.client.Do(r)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	body, _ := ioutil.ReadAll(resp.Body)
+
+	if code := resp.StatusCode; code/100 != 2 {
+		return fmt.Errorf("unexpected response from pre-signed url: %v: %v", code, string(body))
+	}
+
+	return nil
+}
+
+// WithTimeout wraps a Provisioner with a context.WithTimeout.
+func WithTimeout(p Provisioner, timeout time.Duration, grace time.Duration) Provisioner {
+	return &timeoutProvisioner{
+		Provisioner: p,
+		timeout:     timeout,
+		grace:       grace,
+	}
+}
+
+type result struct {
+	id   string
+	data interface{}
+	err  error
+}
+
+type timeoutProvisioner struct {
+	Provisioner
+	timeout time.Duration
+	grace   time.Duration
+}
+
+func (p *timeoutProvisioner) Provision(ctx context.Context, r Request) (string, interface{}, error) {
+	ctx, cancel := context.WithTimeout(ctx, p.timeout)
+	defer cancel()
+
+	done := make(chan result)
+	go func() {
+		id, data, err := p.Provisioner.Provision(ctx, r)
+		done <- result{id, data, err}
+	}()
+
+	select {
+	case r := <-done:
+		return r.id, r.data, r.err
+	case <-ctx.Done():
+		// When the context is canceled, give the provisioner
+		// some extra time to cleanup.
+		<-time.After(p.grace)
+		select {
+		case r := <-done:
+			return r.id, r.data, r.err
+		default:
+			return "", nil, ctx.Err()
+		}
+	}
+}

pkg/cloudformation/customresources/customresources_test.go

@@ -0,0 +1,138 @@
+package customresources
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"golang.org/x/net/context"
+)
+
+var ctx = context.Background()
+
+func TestWithTimeout_NoTimeout(t *testing.T) {
+	m := new(mockProvisioner)
+	p := WithTimeout(m, time.Second, time.Second)
+
+	m.On("Provision", Request{}).Return("id", nil, nil)
+
+	p.Provision(ctx, Request{})
+}
+
+func TestWithTimeout_Timeout_Cleanup(t *testing.T) {
+	m := new(mockProvisioner)
+	p := WithTimeout(m, time.Millisecond*500, time.Millisecond*500)
+
+	m.On("Provision", Request{}).Return("id", nil, nil).Run(func(mock.Arguments) {
+		time.Sleep(time.Millisecond * 750)
+	})
+
+	id, _, err := p.Provision(ctx, Request{})
+	assert.NoError(t, err)
+	assert.Equal(t, "id", id)
+}
+
+func TestWithTimeout_GraceTimeout(t *testing.T) {
+	m := new(mockProvisioner)
+	p := WithTimeout(m, time.Millisecond*500, time.Millisecond*500)
+
+	m.On("Provision", Request{}).Return("id", nil, nil).Run(func(mock.Arguments) {
+		time.Sleep(time.Millisecond * 1500)
+	})
+
+	_, _, err := p.Provision(ctx, Request{})
+	assert.Equal(t, context.DeadlineExceeded, err)
+}
+
+func TestResponse_MarshalJSON(t *testing.T) {
+	tests := map[Response]string{
+		Response{
+			Status: StatusSuccess,
+		}: `{"Status":"SUCCESS","PhysicalResourceId":"","StackId":"","RequestId":"","LogicalResourceId":""}`,
+
+		Response{
+			Status: StatusFailed,
+			Reason: "errored",
+		}: `{"Status":"FAILED","Reason":"errored","PhysicalResourceId":"","StackId":"","RequestId":"","LogicalResourceId":""}`,
+	}
+
+	for resp, expected := range tests {
+		raw, err := json.Marshal(resp)
+		assert.NoError(t, err)
+
+		assert.Equal(t, expected, string(raw))
+	}
+}
+
+func TestResponseClient(t *testing.T) {
+	c := ResponseClient{http.DefaultClient}
+
+	response := Response{
+		Status:             StatusSuccess,
+		PhysicalResourceId: "9001",
+		StackId:            "arn:aws:cloudformation:us-east-1:066251891493:stack/foo/70213b00-0e74-11e6-b4fb-500c28680ac6",
+		RequestId:          "daf3f3f9-79a1-4049-823e-09544e582b06",
+		LogicalResourceId:  "webInstancePort",
+		Data:               map[string]int64{"InstancePort": 9001},
+	}
+
+	var called bool
+	check := func(w http.ResponseWriter, r *http.Request) {
+		called = true
+
+		assert.Equal(t, "PUT", r.Method)
+		assert.Equal(t, "/arn:aws:cloudformation:us-east-1:066251891493:stack/foo/70213b00-0e74-11e6-b4fb-500c28680ac6|webInstancePort|daf3f3f9-79a1-4049-823e-09544e582b06", r.URL.Path)
+
+		expectedRaw, err := json.Marshal(response)
+		assert.NoError(t, err)
+		raw, err := ioutil.ReadAll(r.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, string(expectedRaw), string(raw))
+	}
+	s := httptest.NewServer(http.HandlerFunc(check))
+	defer s.Close()
+
+	req := Request{
+		RequestType:       Create,
+		ResponseURL:       s.URL + "/arn%3Aaws%3Acloudformation%3Aus-east-1%3A066251891493%3Astack/foo/70213b00-0e74-11e6-b4fb-500c28680ac6%7CwebInstancePort%7Cdaf3f3f9-79a1-4049-823e-09544e582b06?AWSAccessKeyId=AKIAJNXHFR7P7YGKLDPQ&Expires=1461987599&Signature=EqV%2BqIUAsZPz5Q%2F%2B75Guvn%2BNREU%3D",
+		StackId:           "arn:aws:cloudformation:us-east-1:066251891493:stack/foo/70213b00-0e74-11e6-b4fb-500c28680ac6",
+		RequestId:         "daf3f3f9-79a1-4049-823e-09544e582b06",
+		LogicalResourceId: "webInstancePort",
+		ResourceType:      "Custom::InstancePort",
+		ResourceProperties: map[string]interface{}{
+			"ServiceToken": "arn:aws:sns:us-east-1:066251891493:empire-e01a8fac-CustomResourcesTopic-9KHPNW7WFKBD",
+		},
+	}
+
+	err := c.SendResponse(req, response)
+	assert.NoError(t, err)
+
+	assert.True(t, called)
+}
+
+type mockProvisioner struct {
+	mock.Mock
+}
+
+func (m *mockProvisioner) Provision(_ context.Context, req Request) (string, interface{}, error) {
+	args := m.Called(req)
+	return args.String(0), args.Get(1), args.Error(2)
+}
+
+func (m *mockProvisioner) Properties() interface{} {
+	return nil
+}
+
+type mockHTTPClient struct {
+	mock.Mock
+}
+
+func (m *mockHTTPClient) Do(req *http.Request) (*http.Response, error) {
+	args := m.Called(req)
+	return args.Get(0).(*http.Response), args.Error(1)
+}