temporary commit

Signed-off-by: Hanwen <[email protected]>

Author: hanwen-cluster

cli/src/pcluster/aws/aws_resources.py

@@ -37,6 +37,7 @@ def __init__(self, stack_data: dict):
         self._params = self._stack_data.get("Parameters", [])
         self.tags = self._stack_data.get("Tags", [])
         self.outputs = self._stack_data.get("Outputs", [])
+        self.__resources = None
 
     @property
     def id(self):
@@ -48,6 +49,15 @@ def name(self):
         """Return the name of the stack."""
         return self._stack_data.get("StackName")
 
+    @property
+    def resources(self):
+        """Return the resources of the stack."""
+        if not self.__resources:
+            from pcluster.aws.aws_api import AWSApi
+
+            self.__resources = AWSApi.instance().cfn.describe_stack_resources(self.name)
+        return self.__resources
+
     @property
     def status(self):
         """Return the status of the stack."""
@@ -90,6 +100,10 @@ def _get_param(self, key_name):
         param_value = next((par["ParameterValue"] for par in self._params if par["ParameterKey"] == key_name), None)
         return None if param_value is None else param_value.strip()
 
+    def get_resource(self, resource_logical_id: str):
+        """Return the resource information."""
+        return self.resources.get(resource_logical_id)
+
 
 class InstanceInfo:
     """Object to store Instance information, initialized with a describe_instances call."""

cli/src/pcluster/aws/cfn.py

@@ -147,6 +147,12 @@ def describe_stack_resource(self, stack_name: str, logic_resource_id: str):
                 function_name="describe_stack_resource", message=f"No resource {logic_resource_id} found."
             )
 
+    @AWSExceptionHandler.handle_client_exception
+    def describe_stack_resources(self, stack_name: str):
+        """Get stack resources information."""
+        response = self._client.describe_stack_resources(StackName=stack_name).get("StackResources")
+        return {resource["LogicalResourceId"]: resource for resource in response}  # Build dictionary for better query.
+
     @AWSExceptionHandler.handle_client_exception
     def get_imagebuilder_stacks(self, next_token=None):
         """List existing imagebuilder stacks."""

cli/src/pcluster/aws/ec2.py

@@ -136,7 +136,7 @@ def get_subnet_vpc(self, subnet_id):
     @Cache.cached
     def get_subnet_cidr(self, subnet_id):
         """Return cidr block  of the given subnet."""
-        subnets = self._client.describe_subnets(SubnetIds=[subnet_id]).get("Subnets")
+        subnets = self.describe_subnets([subnet_id])
         if subnets:
             return subnets[0].get("CidrBlock")
         raise AWSClientError(function_name="describe_subnets", message=f"Subnet {subnet_id} not found")

cli/src/pcluster/config/cluster_config.py

@@ -1261,6 +1261,8 @@ def __init__(
         self._official_ami = None
         self.imds = imds or TopLevelImds(implied="v1.0")
         self.deployment_settings = deployment_settings
+        self.managed_head_node_security_group = None
+        self.managed_compute_security_group = None
 
     def _register_validators(self, context: ValidatorContext = None):  # noqa: D102 #pylint: disable=unused-argument
         self._register_validator(RegionValidator, region=self.region)
@@ -1387,15 +1389,15 @@ def _register_storage_validators(self):
                             EfsIdValidator,
                             efs_id=storage.file_system_id,
                             avail_zones_mapping=self.availability_zones_subnets_mapping,
-                            are_all_security_groups_customized=self.are_all_security_groups_customized,
+                            all_security_groups=self.all_security_groups,
                         )
                     else:
                         new_storage_count["efs"] += 1
             self._register_validator(
                 ExistingFsxNetworkingValidator,
                 file_system_ids=list(existing_fsx),
-                head_node_subnet_id=self.head_node.networking.subnet_id,
-                are_all_security_groups_customized=self.are_all_security_groups_customized,
+                subnet_ids=[self.head_node.networking.subnet_id] + self.compute_subnet_ids,
+                all_security_groups=self.all_security_groups,
             )
 
             self._validate_max_storage_count(ebs_count, existing_storage_count, new_storage_count)
@@ -1624,17 +1626,29 @@ def is_dcv_enabled(self):
         return self.head_node.dcv and self.head_node.dcv.enabled
 
     @property
-    def are_all_security_groups_customized(self):
+    def all_security_groups(self):
         """Return True if all head node and queues have (additional) security groups specified."""
         head_node_networking = self.head_node.networking
+        security_groups_for_head_node = set()
+        if head_node_networking.security_groups:
+            security_groups_for_head_node.update(set(head_node_networking.security_groups))
+        if head_node_networking.additional_security_groups:
+            security_groups_for_head_node.update(set(head_node_networking.additional_security_groups))
         if not (head_node_networking.security_groups or head_node_networking.additional_security_groups):
