From 3d5814b56ffcdfe51027f55819b6611e544b7596 Mon Sep 17 00:00:00 2001 From: Aleksandr Sidoruk Date: Thu, 28 Dec 2023 18:15:24 +0300 Subject: [PATCH] Allow to set permissions per ami in cluster.networks.config --- scripts/autoscaling/aws/nodeup.py | 34 ++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/scripts/autoscaling/aws/nodeup.py b/scripts/autoscaling/aws/nodeup.py index 9d493ca4da..aa6e912e72 100644 --- a/scripts/autoscaling/aws/nodeup.py +++ b/scripts/autoscaling/aws/nodeup.py @@ -146,6 +146,14 @@ def get_preference(preference_name): pipe_log('An error occured while getting preference {}, empty value is going to be used'.format(preference_name)) return None +def get_run_info(run_id): + pipe_api = PipelineAPI(api_url, None) + try: + return pipe_api.load_run_efficiently(run_id) + except: + pipe_log('An error occured while getting info for run id {}'.format(run_id)) + return None + def load_cloud_config(): global __CLOUD_METADATA__ global __CLOUD_TAGS__ @@ -238,7 +246,7 @@ def get_security_groups(aws_region, security_groups): def get_well_known_hosts(aws_region): return get_cloud_config_section(aws_region, "well_known_hosts") -def get_allowed_instance_image(cloud_region, instance_type, default_image, api_token): +def get_allowed_instance_image(cloud_region, instance_type, default_image, api_token, run_id): default_init_script = os.path.dirname(os.path.abspath(__file__)) + '/init.sh' default_embedded_scripts = None default_object = { "instance_mask_ami": default_image, "instance_mask": None, "init_script": default_init_script, @@ -263,6 +271,20 @@ def get_allowed_instance_image(cloud_region, instance_type, default_image, api_t # If something is wrong with the permissions check - do not use a restricted image continue + docker_image_list = image_config["docker_image"] if "docker_image" in image_config else None + try: + if docker_image_list: + pipe_log('Image config with restricted docker image found ({}), checking for match with a current run'.format(docker_image_list)) + run_info = get_run_info(run_id) + if not run_info or \ + not 'dockerImage' in run_info or \ + not run_info['dockerImage'] in docker_image_list: + continue + except: + # If something is wrong with the permissions check - do not use a restricted image + continue + + instance_mask = image_config["instance_mask"] instance_mask_ami = image_config["ami"] init_script = image_config.get("init_script", default_object["init_script"]) @@ -414,7 +436,8 @@ def run_instance(api_url, api_token, bid_price, ec2, aws_region, ins_hdd, kms_en swap_size = get_swap_size(aws_region, ins_type, is_spot) user_data_script = get_user_data_script(api_url, api_token, aws_region, ins_type, ins_img, kube_ip, kubeadm_token, - global_distribution_url, swap_size, pre_pull_images, node_ssh_port) + global_distribution_url, swap_size, pre_pull_images, node_ssh_port, + run_id) if is_spot: ins_id, ins_ip = find_spot_instance(ec2, aws_region, bid_price, run_id, pool_id, ins_img, ins_type, ins_key, ins_hdd, kms_encyr_key_id, user_data_script, num_rep, time_rep, swap_size, kube_client, instance_additional_spec, availability_zone, security_groups, subnet, network_interface, is_dedicated, performance_network) @@ -769,8 +792,9 @@ def replace_docker_images(pre_pull_images, user_data_script): def get_user_data_script(api_url, api_token, aws_region, ins_type, ins_img, kube_ip, kubeadm_token, - global_distribution_url, swap_size, pre_pull_images, node_ssh_port): - allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token) + global_distribution_url, swap_size, pre_pull_images, node_ssh_port, + run_id): + allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token, run_id) if allowed_instance and allowed_instance["init_script"]: init_script = open(allowed_instance["init_script"], 'r') user_data_script = init_script.read() @@ -1516,7 +1540,7 @@ def main(): instance_additional_spec = None if not ins_img or ins_img == 'null': # Redefine default instance image if cloud metadata has specific rules for instance type - allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token) + allowed_instance = get_allowed_instance_image(aws_region, ins_type, ins_img, api_token, run_id) if allowed_instance and allowed_instance["instance_mask"]: pipe_log('Found matching rule {instance_mask}/{ami} for requested instance type {instance_type}\nImage {ami} will be used'.format(instance_mask=allowed_instance["instance_mask"], ami=allowed_instance["instance_mask_ami"],