From fa4dc01cf4b74eeef2571f8fbb7683cf8925010c Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Wed, 22 Mar 2017 16:33:41 -0400
Subject: [PATCH] add named state support to the s3 backend

This adds named state (environment) support to the S3 backend.

A state NAME will prepend the configured s3 key with `env:/NAME/`.
The default state will remain rooted in the bucket for backwards
compatibility.

Locks in DynamoDB use the S3 key as the as the primary key value, so
locking will work as expected for multiple states.
---
 backend/remote-state/s3/backend_state.go | 77 ++++++++++++++++++------
 backend/remote-state/s3/backend_test.go  | 21 +++++--
 2 files changed, 75 insertions(+), 23 deletions(-)

diff --git a/backend/remote-state/s3/backend_state.go b/backend/remote-state/s3/backend_state.go
index 6e9252861..3166cbfb9 100644
--- a/backend/remote-state/s3/backend_state.go
+++ b/backend/remote-state/s3/backend_state.go
@@ -2,46 +2,81 @@ package s3
 
 import (
 	"fmt"
+	"sort"
 	"strings"
 
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/hashicorp/terraform/backend"
 	"github.com/hashicorp/terraform/state"
 	"github.com/hashicorp/terraform/state/remote"
 )
 
 const (
-	// This will be used a directory name, the odd looking colon is to reduce
-	// the chance of name conflicts with existing deployments.
+	// This will be used as directory name, the odd looking colon is simply to
+	// reduce the chance of name conflicts with existing objects.
 	keyEnvPrefix = "env:"
 )
 
 func (b *Backend) States() ([]string, error) {
-	return nil, backend.ErrNamedStatesNotSupported
+	params := &s3.ListObjectsInput{
+		Bucket: &b.bucketName,
+		Prefix: aws.String(keyEnvPrefix + "/"),
+	}
+
+	resp, err := b.s3Client.ListObjects(params)
+	if err != nil {
+		return nil, err
+	}
+
+	var envs []string
+	for _, obj := range resp.Contents {
+		env := keyEnv(*obj.Key)
+		if env != "" {
+			envs = append(envs, env)
+		}
+	}
+
+	sort.Strings(envs)
+	envs = append([]string{backend.DefaultStateName}, envs...)
+	return envs, nil
+}
+
+// extract the env name from the S3 key
+func keyEnv(key string) string {
+	parts := strings.Split(key, "/")
+	if len(parts) < 3 {
+		// no env here
+		return ""
+	}
+
+	if parts[0] != keyEnvPrefix {
+		// not our key, so ignore
+		return ""
+	}
+
+	return parts[1]
 }
 
 func (b *Backend) DeleteState(name string) error {
-	return backend.ErrNamedStatesNotSupported
 	if name == backend.DefaultStateName || name == "" {
 		return fmt.Errorf("can't delete default state")
 	}
 
-	//params := &s3.ListObjectsInput{
-	//    Bucket:       &b.client.bucketName,
-	//    Delimiter:    aws.String("Delimiter"),
-	//    EncodingType: aws.String("EncodingType"),
-	//    Marker:       aws.String("Marker"),
-	//    MaxKeys:      aws.Int64(1),
-	//    Prefix:       aws.String("env"),
-	//    RequestPayer: aws.String("RequestPayer"),
-	//}
+	params := &s3.DeleteObjectInput{
+		Bucket: &b.bucketName,
+		Key:    aws.String(b.path(name)),
+	}
+
+	_, err := b.s3Client.DeleteObject(params)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
 func (b *Backend) State(name string) (state.State, error) {
-	if name != backend.DefaultStateName {
-		return nil, backend.ErrNamedStatesNotSupported
-	}
-
 	client := &RemoteClient{
 		s3Client:             b.s3Client,
 		dynClient:            b.dynClient,
@@ -53,7 +88,13 @@ func (b *Backend) State(name string) (state.State, error) {
 		lockTable:            b.lockTable,
 	}
 
-	// TODO: create new state if it doesn't exist
+	// if this isn't the default state name, we need to create the object so
+	// it's listed by States.
+	if name != backend.DefaultStateName {
+		if err := client.Put([]byte{}); err != nil {
+			return nil, err
+		}
+	}
 
 	return &remote.State{Client: client}, nil
 }
diff --git a/backend/remote-state/s3/backend_test.go b/backend/remote-state/s3/backend_test.go
index 0838fa159..f8b664b80 100644
--- a/backend/remote-state/s3/backend_test.go
+++ b/backend/remote-state/s3/backend_test.go
@@ -127,13 +127,24 @@ func createS3Bucket(t *testing.T, s3Client *s3.S3, bucketName string) {
 }
 
 func deleteS3Bucket(t *testing.T, s3Client *s3.S3, bucketName string) {
-	deleteBucketReq := &s3.DeleteBucketInput{
-		Bucket: &bucketName,
+	warning := "WARNING: Failed to delete the test S3 bucket. It may have been left in your AWS account and may incur storage charges. (error was %s)"
+
+	// first we have to get rid of the env objects, or we can't delete the bucket
+	resp, err := s3Client.ListObjects(&s3.ListObjectsInput{Bucket: &bucketName})
+	if err != nil {
+		t.Logf(warning, err)
+		return
+	}
+	for _, obj := range resp.Contents {
+		if _, err := s3Client.DeleteObject(&s3.DeleteObjectInput{Bucket: &bucketName, Key: obj.Key}); err != nil {
+			// this will need cleanup no matter what, so just warn and exit
+			t.Logf(warning, err)
+			return
+		}
 	}
 
-	_, err := s3Client.DeleteBucket(deleteBucketReq)
-	if err != nil {
-		t.Logf("WARNING: Failed to delete the test S3 bucket. It may have been left in your AWS account and may incur storage charges. (error was %s)", err)
+	if _, err := s3Client.DeleteBucket(&s3.DeleteBucketInput{Bucket: &bucketName}); err != nil {
+		t.Logf(warning, err)
 	}
 }