EC2: improve UserData handling (#9510)

Author: bpandola

moto/autoscaling/responses.py

@@ -1,6 +1,6 @@
 from moto.core.common_types import TYPE_RESPONSE
 from moto.core.responses import ActionResult, BaseResponse, EmptyResult
-from moto.core.types import Base64EncodedString
+from moto.ec2.utils import parse_user_data
 from moto.utilities.aws_headers import amz_crc32
 
 from .models import AutoScalingBackend, autoscaling_backends
@@ -21,9 +21,7 @@ def call_action(self) -> TYPE_RESPONSE:
 
     def create_launch_configuration(self) -> ActionResult:
         params = self._get_params()
-        user_data = params.get("UserData")
-        if user_data is not None:
-            user_data = Base64EncodedString(user_data)
+        user_data = parse_user_data(params.get("UserData"))
         self.autoscaling_backend.create_launch_configuration(
             name=params.get("LaunchConfigurationName"),  # type: ignore[arg-type]
             image_id=params.get("ImageId"),  # type: ignore[arg-type]

moto/ec2/exceptions.py

@@ -913,3 +913,11 @@ def __init__(self) -> None:
             "AuthFailure",
             "Unauthorized attempt to access restricted resource",
         )
+
+
+class InvalidUserDataError(EC2ClientError):
+    def __init__(self, message: str):
+        super().__init__(
+            "InvalidUserData.Malformed",
+            message,
+        )

moto/ec2/models/launch_templates.py

@@ -13,6 +13,7 @@
 from ..utils import (
     convert_tag_spec,
     generic_filter,
+    parse_user_data,
     random_launch_template_id,
     random_launch_template_name,
     utc_date_and_time,
@@ -51,9 +52,10 @@ def security_groups(self) -> list[str]:
 
     @property
     def user_data(self) -> Optional[Base64EncodedString]:
-        user_data = self.data.get("UserData", None)
-        if user_data is not None:
-            user_data = Base64EncodedString(user_data)
+        user_data = self.data.get("UserData")
+        # UserData can be specified via multiple services/api endpoints,
+        # so we make an assertion here that it's in the format we expect.
+        assert user_data is None or isinstance(user_data, Base64EncodedString)
         return user_data
 
 
@@ -143,6 +145,7 @@ def create_from_cloudformation_json(  # type: ignore[misc]
         properties = cloudformation_json["Properties"]
         name = properties.get("LaunchTemplateName")
         data = properties.get("LaunchTemplateData")
+        data["UserData"] = parse_user_data(data.get("UserData"))
         description = properties.get("VersionDescription")
         tag_spec = convert_tag_spec(
             properties.get("TagSpecifications", {}), tag_key="Tags"
@@ -173,6 +176,7 @@ def update_from_cloudformation_json(  # type: ignore[misc]
         properties = cloudformation_json["Properties"]
 
         data = properties.get("LaunchTemplateData")
+        data["UserData"] = parse_user_data(data.get("UserData"))
         description = properties.get("VersionDescription")
 
         launch_template = backend.get_launch_template(original_resource.id)

moto/ec2/models/spot_requests.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 from moto.core.common_models import BaseModel, CloudFormationModel
+from moto.core.types import Base64EncodedString
 from moto.ec2.exceptions import InvalidParameterValueErrorTagSpotFleetRequest
 
 if TYPE_CHECKING:
@@ -56,7 +57,7 @@ def __init__(
         availability_zone_group: Optional[str],
         key_name: str,
         security_groups: list[str],
-        user_data: dict[str, Any],
+        user_data: Optional[Base64EncodedString],
         instance_type: str,
         placement: Optional[str],
         kernel_id: Optional[str],
@@ -411,7 +412,7 @@ def request_spot_instances(
         availability_zone_group: Optional[str],
         key_name: str,
         security_groups: list[str],
-        user_data: dict[str, Any],
+        user_data: Optional[Base64EncodedString],
         instance_type: str,
         placement: Optional[str],
         kernel_id: Optional[str],

moto/ec2/responses/instances.py

@@ -2,14 +2,13 @@
 from typing import Any
 
 from moto.core.responses import ActionResult, EmptyResult
-from moto.core.types import Base64EncodedString
 from moto.core.utils import camelcase_to_underscores
 from moto.ec2.exceptions import (
     InvalidParameterCombination,
     InvalidRequest,
     MissingParameterError,
 )
-from moto.ec2.utils import filter_iam_instance_profiles
+from moto.ec2.utils import filter_iam_instance_profiles, parse_user_data
 
 from ._base_response import EC2BaseResponse
 
@@ -49,9 +48,7 @@ def describe_instances(self) -> ActionResult:
     def run_instances(self) -> ActionResult:
         min_count = int(self._get_param("MinCount", if_none="1"))
         image_id = self._get_param("ImageId")
-        user_data = self._get_param("UserData")
-        if user_data is not None:
-            user_data = Base64EncodedString(user_data)
+        user_data = parse_user_data(self._get_param("UserData"))
         security_group_names = self._get_param("SecurityGroups", [])
         kwargs = {
             "instance_type": self._get_param("InstanceType", "m1.small"),
@@ -340,7 +337,7 @@ def _dot_value_instance_attribute_handler(self) -> bool:
                 attr_name = camelcase_to_underscores(attribute)
                 attr_value = self._get_param(f"{attribute}.Value")
                 if attribute == "UserData" and attr_value:
-                    attr_value = Base64EncodedString.from_encoded_bytes(attr_value)
+                    attr_value = parse_user_data(attr_value)
                 self.ec2_backend.modify_instance_attribute(
                     instance_id, attr_name, attr_value
                 )

moto/ec2/responses/launch_templates.py

@@ -4,6 +4,7 @@
 
 from moto.core.types import Base64EncodedString
 from moto.ec2.exceptions import FilterNotImplementedError
+from moto.ec2.utils import parse_user_data
 from moto.moto_api._internal import mock_random
 
 from ._base_response import EC2BaseResponse
@@ -20,6 +21,8 @@ def xml_root(name: str) -> ElementTree.Element:
 
 
 def xml_serialize(tree: ElementTree.Element, key: str, value: Any) -> None:
+    if value is None:
+        return
     name = key[0].lower() + key[1:]
     if isinstance(value, list):
         if name[-1] == "s":
@@ -29,7 +32,7 @@ def xml_serialize(tree: ElementTree.Element, key: str, value: Any) -> None:
 
     node = ElementTree.SubElement(tree, name)
 
-    if isinstance(value, (str, int, float, str)):
+    if isinstance(value, (str, int, float, str, Base64EncodedString)):
         node.text = str(value)
     elif isinstance(value, dict):
         for dictkey, dictvalue in value.items():
@@ -58,10 +61,9 @@ def create_launch_template(self) -> str:
         tag_spec = self._parse_tag_specification()
 
         parsed_template_data = self._get_param("LaunchTemplateData", {})
-        if parsed_template_data.get("UserData"):
-            parsed_template_data["UserData"] = Base64EncodedString(
-                parsed_template_data["UserData"]
-            )
+        parsed_template_data["UserData"] = parse_user_data(
+            self._get_param("LaunchTemplateData.UserData")
+        )
         self.error_on_dryrun()
 
         if tag_spec:

moto/ec2/responses/spot_instances.py

@@ -1,3 +1,4 @@
+from ..utils import parse_user_data
 from ._base_response import EC2BaseResponse
 
 
@@ -59,7 +60,7 @@ def request_spot_instances(self) -> str:
         availability_zone_group = self._get_param("AvailabilityZoneGroup")
         key_name = self._get_param("LaunchSpecification.KeyName")
         security_groups = self._get_param("LaunchSpecification.SecurityGroups", [])
-        user_data = self._get_param("LaunchSpecification.UserData")
+        user_data = parse_user_data(self._get_param("LaunchSpecification.UserData"))
         instance_type = self._get_param("LaunchSpecification.InstanceType", "m1.small")
         placement = self._get_param("LaunchSpecification.Placement.AvailabilityZone")
         kernel_id = self._get_param("LaunchSpecification.KernelId")

moto/ec2/utils.py

@@ -22,7 +22,9 @@
 )
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
 
+from moto.core.types import Base64EncodedString
 from moto.core.utils import utcnow
+from moto.ec2.exceptions import InvalidUserDataError
 from moto.iam import iam_backends
 from moto.moto_api._internal import mock_random as random
 from moto.utilities.utils import md5_hash
@@ -931,3 +933,16 @@ def convert_tag_spec(
             {tag["Key"]: tag["Value"] for tag in tag_spec[tag_key]}
         )
     return tags
+
+
+def parse_user_data(value: Any) -> Optional[Base64EncodedString]:
+    if value is None:
+        return None
+    try:
+        if isinstance(value, bytes):
+            user_data = Base64EncodedString.from_encoded_bytes(value)
+        else:
+            user_data = Base64EncodedString(value)
+    except ValueError:
+        raise InvalidUserDataError("Invalid BASE64 encoding of user data.")
+    return user_data

tests/test_autoscaling/test_autoscaling_groups.py

@@ -4,6 +4,7 @@
 
 from moto import mock_aws
 from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID
+from moto.core.types import Base64EncodedString
 from tests import EXAMPLE_AMI_ID
 
 from .utils import setup_networking
@@ -363,3 +364,35 @@ def test_launch_template_with_tags():
     tags = instances["Reservations"][0]["Instances"][0]["Tags"]
     assert {"Value": "TestTagValue1", "Key": "TestTagKey1"} in tags
     assert {"Key": "from_lt", "Value": "val"} in tags
+
+
+@mock_aws
+def test_launch_template_with_user_data():
+    mocked_networking = setup_networking()
+    user_data = Base64EncodedString.from_raw_string("test user data")
+    ec2_client = boto3.client("ec2", region_name="us-east-1")
+    template = ec2_client.create_launch_template(
+        LaunchTemplateName="test_launch_template",
+        LaunchTemplateData={
+            "ImageId": EXAMPLE_AMI_ID,
+            "InstanceType": "t2.micro",
+            "UserData": str(user_data),
+        },
+    )["LaunchTemplate"]
+    as_client = boto3.client("autoscaling", region_name="us-east-1")
+    as_client.create_auto_scaling_group(
+        AutoScalingGroupName="myasgroup",
+        MinSize=2,
+        MaxSize=3,
+        LaunchTemplate={"LaunchTemplateId": template["LaunchTemplateId"]},
+        VPCZoneIdentifier=mocked_networking["subnet1"],
+    )
+    resp = ec2_client.describe_instances(
+        Filters=[{"Name": "tag:aws:autoscaling:groupName", "Values": ["myasgroup"]}]
+    )
+    instances = resp["Reservations"][0]["Instances"]
+    for instance in instances:
+        attr_resp = ec2_client.describe_instance_attribute(
+            InstanceId=instance["InstanceId"], Attribute="userData"
+        )
+        assert attr_resp["UserData"]["Value"] == str(user_data)

tests/test_ec2/test_launch_templates.py

@@ -763,3 +763,20 @@ def test_modify_launch_template_by_name():
     assert (
         version["LaunchTemplateId"] == create_resp["LaunchTemplate"]["LaunchTemplateId"]
     )
+
+
+@mock_aws
+def test_create_launch_template_with_non_base64_encoded_user_data_fails():
+    client = boto3.client("ec2", region_name="us-east-1")
+    with pytest.raises(ClientError) as exc:
+        client.create_launch_template(
+            LaunchTemplateName="test-template",
+            LaunchTemplateData={
+                "ImageId": "ami-12345678",
+                "InstanceType": "t2.nano",
+                "UserData": "not base64 encoded",
+            },
+        )
+    error = exc.value.response["Error"]
+    assert error["Code"] == "InvalidUserData.Malformed"
+    assert error["Message"] == "Invalid BASE64 encoding of user data."