-            return False
+            security_groups_for_head_node.add(self.managed_head_node_security_group)
+        security_groups_for_all_nodes = {frozenset(security_groups_for_head_node)}
         for queue in self.scheduling.queues:
             queue_networking = queue.networking
             if isinstance(queue_networking, _QueueNetworking):
-                if not (queue_networking.security_groups or queue_networking.additional_security_groups):
-                    return False
-        return True
+                security_groups_for_compute_node = set()
+                if queue_networking.security_groups:
+                    security_groups_for_compute_node.update(set(queue_networking.security_groups))
+                else:
+                    security_groups_for_compute_node.add(self.managed_compute_security_group)
+                if queue_networking.additional_security_groups:
+                    security_groups_for_compute_node.update(set(queue_networking.additional_security_groups))
+                security_groups_for_all_nodes.add(frozenset(security_groups_for_compute_node))
+        return security_groups_for_all_nodes
 
     @property
     def extra_chef_attributes(self):

cli/src/pcluster/models/cluster.py

@@ -440,6 +440,9 @@ def _validate_and_parse_config(
             Cluster._load_additional_instance_type_data(cluster_config_dict)
             config = self._load_config(cluster_config_dict)
             config.official_ami = self.__official_ami
+            if context.during_update:
+                config.managed_head_node_security_group = self.stack.get_resource("HeadNodeSecurityGroup")
+                config.managed_compute_security_group = self.stack.get_resource("ComputeSecurityGroup")
 
             validation_failures = config.validate(validator_suppressors, context)
             if any(f.level.value >= FailureLevel(validation_failure_level).value for f in validation_failures):
@@ -851,7 +854,7 @@ def validate_update_request(
             validator_suppressors=validator_suppressors,
             validation_failure_level=validation_failure_level,
             config_text=target_source_config,
-            context=ValidatorContext(head_node_instance_id=self.head_node_instance.id),
+            context=ValidatorContext(head_node_instance_id=self.head_node_instance.id, during_update=True),
         )
         changes = self._validate_patch(force, target_config)
 

cli/src/pcluster/validators/cluster_validators.py

@@ -12,7 +12,7 @@
 import re
 from collections import defaultdict
 from enum import Enum
-from ipaddress import ip_network
+from ipaddress import collapse_addresses, ip_network
 from itertools import combinations, product
 from typing import List
 
@@ -417,42 +417,64 @@ def _validate(
 # --------------- Storage validators --------------- #
 
 
-def _check_in_out_access(security_groups_ids, port, is_cidr_optional, protocol="tcp"):
+def _is_access_allowed(security_groups_ids, subnets, port, all_security_groups, protocol="tcp"):
     """
     Verify given list of security groups to check if they allow in and out access on the given port.
 
     :param security_groups_ids: list of security groups to verify
     :param port: port to verify
-    :param is_cidr_optional: if it is True, don't enforce check on CIDR.
+    :param all_security_groups: all security groups from cluster. This is a set of frozen sets.
+    Each frozen set contains sg combination of a queue.
     :param protocol: the IP protocol to be checked.
     :return: True if both in and out access are allowed
     :raise: ClientError if a given security group doesn't exist
     """
     in_access = False
     out_access = False
 
+    in_ip_ranges = []
+    out_ip_ranges = []
     for sec_group in AWSApi.instance().ec2.describe_security_groups(security_groups_ids):
         # Check all inbound rules
         for rule in sec_group.get("IpPermissions"):
-            if _check_sg_rules_for_port(rule, port, protocol):
-                if is_cidr_optional or rule.get("IpRanges") or rule.get("PrefixListIds"):
-                    in_access = True
-                    break
+            if in_access:
+                break
+            if _is_port_allowed_by_sg_rule(rule, port, protocol):
+                in_access = _is_src_or_dst_allowed_by_sg_rule(rule, all_security_groups, in_ip_ranges)
 
         # Check all outbound rules
         for rule in sec_group.get("IpPermissionsEgress"):
-            if _check_sg_rules_for_port(rule, port, protocol):
-                if is_cidr_optional or rule.get("IpRanges") or rule.get("PrefixListIds"):
-                    out_access = True
-                    break
+            if out_access:
+                break
+            if _is_port_allowed_by_sg_rule(rule, port, protocol):
+                out_access = _is_src_or_dst_allowed_by_sg_rule(rule, all_security_groups, out_ip_ranges)
 
         if in_access and out_access:
             return True
 
+    # Rules of ip ranges have to be checked at the end because the union of all ip ranges may cover the subnets,
+    # even when individual ip ranges do not cover the subnets
+    in_access = in_access or _are_subnets_covered_by_cidrs(in_ip_ranges, subnets)
+    out_access = out_access or _are_subnets_covered_by_cidrs(out_ip_ranges, subnets)
+    return in_access and out_access
+
+
+def _is_src_or_dst_allowed_by_sg_rule(rule, all_security_groups, ip_ranges):
+    if rule.get("IpRanges"):
+        ip_ranges.extend(rule.get("IpRanges"))
+        return False  # Ip Ranges have to be checked later. Return False because the rule allowance is not determined.
+    elif rule.get("PrefixListIds"):
+        return True  # Always assume prefix list is properly set for code simplicity
+    elif rule.get("UserIdGroupPairs"):
+        allowed_security_groups = {
+            user_id_group_pair.get("GroupId") for user_id_group_pair in rule.get("UserIdGroupPairs")
+        }
+        # For all cluster nodes, at least one of the security groups attached need to be in the UserIdGroupPairs.
+        return all(node_security_groups & allowed_security_groups for node_security_groups in all_security_groups)
     return False
 
 
-def _check_sg_rules_for_port(rule, port_to_check, protocol):
+def _is_port_allowed_by_sg_rule(rule, port_to_check, protocol):
     """
     Verify if the security group rule accepts connections on the given port.
 
@@ -483,6 +505,23 @@ def _check_sg_rules_for_port(rule, port_to_check, protocol):
     return False
 
 
+def _are_subnets_covered_by_cidrs(ip_ranges, subnets):
+    """Verify given list of security groups to check if they allow in and out access on cluster subnet CIDRs."""
+    # Collapse ip ranges for better performance and correctness
+    collapsed_ip_ranges = list(collapse_addresses([ip_network(ip_range["CidrIp"]) for ip_range in ip_ranges]))
+
+    for subnet in subnets:
+        subnet_cidr = ip_network(AWSApi.instance().ec2.get_subnet_cidr(subnet))
+        covered = False
+        for ip_range in collapsed_ip_ranges:
+            if ip_range.supernet_of(subnet_cidr):
+                covered = True
+                break
+        if not covered:
+            return False
+    return True
+
+
 class ExistingFsxNetworkingValidator(Validator):
     """
     FSx networking validator.
@@ -504,27 +543,23 @@ def _describe_network_interfaces(self, file_systems):
         else:
             return {}
 
-    def _validate(self, file_system_ids, head_node_subnet_id, are_all_security_groups_customized):
+    def _validate(self, file_system_ids, subnet_ids, all_security_groups):
         try:
-            # Check to see if there is any existing mt on the fs
             file_systems = AWSApi.instance().fsx.get_file_systems_info(file_system_ids)
-
-            vpc_id = AWSApi.instance().ec2.get_subnet_vpc(head_node_subnet_id)
-
-            network_interfaces_data = self._describe_network_interfaces(file_systems)
-
-            self._check_file_systems(are_all_security_groups_customized, file_systems, network_interfaces_data, vpc_id)
+            self._check_file_systems(all_security_groups, file_systems, subnet_ids)
         except AWSClientError as e:
             self._add_failure(str(e), FailureLevel.ERROR)
 
-    def _check_file_systems(self, are_all_security_groups_customized, file_systems, network_interfaces_data, vpc_id):
+    def _check_file_systems(self, all_security_groups, file_systems, subnet_ids):
+        vpc_id = AWSApi.instance().ec2.get_subnet_vpc(subnet_ids[0])
+        network_interfaces_data = self._describe_network_interfaces(file_systems)
         for file_system in file_systems:
             # Check to see if fs is in the same VPC as the stack
             file_system_id = file_system.file_system_id
             if file_system.vpc_id != vpc_id:
                 self._add_failure(
                     "Currently only support using FSx file system that is in the same VPC as the cluster. "
-                    "The file system provided is in {0}.".format(file_system.vpc_id),
+                    f"The file system {file_system_id} is in {file_system.vpc_id}.",
                     FailureLevel.ERROR,
                 )
 
@@ -545,7 +580,7 @@ def _check_file_systems(self, are_all_security_groups_customized, file_systems,
 
                 for protocol, ports in FSX_PORTS[file_system.file_system_type].items():
                     missing_ports = self._get_missing_ports(
-                        are_all_security_groups_customized, network_interfaces, ports, protocol
+                        all_security_groups, subnet_ids, network_interfaces, ports, protocol
                     )
 
                     if missing_ports:
@@ -557,17 +592,18 @@ def _check_file_systems(self, are_all_security_groups_customized, file_systems,
                             FailureLevel.ERROR,
                         )
 
-    def _get_missing_ports(self, are_all_security_groups_customized, network_interfaces, ports, protocol):
+    def _get_missing_ports(self, all_security_groups, subnet_ids, network_interfaces, ports, protocol):
         missing_ports = []
         for port in ports:
             fs_access = False
             for network_interface in network_interfaces:
                 # Get list of security group IDs
                 sg_ids = [sg.get("GroupId") for sg in network_interface.get("Groups")]
-                if _check_in_out_access(
+                if _is_access_allowed(
                     sg_ids,
+                    subnet_ids,
                     port=port,
-                    is_cidr_optional=are_all_security_groups_customized,
+                    all_security_groups=all_security_groups,
                     protocol=protocol,
                 ):
                     fs_access = True
@@ -743,7 +779,7 @@ class EfsIdValidator(Validator):  # TODO add tests
     Validate if there are existing mount target in the cluster (head and computes) availability zone
     """
 
-    def _validate(self, efs_id, avail_zones_mapping: dict, are_all_security_groups_customized):
+    def _validate(self, efs_id, avail_zones_mapping: dict, all_security_groups):
         availability_zones = avail_zones_mapping.keys()
         if len(availability_zones) > 1 and not AWSApi.instance().efs.is_efs_standard(efs_id):
             self._add_failure(
@@ -760,7 +796,7 @@ def _validate(self, efs_id, avail_zones_mapping: dict, are_all_security_groups_c
             if head_node_target_id:
                 # Get list of security group IDs of the mount target
                 sg_ids = AWSApi.instance().efs.get_efs_mount_target_security_groups(head_node_target_id)
-                if not _check_in_out_access(sg_ids, port=2049, is_cidr_optional=are_all_security_groups_customized):
+                if not _is_access_allowed(sg_ids, subnets, port=2049, all_security_groups=all_security_groups):
                     self._add_failure(
                         "There is an existing Mount Target {0} in the Availability Zone {1} for EFS {2}, "
                         "but it does not have a security group that allows inbound and outbound rules to support NFS. "
@@ -769,8 +805,6 @@ def _validate(self, efs_id, avail_zones_mapping: dict, are_all_security_groups_c
                         ),
                         FailureLevel.ERROR,
                     )
-                if not are_all_security_groups_customized:
-                    self._check_cidrs_cover_subnets(head_node_target_id, avail_zone, sg_ids, efs_id, subnets)
             else:
                 if AWSApi.instance().efs.is_efs_standard(efs_id):
                     avail_zones_missing_mount_target_for_efs_standard.append(avail_zone)
@@ -783,42 +817,6 @@ def _validate(self, efs_id, avail_zones_mapping: dict, are_all_security_groups_c
                 FailureLevel.ERROR,
             )
 
-    def _check_subnet_access(self, security_groups_ids, subnet_cidr, access_type):
-        permission = "IpPermissions" if access_type == "in" else "IpPermissionsEgress"
-        access = False
-        for sec_group in security_groups_ids:
-            for rule in sec_group.get(permission):
-                if rule.get("PrefixListIds"):
-                    access = True
-                    break
-                if rule.get("IpRanges"):
-                    for ip_range in rule.get("IpRanges"):
-                        if ip_network(ip_range.get("CidrIp")).supernet_of(subnet_cidr):
-                            access = True
-                            break
-        return access
-
-    def _check_cidrs_cover_subnets(self, head_node_target_id, avail_zone, security_groups_ids, efs_id, subnets):
-        """Verify given list of security groups to check if they allow in and out access on cluster subnet CIDRs."""
-        security_groups_ids = AWSApi.instance().ec2.describe_security_groups(security_groups_ids)
-        for subnet in subnets:
-            subnet_cidr = ip_network(AWSApi.instance().ec2.get_subnet_cidr(subnet))
-            in_access, out_access = self._check_subnet_access(
-                security_groups_ids, subnet_cidr, "in"
-            ), self._check_subnet_access(security_groups_ids, subnet_cidr, "out")
-
-            if not in_access or not out_access:
-                self._add_failure(
-                    "There is an existing Mount Target {0} in the Availability Zone {1} for EFS {2}, "
-                    "but it does not have a security group that allows inbound and outbound rules to allow traffic of "
-                    "subnet {3}. Please modify the Mount Target's security group, to allow traffic on subnet.".format(
-                        head_node_target_id, avail_zone, efs_id, subnet
-                    ),
-                    FailureLevel.WARNING,
-                )
-
-        return False
-
 
 class SharedStorageNameValidator(Validator):
     """

cli/src/pcluster/validators/common.py

@@ -69,5 +69,6 @@ def _validate(self, *args, **kwargs):
 class ValidatorContext:
     """Context containing information about cluster environment meant to be passed to validators."""
 
-    def __init__(self, head_node_instance_id: str = None):
+    def __init__(self, head_node_instance_id: str = None, during_update: bool = None):
         self.head_node_instance_id = head_node_instance_id
+        self.during_update = during_update