Add Lambda handler wrapper for SNS-based custom resources

cfn/wrap.go

@@ -4,15 +4,22 @@ package cfn
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"log"
 	"net/http"
 
+	"github.com/aws/aws-lambda-go/events"
 	"github.com/aws/aws-lambda-go/lambdacontext"
 )
 
 // CustomResourceLambdaFunction is a standard form Lambda for a Custom Resource.
 type CustomResourceLambdaFunction func(context.Context, Event) (reason string, err error)
 
+// SNSCustomResourceLambdaFunction is a standard form Lambda for a Custom Resource
+// that is triggered via a SNS topic.
+type SNSCustomResourceLambdaFunction func(context.Context, events.SNSEvent) (reason string, err error)
+
 // CustomResourceFunction is a representation of the customer's Custom Resource function.
 // LambdaWrap will take the returned values and turn them into a response to be sent
 // to CloudFormation.
@@ -73,3 +80,24 @@ func lambdaWrapWithClient(lambdaFunction CustomResourceFunction, client httpClie
 func LambdaWrap(lambdaFunction CustomResourceFunction) (fn CustomResourceLambdaFunction) {
 	return lambdaWrapWithClient(lambdaFunction, http.DefaultClient)
 }
+
+// LambdaWrapSNS wraps a Lambda handler with support for SNS-based custom
+// resources. Usage and purpose otherwise same as LambdaWrap().
+func LambdaWrapSNS(lambdaFunction CustomResourceFunction) SNSCustomResourceLambdaFunction {
+	inner := LambdaWrap(lambdaFunction)
+	return func(ctx context.Context, event events.SNSEvent) (reason string, err error) {
+		if len(event.Records) != 1 {
+			err = errors.New("expected exactly 1 incoming record")
+			return
+		}
+
+		message := event.Records[0].SNS.Message
+
+		var innerEvent Event
+		if err = json.Unmarshal([]byte(message), &innerEvent); err != nil {
+			return
+		}
+
+		return inner(ctx, innerEvent)
+	}
+}