Adds support for AMI sharing to Orgs and OUs

internal/service/ec2/ami_launch_permission.go

@@ -9,6 +9,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
 	"github.com/hashicorp/terraform-provider-aws/internal/conns"
 )
 
@@ -19,15 +20,36 @@ func ResourceAMILaunchPermission() *schema.Resource {
 		Delete: resourceAMILaunchPermissionDelete,
 		Importer: &schema.ResourceImporter{
 			State: func(d *schema.ResourceData, meta interface{}) ([]*schema.ResourceData, error) {
-				idParts := strings.Split(d.Id(), "/")
-				if len(idParts) != 2 || idParts[0] == "" || idParts[1] == "" {
-					return nil, fmt.Errorf("Unexpected format of ID (%q), expected ACCOUNT-ID/IMAGE-ID", d.Id())
+				if strings.HasPrefix(d.Id(), "arn") {
+					idParts := strings.Split(d.Id(), ",")
+					if len(idParts) != 3 || idParts[0] == "" || idParts[1] == "" || idParts[2] == "" {
+						return nil, fmt.Errorf("Unexpected format of ID (%q), expected ARN,ARN_TYPE,IMAGE-ID", d.Id())
+					}
+
+					arn := idParts[0]
+					arn_type := idParts[1]
+					imageId := idParts[2]
+					d.Set("arn", arn)
+					d.Set("arn_type", arn_type)
+					d.Set("image_id", imageId)
+					d.SetId(fmt.Sprintf("%s-%s", imageId, arn))
+				} else {
+					sep := ","
+					if !strings.Contains(d.Id(), sep) {
+						sep = "/"
+					}
+
+					idParts := strings.Split(d.Id(), sep)
+					if len(idParts) != 2 || idParts[0] == "" || idParts[1] == "" {
+						return nil, fmt.Errorf("Unexpected format of ID (%q), expected ACCOUNT-ID,IMAGE-ID", d.Id())
+					}
+
+					accountId := idParts[0]
+					imageId := idParts[1]
+					d.Set("account_id", accountId)
+					d.Set("image_id", imageId)
+					d.SetId(fmt.Sprintf("%s-%s", imageId, accountId))
 				}
-				accountId := idParts[0]
-				imageId := idParts[1]
-				d.Set("account_id", accountId)
-				d.Set("image_id", imageId)
-				d.SetId(fmt.Sprintf("%s-%s", imageId, accountId))
 				return []*schema.ResourceData{d}, nil
 			},
 		},
@@ -39,9 +61,26 @@ func ResourceAMILaunchPermission() *schema.Resource {
 				ForceNew: true,
 			},
 			"account_id": {
-				Type:     schema.TypeString,
-				Required: true,
-				ForceNew: true,
+				Type:         schema.TypeString,
+				Optional:     true,
+				ForceNew:     true,
+				ExactlyOneOf: []string{"account_id", "arn"},
+			},
+			"arn": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ForceNew:     true,
+				ExactlyOneOf: []string{"account_id", "arn"},
+			},
+			"arn_type": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ForceNew:     true,
+				RequiredWith: []string{"arn"},
+				ValidateFunc: validation.StringInSlice([]string{
+					"OrganizationArn",
+					"OrganizationalUnitArn",
+				}, false),
 			},
 		},
 	}
@@ -51,29 +90,43 @@ func resourceAMILaunchPermissionCreate(d *schema.ResourceData, meta interface{})
 	conn := meta.(*conns.AWSClient).EC2Conn
 
 	image_id := d.Get("image_id").(string)
-	account_id := d.Get("account_id").(string)
+
+	launch_permission := BuildLaunchPermission(d)
 
 	_, err := conn.ModifyImageAttribute(&ec2.ModifyImageAttributeInput{
 		ImageId:   aws.String(image_id),
 		Attribute: aws.String(ec2.ImageAttributeNameLaunchPermission),
 		LaunchPermission: &ec2.LaunchPermissionModifications{
-			Add: []*ec2.LaunchPermission{
-				{UserId: aws.String(account_id)},
-			},
+			Add: launch_permission,
 		},
 	})
 	if err != nil {
 		return fmt.Errorf("error creating AMI launch permission: %w", err)
 	}
 
