Add VPC support to lambda functions

CHANGELOG.rst

@@ -11,6 +11,7 @@ CHANGELOG
 * Update ``policies.json`` file
   (`#817 <https://github.com/aws/chalice/issues/817>`__)
 
+
 1.2.2
 =====
 

chalice/app.py

@@ -9,8 +9,10 @@
 import base64
 from collections import defaultdict, Mapping
 
+
 __version__ = '1.2.3'
 
+
 # Implementation note:  This file is intended to be a standalone file
 # that gets copied into the lambda deployment package.  It has no dependencies
 # on other parts of chalice so it can stay small and lightweight, with minimal

chalice/awsclient.py

@@ -36,6 +36,7 @@
 _STR_MAP = Optional[Dict[str, str]]
 _OPT_STR = Optional[str]
 _OPT_INT = Optional[int]
+_OPT_STR_LIST = Optional[List[str]]
 _CLIENT_METHOD = Callable[..., Dict[str, Any]]
 
 
@@ -104,6 +105,22 @@ def get_function_configuration(self, name):
             FunctionName=name)
         return response
 
+    def _create_vpc_config(self, security_group_ids, subnet_ids):
+        # type: (_OPT_STR_LIST, _OPT_STR_LIST) -> Dict[str, List[str]]
+        # We always set the SubnetIds and SecurityGroupIds to an empty
+        # list to ensure that we properly remove Vpc configuration
+        # if you remove these values from your config.json.  Omitting
+        # the VpcConfig key or just setting to {} won't actually remove
+        # the VPC configuration.
+        vpc_config = {
+            'SubnetIds': [],
+            'SecurityGroupIds': [],
+        }  # type: Dict[str, List[str]]
+        if security_group_ids is not None and subnet_ids is not None:
+            vpc_config['SubnetIds'] = subnet_ids
+            vpc_config['SecurityGroupIds'] = security_group_ids
+        return vpc_config
+
     def create_function(self,
                         function_name,               # type: str
                         role_arn,                    # type: str
@@ -113,7 +130,9 @@ def create_function(self,
                         environment_variables=None,  # type: _STR_MAP
                         tags=None,                   # type: _STR_MAP
                         timeout=None,                # type: _OPT_INT
-                        memory_size=None             # type: _OPT_INT
+                        memory_size=None,            # type: _OPT_INT
+                        security_group_ids=None,     # type: _OPT_STR_LIST
+                        subnet_ids=None,             # type: _OPT_STR_LIST
                         ):
         # type: (...) -> str
         kwargs = {
@@ -131,14 +150,25 @@ def create_function(self,
             kwargs['Timeout'] = timeout
         if memory_size is not None:
             kwargs['MemorySize'] = memory_size
+        if security_group_ids is not None and subnet_ids is not None:
+            kwargs['VpcConfig'] = self._create_vpc_config(
+                security_group_ids=security_group_ids,
+                subnet_ids=subnet_ids,
+            )
+        return self._create_lambda_function(kwargs)
+
+    def _create_lambda_function(self, api_args):
+        # type: (Dict[str, Any]) -> str
         try:
             return self._call_client_method_with_retries(
-                self._client('lambda').create_function, kwargs)['FunctionArn']
+                self._client('lambda').create_function,
+                api_args
+            )['FunctionArn']
         except _REMOTE_CALL_ERRORS as e:
             context = LambdaErrorContext(
-                function_name,
+                api_args['FunctionName'],
                 'create_function',
-                len(zip_contents)
+                len(api_args['Code']['ZipFile']),
             )
             raise self._get_lambda_code_deployment_error(e, context)
 
@@ -208,7 +238,9 @@ def update_function(self,
                         tags=None,                   # type: _STR_MAP
                         timeout=None,                # type: _OPT_INT
                         memory_size=None,            # type: _OPT_INT
-                        role_arn=None                # type: _OPT_STR
+                        role_arn=None,               # type: _OPT_STR
+                        subnet_ids=None,             # type: _OPT_STR_LIST
+                        security_group_ids=None,     # type: _OPT_STR_LIST
                         ):
         # type: (...) -> Dict[str, Any]
         """Update a Lambda function's code and configuration.
@@ -217,9 +249,27 @@ def update_function(self,
         is not provided, no changes will be made for that that parameter on
         the targeted lambda function.
         """
+        return_value = self._update_function_code(function_name=function_name,
+                                                  zip_contents=zip_contents)
+        self._update_function_config(
+            environment_variables=environment_variables,
+            runtime=runtime,
+            timeout=timeout,
+            memory_size=memory_size,
+            role_arn=role_arn,
+            subnet_ids=subnet_ids,
+            security_group_ids=security_group_ids,
+            function_name=function_name
+        )
+        if tags is not None:
+            self._update_function_tags(return_value['FunctionArn'], tags)
+        return return_value
+
+    def _update_function_code(self, function_name, zip_contents):
+        # type: (str, str) -> Dict[str, Any]
         lambda_client = self._client('lambda')
         try:
-            return_value = lambda_client.update_function_code(
+            return lambda_client.update_function_code(
                 FunctionName=function_name, ZipFile=zip_contents)
         except _REMOTE_CALL_ERRORS as e:
             context = LambdaErrorContext(
@@ -229,6 +279,17 @@ def update_function(self,
             )
             raise self._get_lambda_code_deployment_error(e, context)
 
+    def _update_function_config(self,
+                                environment_variables,  # type: _STR_MAP
+                                runtime,                # type: _OPT_STR
+                                timeout,                # type: _OPT_INT
+                                memory_size,            # type: _OPT_INT
+                                role_arn,               # type: _OPT_STR
+                                subnet_ids,             # type: _OPT_STR_LIST
+                                security_group_ids,     # type: _OPT_STR_LIST
+                                function_name,          # type: str
+                                ):
+        # type: (...) -> None
         kwargs = {}  # type: Dict[str, Any]
         if environment_variables is not None:
             kwargs['Environment'] = {'Variables': environment_variables}
@@ -240,13 +301,16 @@ def update_function(self,
             kwargs['MemorySize'] = memory_size
         if role_arn is not None:
             kwargs['Role'] = role_arn
+        if security_group_ids is not None and subnet_ids is not None:
+            kwargs['VpcConfig'] = self._create_vpc_config(
+                subnet_ids=subnet_ids,
+                security_group_ids=security_group_ids
+            )
         if kwargs:
             kwargs['FunctionName'] = function_name
+            lambda_client = self._client('lambda')
             self._call_client_method_with_retries(
                 lambda_client.update_function_configuration, kwargs)
-        if tags is not None:
-            self._update_function_tags(return_value['FunctionArn'], tags)
-        return return_value
 
     def _update_function_tags(self, function_arn, requested_tags):
         # type: (str, Dict[str, str]) -> None

chalice/config.py

@@ -275,6 +275,20 @@ def tags(self):
             current_chalice_version, self.chalice_stage, self.app_name)
         return tags
 
+    @property
+    def security_group_ids(self):
+        # type: () -> List[str]
+        return self._chain_lookup('security_group_ids',
+                                  varies_per_chalice_stage=True,
+                                  varies_per_function=True)
+
+    @property
+    def subnet_ids(self):
+        # type: () -> List[str]
+        return self._chain_lookup('subnet_ids',
+                                  varies_per_chalice_stage=True,
+                                  varies_per_function=True)
+
     def scope(self, chalice_stage, function_name):
         # type: (str, str) -> Config
         # Used to create a new config object that's scoped to a different

chalice/constants.py

@@ -80,6 +80,18 @@ def index():
 }
 
 
+VPC_ATTACH_POLICY = {
+    "Effect": "Allow",
+    "Action": [
+        "ec2:CreateNetworkInterface",
+        "ec2:DescribeNetworkInterfaces",
+        "ec2:DetachNetworkInterface",
+        "ec2:DeleteNetworkInterface"
+    ],
+    "Resource": "*"
+}
+
+
 CODEBUILD_POLICY = {
     "Version": "2012-10-17",
     # This is the policy straight from the console.

chalice/deploy/deployer.py

@@ -100,8 +100,8 @@
 from chalice.compat import is_broken_pipe_error
 from chalice.awsclient import DeploymentPackageTooLargeError, TypedAWSClient
 from chalice.awsclient import LambdaClientError, AWSClientError
-from chalice.constants import MAX_LAMBDA_DEPLOYMENT_SIZE, \
-    LAMBDA_TRUST_POLICY, DEFAULT_LAMBDA_TIMEOUT, DEFAULT_LAMBDA_MEMORY_SIZE
+from chalice.constants import MAX_LAMBDA_DEPLOYMENT_SIZE, VPC_ATTACH_POLICY, \
+    DEFAULT_LAMBDA_TIMEOUT, DEFAULT_LAMBDA_MEMORY_SIZE, LAMBDA_TRUST_POLICY
 from chalice.deploy import models
 from chalice.deploy.packager import PipRunner, SubprocessPip, \
     DependencyBuilder as PipDependencyBuilder, LambdaDeploymentPackager
@@ -537,7 +537,9 @@ def _create_role_reference(self, config, stage_name, function_name):
             resource_name = 'default-role'
             role_name = '%s-%s' % (config.app_name, stage_name)
             policy = models.AutoGenIAMPolicy(
-                document=models.Placeholder.BUILD_STAGE)
+                document=models.Placeholder.BUILD_STAGE,
+                traits=set([]),
+            )
         return models.ManagedIAMRole(
             resource_name=resource_name,
             role_name=role_name,
@@ -555,7 +557,12 @@ def _build_lambda_function(self,
         # type: (...) -> models.LambdaFunction
         function_name = '%s-%s-%s' % (
             config.app_name, config.chalice_stage, name)
-        return models.LambdaFunction(
+        security_group_ids = config.security_group_ids
+        subnet_ids = config.subnet_ids
+        if security_group_ids is None or subnet_ids is None:
+            security_group_ids = []
+            subnet_ids = []
+        function = models.LambdaFunction(
             resource_name=name,
             function_name=function_name,
             environment_variables=config.environment_variables,
@@ -566,7 +573,21 @@ def _build_lambda_function(self,
             memory_size=config.lambda_memory_size,
             deployment_package=deployment,
             role=role,
+            security_group_ids=security_group_ids,
+            subnet_ids=subnet_ids,
         )
+        self._inject_role_traits(function, role)
+        return function
+
+    def _inject_role_traits(self, function, role):
+        # type: (models.LambdaFunction, models.IAMRole) -> None
+        if not isinstance(role, models.ManagedIAMRole):
+            return
+        policy = role.policy
+        if not isinstance(policy, models.AutoGenIAMPolicy):
+            return
+        if function.security_group_ids and function.subnet_ids:
+            policy.traits.add(models.RoleTraits.VPC_NEEDED)
 
 
 class DependencyBuilder(object):
@@ -664,7 +685,10 @@ def handle_filebasediampolicy(self, config, resource):
     def handle_autogeniampolicy(self, config, resource):
         # type: (Config, models.AutoGenIAMPolicy) -> None
         if isinstance(resource.document, models.Placeholder):
-            resource.document = self._policy_gen.generate_policy(config)
+            policy = self._policy_gen.generate_policy(config)
+            if models.RoleTraits.VPC_NEEDED in resource.traits:
+                policy['Statement'].append(VPC_ATTACH_POLICY)
+            resource.document = policy
 
 
 class BuildStage(object):
@@ -891,94 +915,3 @@ def display_report(self, deployed_values):
         # type: (Dict[str, Any]) -> None
         report = self.generate_report(deployed_values)
         self._ui.write(report)
-
-
-class ApplicationPolicyHandler(object):
-    """Manages the IAM policy for an application.
-
-    This class handles returning the policy that used by
-    used for the API handler lambda function for a given
-    stage.
-
-    It has several possible outcomes:
-
-        * By default, it will autogenerate a policy based on
-          analyzing the application source code.
-        * It will return a policy from a file on disk that's been
-          configured as the policy for the given stage.
-
-    This class has a precondition that we should be loading
-    some IAM policy for the the API handler function.
-
-    If a user has indicated that there's a pre-existing
-    role that they'd like to use for the API handler function,
-    this class should never be invoked.  In other words,
-    the logic of whether or not we even need to bother with
-    loading an IAM policy is handled a layer above where
-    this class should be used.
-
-    """
-
-    _EMPTY_POLICY = {
-        'Version': '2012-10-17',
-        'Statement': [],
-    }
-
-    def __init__(self, osutils, policy_generator):
-        # type: (OSUtils, AppPolicyGenerator) -> None
-        self._osutils = osutils
-        self._policy_gen = policy_generator
-
-    def generate_policy_from_app_source(self, config):
-        # type: (Config) -> Dict[str, Any]
-        """Generate a policy from application source code.
-
-        If the ``autogen_policy`` value is set to false, then
-        the .chalice/policy.json file will be used instead of generating
-        the policy from the source code.
-
-        """
-        if config.autogen_policy:
-            app_policy = self._do_generate_from_source(config)
-        else:
-            app_policy = self.load_last_policy(config)
-        return app_policy
-
-    def _do_generate_from_source(self, config):
-        # type: (Config) -> Dict[str, Any]
-        return self._policy_gen.generate_policy(config)
-
-    def load_last_policy(self, config):
-        # type: (Config) -> Dict[str, Any]
-        """Load the last recorded policy file for the app."""
-        filename = self._app_policy_file(config)
-        if config.autogen_policy and not self._osutils.file_exists(filename):
-            return self._EMPTY_POLICY
-        elif not self._osutils.file_exists(filename):
-            raise RuntimeError("Unable to load the policy file. Are you sure "
-                               "it exists?")
-        try:
-            return json.loads(
-                self._osutils.get_file_contents(filename, binary=False)
-            )
-        except ValueError as err:
-            raise RuntimeError("Unable to load the project policy file: %s"
-                               % err)
-
-    def record_policy(self, config, policy_document):
-        # type: (Config, Dict[str, Any]) -> None
-        policy_file = self._app_policy_file(config)
-        policy_json = serialize_to_json(policy_document)
-        self._osutils.set_file_contents(policy_file, policy_json, binary=False)
-
-    def _app_policy_file(self, config):
-        # type: (Config) -> str
-        if config.iam_policy_file:
-            filename = os.path.join(config.project_dir, '.chalice',
-                                    config.iam_policy_file)
-        else:
-            # Otherwise if the user doesn't specify a file it defaults
-            # to a fixed name based on the stage.
-            basename = 'policy-%s.json' % config.chalice_stage
-            filename = os.path.join(config.project_dir, '.chalice', basename)
-        return filename