+	account_id := d.Get("account_id").(string)
 	d.SetId(fmt.Sprintf("%s-%s", image_id, account_id))
 	return nil
 }
 
 func resourceAMILaunchPermissionRead(d *schema.ResourceData, meta interface{}) error {
 	conn := meta.(*conns.AWSClient).EC2Conn
 
-	exists, err := HasLaunchPermission(conn, d.Get("image_id").(string), d.Get("account_id").(string))
+	account_id := d.Get("account_id").(string)
+	arn := d.Get("arn").(string)
+	arn_type := d.Get("arn_type").(string)
+
+	var read_value, read_type string
+
+	if len(account_id) > 0 {
+		read_value = account_id
+		read_type = "UserId"
+	} else {
+		read_value = arn
+		read_type = arn_type
+	}
+
+	exists, err := HasLaunchPermission(conn, d.Get("image_id").(string), read_value, read_type)
 	if err != nil {
 		return fmt.Errorf("error reading AMI launch permission (%s): %w", d.Id(), err)
 	}
@@ -94,15 +147,14 @@ func resourceAMILaunchPermissionDelete(d *schema.ResourceData, meta interface{})
 	conn := meta.(*conns.AWSClient).EC2Conn
 
 	image_id := d.Get("image_id").(string)
-	account_id := d.Get("account_id").(string)
+
+	launch_permission := BuildLaunchPermission(d)
 
 	_, err := conn.ModifyImageAttribute(&ec2.ModifyImageAttributeInput{
 		ImageId:   aws.String(image_id),
 		Attribute: aws.String(ec2.ImageAttributeNameLaunchPermission),
 		LaunchPermission: &ec2.LaunchPermissionModifications{
-			Remove: []*ec2.LaunchPermission{
-				{UserId: aws.String(account_id)},
-			},
+			Remove: launch_permission,
 		},
 	})
 	if err != nil {
@@ -112,7 +164,7 @@ func resourceAMILaunchPermissionDelete(d *schema.ResourceData, meta interface{})
 	return nil
 }
 
-func HasLaunchPermission(conn *ec2.EC2, image_id string, account_id string) (bool, error) {
+func HasLaunchPermission(conn *ec2.EC2, image_id string, read_value string, read_type string) (bool, error) {
 	attrs, err := conn.DescribeImageAttribute(&ec2.DescribeImageAttributeInput{
 		ImageId:   aws.String(image_id),
 		Attribute: aws.String(ec2.ImageAttributeNameLaunchPermission),
@@ -121,16 +173,54 @@ func HasLaunchPermission(conn *ec2.EC2, image_id string, account_id string) (boo
 		// When an AMI disappears out from under a launch permission resource, we will
 		// see either InvalidAMIID.NotFound or InvalidAMIID.Unavailable.
 		if ec2err, ok := err.(awserr.Error); ok && strings.HasPrefix(ec2err.Code(), "InvalidAMIID") {
-			log.Printf("[DEBUG] %s no longer exists, so we'll drop launch permission for %s from the state", image_id, account_id)
+			log.Printf("[DEBUG] %s no longer exists, so we'll drop launch permission for %s from the state", image_id, read_value)
 			return false, nil
 		}
 		return false, err
 	}
 
 	for _, lp := range attrs.LaunchPermissions {
-		if aws.StringValue(lp.UserId) == account_id {
-			return true, nil
+		switch read_type {
+		case "UserId":
+			if aws.StringValue(lp.UserId) == read_value {
+				return true, nil
+			}
+		case "OrganizationArn":
+			if aws.StringValue(lp.OrganizationArn) == read_value {
+				return true, nil
+			}
+		case "OrganizationalUnitArn":
+			if aws.StringValue(lp.OrganizationalUnitArn) == read_value {
+				return true, nil
+			}
 		}
 	}
 	return false, nil
 }
+
+func BuildLaunchPermission(d *schema.ResourceData) []*ec2.LaunchPermission {
+	account_id := d.Get("account_id").(string)
+	arn := d.Get("arn").(string)
+	arn_type := d.Get("arn_type").(string)
+
+	var launch_permission []*ec2.LaunchPermission
+
+	if len(account_id) > 0 {
+		log.Printf("[DEBUG] Building LaunchPermission of type UserId: %s", account_id)
+		launch_permission = []*ec2.LaunchPermission{
+			{UserId: aws.String(account_id)},
+		}
+	} else if arn_type == "OrganizationArn" {
+		log.Printf("[DEBUG] Building LaunchPermission of type OrganizationArn: %s", arn)
+		launch_permission = []*ec2.LaunchPermission{
+			{OrganizationArn: aws.String(arn)},
+		}
+	} else if arn_type == "OrganizationalUnitArn" {
+		log.Printf("[DEBUG] Building LaunchPermission of type OrganizationalUnitArn: %s", arn)
+		launch_permission = []*ec2.LaunchPermission{
+			{OrganizationalUnitArn: aws.String(arn)},
+		}
+	}
+
+	return launch_permission
+}

internal/service/ec2/ami_launch_permission_test.go

@@ -40,6 +40,32 @@ func TestAccEC2AMILaunchPermission_basic(t *testing.T) {
 	})
 }
 
+func TestAccEC2AMILaunchPermission_org(t *testing.T) {
+	resourceName := "aws_ami_launch_permission.test"
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:     func() { acctest.PreCheck(t) },
+		ErrorCheck:   acctest.ErrorCheck(t, ec2.EndpointsID),
+		Providers:    acctest.Providers,
+		CheckDestroy: testAccCheckAMILaunchPermissionDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testOrgAMILaunchPermissionConfig(rName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAMILaunchPermissionExists(resourceName),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateIdFunc: testAccAMILaunchPermissionImportStateIdFunc(resourceName),
+				ImportStateVerify: true,
+			},
+		},
+	})
+}
+
 func TestAccEC2AMILaunchPermission_Disappears_launchPermission(t *testing.T) {
 	resourceName := "aws_ami_launch_permission.test"
 	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
@@ -149,12 +175,24 @@ func testAccCheckAMILaunchPermissionExists(resourceName string) resource.TestChe
 
 		conn := acctest.Provider.Meta().(*conns.AWSClient).EC2Conn
 		accountID := rs.Primary.Attributes["account_id"]
+		arn := rs.Primary.Attributes["arn"]
+		arn_type := rs.Primary.Attributes["arn_type"]
+
+		var read_value, read_type string
+		if len(accountID) > 0 {
+			read_value = accountID
+			read_type = "UserId"
+		} else {
+			read_value = arn
+			read_type = arn_type
+		}
+
 		imageID := rs.Primary.Attributes["image_id"]
 
-		if has, err := tfec2.HasLaunchPermission(conn, imageID, accountID); err != nil {
+		if has, err := tfec2.HasLaunchPermission(conn, imageID, read_value, read_type); err != nil {
 			return err
 		} else if !has {
-			return fmt.Errorf("launch permission does not exist for '%s' on '%s'", accountID, imageID)
+			return fmt.Errorf("launch permission does not exist for '%s' on '%s'", read_value, imageID)
 		}
 		return nil
 	}
@@ -168,12 +206,24 @@ func testAccCheckAMILaunchPermissionDestroy(s *terraform.State) error {
 
 		conn := acctest.Provider.Meta().(*conns.AWSClient).EC2Conn
 		accountID := rs.Primary.Attributes["account_id"]
+		arn := rs.Primary.Attributes["arn"]
+		arn_type := rs.Primary.Attributes["arn_type"]
+
+		var read_value, read_type string
+		if len(accountID) > 0 {
+			read_value = accountID
+			read_type = "UserId"
+		} else {
+			read_value = arn
+			read_type = arn_type
+		}
+
 		imageID := rs.Primary.Attributes["image_id"]
 
-		if has, err := tfec2.HasLaunchPermission(conn, imageID, accountID); err != nil {
+		if has, err := tfec2.HasLaunchPermission(conn, imageID, read_value, read_type); err != nil {
 			return err
 		} else if has {
-			return fmt.Errorf("launch permission still exists for '%s' on '%s'", accountID, imageID)
+			return fmt.Errorf("launch permission still exists for '%s' on '%s'", read_value, imageID)
 		}
 	}
 
@@ -298,13 +348,58 @@ resource "aws_ami_launch_permission" "test" {
 `, rName, rName)
 }
 
+func testOrgAMILaunchPermissionConfig(rName string) string {
+	return fmt.Sprintf(`
+data "aws_ami" "amzn-ami-minimal-hvm" {
+  most_recent = true
+  owners      = ["amazon"]
+
+  filter {
+    name   = "name"
+    values = ["amzn-ami-minimal-hvm-*"]
+  }
+
+  filter {
+    name   = "root-device-type"
+    values = ["ebs"]
+  }
+}
+
+data "aws_organizations_organization" "current" {}
+
+data "aws_region" "current" {}
+
+resource "aws_ami_copy" "test" {
+  description       = %q
+  name              = %q
+  source_ami_id     = data.aws_ami.amzn-ami-minimal-hvm.id
+  source_ami_region = data.aws_region.current.name
+}
+
+resource "aws_ami_launch_permission" "test" {
+  arn = data.aws_organizations_organization.current.arn
+  arn_type = "OrganizationArn"
+  image_id   = aws_ami_copy.test.id
+}
+`, rName, rName)
+}
+
 func testAccAMILaunchPermissionImportStateIdFunc(resourceName string) resource.ImportStateIdFunc {
 	return func(s *terraform.State) (string, error) {
 		rs, ok := s.RootModule().Resources[resourceName]
 		if !ok {
 			return "", fmt.Errorf("Not found: %s", resourceName)
 		}
 
-		return fmt.Sprintf("%s/%s", rs.Primary.Attributes["account_id"], rs.Primary.Attributes["image_id"]), nil
+		accountID := rs.Primary.Attributes["account_id"]
+		imageId := rs.Primary.Attributes["image_id"]
+
+		if len(accountID) > 0 {
+			return fmt.Sprintf("%s,%s", accountID, imageId), nil
+		} else {
+			arn := rs.Primary.Attributes["arn"]
+			arn_type := rs.Primary.Attributes["arn_type"]
+			return fmt.Sprintf("%s,%s,%s", arn, arn_type, imageId), nil
+		}
 	}
 }