lambda_image.py

"""
Generates a Docker Image to be used for invoking a function locally
"""

import hashlib
import logging
import os
import platform
import re
import sys
import tempfile
import uuid
from enum import Enum
from pathlib import Path
from typing import Optional

import docker

from samcli.commands.local.cli_common.user_exceptions import (
    DockerDistributionAPIError,
    ImageBuildException,
)
from samcli.commands.local.lib.exceptions import InvalidIntermediateImageError
from samcli.lib.utils.architecture import has_runtime_multi_arch_image
from samcli.lib.utils.packagetype import IMAGE, ZIP
from samcli.lib.utils.stream_writer import StreamWriter
from samcli.lib.utils.tar import create_tarball
from samcli.local.common.file_lock import FileLock, cleanup_stale_locks
from samcli.local.docker.utils import (
    get_docker_platform,
    get_rapid_name,
    get_tar_filter_for_windows,
    get_validated_container_client,
)

LOG = logging.getLogger(__name__)

RAPID_IMAGE_TAG_PREFIX = "rapid"

TEST_RUNTIMES: list[str] = []


class Runtime(Enum):
    nodejs16x = "nodejs16.x"
    nodejs18x = "nodejs18.x"
    nodejs20x = "nodejs20.x"
    nodejs22x = "nodejs22.x"
    nodejs24x = "nodejs24.x"
    python38 = "python3.8"
    python39 = "python3.9"
    python310 = "python3.10"
    python311 = "python3.11"
    python312 = "python3.12"
    python313 = "python3.13"
    python314 = "python3.14"
    ruby32 = "ruby3.2"
    ruby33 = "ruby3.3"
    ruby34 = "ruby3.4"
    ruby40 = "ruby4.0"
    java8al2 = "java8.al2"
    java11 = "java11"
    java17 = "java17"
    java21 = "java21"
    java25 = "java25"
    go1x = "go1.x"
    dotnet6 = "dotnet6"
    dotnet8 = "dotnet8"
    dotnet10 = "dotnet10"
    provided = "provided"
    providedal2 = "provided.al2"
    providedal2023 = "provided.al2023"

    @classmethod
    def has_value(cls, value):
        """
        Checks if the enum has this value

        :param string value: Value to check
        :return bool: True, if enum has the value
        """
        return any(value == item.value for item in cls)

    @classmethod
    def get_image_name_tag(cls, runtime: str, architecture: str, is_preview: bool = False) -> str:
        """
        Returns the image name and tag for a particular runtime

        Parameters
        ----------
        runtime : str
            AWS Lambda runtime
        architecture : str
            Architecture for the runtime
        is_preview : bool
            Flag to use preview tag

        Returns
        -------
        str
            Image name and tag for the runtime's base image, like `python:3.12` or `provided:al2`
        """
        runtime_image_tag = ""
        if runtime == cls.provided.value:
            # There's a special tag for `provided` not al2 (provided:alami)
            runtime_image_tag = "provided:alami"
        elif runtime.startswith("provided"):
            # `provided.al2` becomes `provided:al2``
            runtime_image_tag = runtime.replace(".", ":")
        elif runtime.startswith("dotnet"):
            # dotnet6 becomes dotnet:6
            runtime_image_tag = runtime.replace("dotnet", "dotnet:")
        else:
            # This fits most runtimes format: `nameN.M` becomes `name:N.M` (python3.9 -> python:3.9)
            runtime_image_tag = re.sub(r"^([a-z]+)([0-9][a-z0-9\.]*)$", r"\1:\2", runtime)
            # nodejs20.x, go1.x, etc don't have the `.x` part.
            runtime_image_tag = runtime_image_tag.replace(".x", "")

        if is_preview:
            runtime_image_tag = f"{runtime_image_tag}-preview"

        # Runtime image tags contain the architecture only if more than one is supported for that runtime
        if has_runtime_multi_arch_image(runtime):
            runtime_image_tag = f"{runtime_image_tag}-{architecture}"
        return runtime_image_tag


class LambdaImage:
    _LAYERS_DIR = "/opt"
    _INVOKE_REPO_PREFIX = "public.ecr.aws/lambda"
    _SAM_INVOKE_REPO_PREFIX = "public.ecr.aws/sam/emulation"
    _SAM_CLI_REPO_NAME = "samcli/lambda"
    _RAPID_SOURCE_PATH = Path(__file__).parent.joinpath("..", "rapid").resolve()

    def __init__(self, layer_downloader, skip_pull_image, force_image_build, docker_client=None, invoke_images=None):
        """

        Parameters
        ----------
        layer_downloader samcli.local.layers.layer_downloader.LayerDownloader
            LayerDownloader to download layers locally
        skip_pull_image bool
            True if the image should not be pulled from DockerHub
        force_image_build bool
            True to download the layer and rebuild the image even if it exists already on the system
        docker_client docker.DockerClient
            Optional docker client object
        """
        self.layer_downloader = layer_downloader
        self.skip_pull_image = skip_pull_image
        self.force_image_build = force_image_build

        # Store docker_client parameter for lazy initialization
        # Only validate container runtime when Docker client is actually accessed
        self._docker_client_param = docker_client
        self._validated_docker_client = None
        self.invoke_images = invoke_images

        # Clean up old lock files on initialization
        cleanup_stale_locks(Path(tempfile.gettempdir()), "building")

    @property
    def docker_client(self):
        """
        Lazy initialization of Docker client. Only validates container runtime when actually accessed.
        This prevents unnecessary container runtime validation during object creation.
        """
        if self._validated_docker_client is None:
            self._validated_docker_client = self._docker_client_param or get_validated_container_client()
        return self._validated_docker_client

    def build(self, runtime, packagetype, image, layers, architecture, stream=None, function_name=None):
        """
        Build the image if one is not already on the system that matches the runtime and layers

        Parameters
        ----------
        runtime : str
            Name of the Lambda runtime
        packagetype : str
            Packagetype for the Lambda
        image : str
            Pre-defined invocation image.
        layers : list(samcli.commands.local.lib.provider.Layer)
            List of layers
        architecture : str
            Architecture type either x86_64 or arm64 on AWS lambda
        stream : io.RawIOBase
            stream to write
        function_name : str
            The name of the function that the image is building for

        Returns
        -------
        str
            The image to be used (REPOSITORY:TAG)
        """
        base_image = None
        tag_prefix = ""

        if packagetype == IMAGE:
            if self.invoke_images:
                base_image = self.invoke_images.get(function_name, self.invoke_images.get(None))
            if not base_image:
                base_image = image
        elif packagetype == ZIP:
            is_preview = runtime in TEST_RUNTIMES
            runtime_image_tag = Runtime.get_image_name_tag(runtime, architecture, is_preview=is_preview)
            if self.invoke_images:
                base_image = self.invoke_images.get(function_name, self.invoke_images.get(None))
            if not base_image:
                # Gets the ECR image format like `python:3.12` or `nodejs:16-x86_64`
                runtime_only_number = re.split("[:-]", runtime_image_tag)[1]
                tag_prefix = f"{runtime_only_number}-"
                base_image = f"{self._INVOKE_REPO_PREFIX}/{runtime_image_tag}"

                # Temporarily add a version tag to the emulation image so that we don't pull a broken image
                if platform.system().lower() == "windows" and runtime in [Runtime.go1x.value]:
                    LOG.info("Falling back to a previous version of the emulation image")
                    base_image = f"{base_image}.2023.08.02.10"

        if not base_image:
            raise InvalidIntermediateImageError(f"Invalid PackageType, PackageType needs to be one of [{ZIP}, {IMAGE}]")

        if image:
            self.skip_pull_image = True

        # If the image name had a digest, removing the @ so that a valid image name can be constructed
        # to use for the local invoke image name.
        image_repo = base_image.split(":")[0].replace("@", "")
        rapid_image = f"{image_repo}:{tag_prefix}{RAPID_IMAGE_TAG_PREFIX}-{architecture}"

        downloaded_layers = []

        if layers and packagetype == ZIP:
            downloaded_layers = self.layer_downloader.download_all(layers, self.force_image_build)

            docker_image_version = self._generate_docker_image_version(downloaded_layers, runtime_image_tag)
            rapid_image = f"{self._SAM_CLI_REPO_NAME}-{docker_image_version}"

        image_not_found = False

        # If we are not using layers, build anyways to ensure any updates to rapid get added
        try:
            self.docker_client.images.get(rapid_image)
            # Check if the base image is up-to-date locally and modify build/pull parameters accordingly
            self._check_base_image_is_current(base_image)
        except docker.errors.ImageNotFound:
            LOG.info("Local image was not found.")
            image_not_found = True
        except docker.errors.APIError as e:
            if e.__class__ is docker.errors.NotFound:
                # A generic "NotFound" is raised when we aren't able to check the image version
                # for example when the docker daemon's api doesn't support this action.
                #
                # See Also: https://github.com/containers/podman/issues/17726
                LOG.warning(
                    "Unknown 404 - Unable to check if base image is current.\n\nPossible incompatible "
                    "Docker engine clone employed. Consider `--skip-pull-image` for improved speed, the "
                    "tradeoff being not running the latest image."
                )
                image_not_found = True
            else:
                raise DockerDistributionAPIError(str(e)) from e

        # If building a new rapid image, delete older rapid images
        if image_not_found and rapid_image == f"{image_repo}:{tag_prefix}{RAPID_IMAGE_TAG_PREFIX}-{architecture}":
            if tag_prefix:
                # ZIP functions with new RAPID format. Delete images from the old ecr/sam repository
                self._remove_rapid_images(f"{self._SAM_INVOKE_REPO_PREFIX}-{runtime}")
            else:
                self._remove_rapid_images(image_repo)

        if (
            self.force_image_build
            or image_not_found
            or any(layer.is_defined_within_template for layer in downloaded_layers)
            or not runtime
        ):
            stream_writer = stream or StreamWriter(sys.stderr)

            # Use build lock to prevent concurrent builds of the same image
            # Use system temp directory for lock files instead of layer cache to avoid polluting cache
            build_lock = FileLock(Path(tempfile.gettempdir()), rapid_image, "building")

            if build_lock.acquire_lock():
                # We got the lock, proceed with build
                try:
                    stream_writer.write_str("Building image...")
                    stream_writer.flush()
                    self._build_image(
                        image if image else base_image,
                        rapid_image,
                        downloaded_layers,
                        architecture,
                        stream=stream_writer,
                    )
                    build_lock.release_lock(success=True)
                except Exception:
                    build_lock.release_lock(success=False)
                    raise
            else:
                # Another process is building, wait for it to complete
                stream_writer.write_str("Another process is building the same image, waiting...")
                stream_writer.flush()

                if build_lock.wait_for_operation():
                    # Build completed successfully by another process
                    stream_writer.write_str("Image build completed by another process.")
                    stream_writer.flush()

                    # Verify the image actually exists
                    try:
                        self.docker_client.images.get(rapid_image)
                    except docker.errors.ImageNotFound:
                        # Image doesn't exist, fallback to building ourselves
                        LOG.warning("Expected image not found after concurrent build, building ourselves")
                        stream_writer.write_str("Building image...")
                        stream_writer.flush()
                        self._build_image(
                            image if image else base_image,
                            rapid_image,
                            downloaded_layers,
                            architecture,
                            stream=stream_writer,
                        )
                else:
                    # Build failed or timed out, build ourselves
                    LOG.warning("Concurrent build failed or timed out, building ourselves")
                    stream_writer.write_str("Building image...")
                    stream_writer.flush()
                    self._build_image(
                        image if image else base_image,
                        rapid_image,
                        downloaded_layers,
                        architecture,
                        stream=stream_writer,
                    )

        return rapid_image

    def get_config(self, image_tag):
        config = {}
        try:
            image = self.docker_client.images.get(image_tag)
            return image.attrs.get("Config")
        except docker.errors.ImageNotFound:
            return config

    @staticmethod
    def _generate_docker_image_version(layers, runtime_image_tag):
        """
        Generate the Docker TAG that will be used to create the image

        Parameters
        ----------
        layers list(samcli.commands.local.lib.provider.Layer)
            List of the layers

        runtime_image_tag str
            Runtime version format to generate image name and tag (including architecture, e.g. "python:3.12-x86_64")

        Returns
        -------
        str
            String representing the TAG to be attached to the image
        """

        # Docker has a concept of a TAG on an image. This is plus the REPOSITORY is a way to determine
        # a version of the image. We will produced a TAG for a combination of the runtime with the layers
        # specified in the template. This will allow reuse of the runtime and layers across different
        # functions that are defined. If two functions use the same runtime with the same layers (in the
        # same order), SAM CLI will only produce one image and use this image across both functions for invoke.

        return (
            runtime_image_tag
            + "-"
            + hashlib.sha256("-".join([layer.name for layer in layers]).encode("utf-8")).hexdigest()[0:25]
        )

    def _build_image(self, base_image, docker_tag, layers, architecture, stream=None):
        """
        Builds the image

        Parameters
        ----------
        base_image str
            Base Image to use for the new image
        docker_tag str
            Docker tag (REPOSITORY:TAG) to use when building the image
        layers list(samcli.commands.local.lib.provider.Layer)
            List of Layers to be use to mount in the image
        architecture str
            Architecture, either x86_64 or arm64
        stream samcli.lib.utils.stream_writer.StreamWriter
            Stream to write the build output

        Raises
        ------
        samcli.commands.local.cli_common.user_exceptions.ImageBuildException
            When docker fails to build the image
        """
        dockerfile_content = self._generate_dockerfile(base_image, layers, architecture)

        # Create dockerfile in the same directory of the layer cache
        dockerfile_name = "dockerfile_" + str(uuid.uuid4())
        full_dockerfile_path = Path(self.layer_downloader.layer_cache, dockerfile_name)
        stream_writer = stream or StreamWriter(sys.stderr)

        try:
            with open(str(full_dockerfile_path), "w") as dockerfile:
                dockerfile.write(dockerfile_content)

            # add dockerfile and rapid source paths
            tar_paths = {
                str(full_dockerfile_path): "Dockerfile",
                self._RAPID_SOURCE_PATH: "/" + get_rapid_name(architecture),
            }

            for layer in layers:
                tar_paths[layer.codeuri] = "/" + layer.name

            # Use shared tar filter for Windows compatibility
            tar_filter = get_tar_filter_for_windows()

            with create_tarball(tar_paths, tar_filter=tar_filter, dereference=True) as tarballfile:
                try:
                    resp_stream = self.docker_client.api.build(
                        fileobj=tarballfile,
                        custom_context=True,
                        rm=True,
                        tag=docker_tag,
                        pull=not self.skip_pull_image,
                        decode=True,
                        platform=get_docker_platform(architecture),
                    )
                    # Track errors and check for Finch "already exists" pattern (minimal memory usage)
                    error_details = []
                    finch_already_exists_detected = False

                    for log in resp_stream:
                        if log is None:
                            continue

                        stream_writer.write_str(".")
                        stream_writer.flush()

                        if isinstance(log, dict):
                            # Check for errors
                            if "error" in log:
                                error_msg = log["error"]
                                error_details.append(error_msg)
                                LOG.error(f"Docker build error: {error_msg}")

                                # Check for Finch "already exists" pattern immediately
                                if "already exists" in error_msg.lower() and "image" in error_msg.lower():
                                    finch_already_exists_detected = True
                                    LOG.debug(f"Finch 'already exists' pattern detected in error: {error_msg}")

                            # Process stream messages
                            if "stream" in log:
                                stream_msg = log["stream"].strip()
                                if stream_msg:
                                    LOG.debug(f"Build stream: {stream_msg}")

                                    # Check for Finch "already exists" pattern in stream
                                    if "already exists" in stream_msg.lower() and "image" in stream_msg.lower():
                                        finch_already_exists_detected = True
                                        LOG.debug(f"Finch 'already exists' pattern detected in stream: {stream_msg}")

                            elif "status" in log:
                                status_msg = log["status"]
                                LOG.debug(f"Build status: {status_msg}")

                            elif "progress" in log:
                                LOG.debug(f"Build progress: {log['progress']}")
                        else:
                            LOG.debug(f"Non-dict build log: {log}")

                    stream_writer.write_str(os.linesep)

                    # Handle Finch "already exists" case
                    if finch_already_exists_detected:
                        LOG.info("Finch reported image already exists - treating as successful build")
                    # If we have errors that aren't Finch "already exists", raise an exception
                    elif error_details:
                        full_error = "; ".join(error_details)
                        LOG.error(f"Docker build failed with errors: {full_error}")
                        raise ImageBuildException(f"Error building docker image: {full_error}")
                except (docker.errors.BuildError, docker.errors.APIError) as ex:
                    stream_writer.write_str(os.linesep)
                    LOG.exception("Failed to build Docker Image")
                    raise ImageBuildException("Building Image failed.") from ex
        finally:
            if full_dockerfile_path.exists():
                full_dockerfile_path.unlink()

    @staticmethod
    def _generate_dockerfile(base_image, layers, architecture):
        """
        FROM public.ecr.aws/lambda/python:3.9-x86_64

        ADD aws-lambda-rie /var/rapid

        ADD layer1 /opt
        ADD layer2 /tmp/layer1
        RUN cp -rf /tmp/layer1/. /opt/ && rm -rf /tmp/layer1

        Parameters
        ----------
        base_image : str
            Base Image to use for the new image
        layers : list
            List of Layers to be use to mount in the image
        architecture : str
            Architecture type either x86_64 or arm64 on AWS lambda

        Returns
        -------
        str
            String representing the Dockerfile contents for the image
        """
        rie_name = get_rapid_name(architecture)
        rie_path = "/var/rapid/"
        dockerfile_content = (
            f"FROM {base_image}\n"
            + f"ADD {rie_name} {rie_path}\n"
            + f"RUN mv {rie_path}{rie_name} {rie_path}aws-lambda-rie && chmod +x {rie_path}aws-lambda-rie\n"
        )
        for idx, layer in enumerate(layers):
            if idx == 0:
                # First layer can be added directly — no existing files to conflict with
                dockerfile_content += f"ADD {layer.name} {LambdaImage._LAYERS_DIR}\n"
            else:
                # Subsequent layers: stage in temp dir, then cp -rf to ensure overwrite.
                # Docker ADD/COPY may merge directories without overwriting in BuildKit.
                stage_dir = f"/tmp/layer{idx}"
                dockerfile_content += f"ADD {layer.name} {stage_dir}\n"
                dockerfile_content += f"RUN cp -rf {stage_dir}/. {LambdaImage._LAYERS_DIR}/ && rm -rf {stage_dir}\n"
        return dockerfile_content

    def _remove_rapid_images(self, repo: str) -> None:
        """
        Remove all rapid images for given repo

        Parameters
        ----------
        repo string
            Repo for which rapid images will be removed
        """
        LOG.info("Removing rapid images for repo %s", repo)
        try:
            for image in self.docker_client.images.list(name=repo):
                for tag in image.tags:
                    if self.is_rapid_image(tag) and not self.is_rapid_image_current(tag):
                        try:
                            self.docker_client.images.remove(image.id, force=True)
                        except docker.errors.APIError as ex:
                            LOG.warning("Failed to remove rapid image with ID: %s", image.id, exc_info=ex)
                        break
        except docker.errors.APIError as ex:
            LOG.warning("Failed getting images from repo %s", repo, exc_info=ex)

    @staticmethod
    def is_rapid_image(image_name: str) -> bool:
        """
        Is the image tagged as a RAPID clone?

        Parameters
        ----------
        image_name : str
            Name of the image

        Returns
        -------
        bool
            True if the image tag starts with the rapid prefix or contains it in between. False, otherwise
        """

        try:
            tag = image_name.split(":")[1]
            return tag.startswith(f"{RAPID_IMAGE_TAG_PREFIX}-") or f"-{RAPID_IMAGE_TAG_PREFIX}-" in tag
        except (IndexError, AttributeError):
            # split() returned 1 or less items or image_name is None
            return False

    @staticmethod
    def is_rapid_image_current(image_name: str) -> bool:
        """
        Verify if an image has the latest format.
        The current format doesn't include the SAM version and has the RAPID prefix between dashes.

        Parameters
        ----------
        image_name : str
            name the image

        Returns
        -------
        bool
            return True if it is current and vice versa
        """
        return f"-{RAPID_IMAGE_TAG_PREFIX}-" in image_name

    def _check_base_image_is_current(self, image_name: str) -> None:
        """
        Check if the existing base image is up-to-date and update modifier parameters
        (skip_pull_image, force_image_build) accordingly, printing an informative
        message depending on the case.

        Parameters
        ----------
        image_name : str
            Base image name to check
        """

        # No need to check if the overriding parameters are already set
        if self.skip_pull_image or self.force_image_build:
            return

        if self.is_base_image_current(image_name):
            self.skip_pull_image = True
            LOG.info("Local image is up-to-date")
        else:
            self.force_image_build = True
            LOG.info(
                "Local image is out of date and will be updated to the latest runtime. "
                "To skip this, pass in the parameter --skip-pull-image"
            )

    def is_base_image_current(self, image_name: str) -> bool:
        """
        Return True if the base image is up-to-date with the remote environment by comparing the image digests

        Parameters
        ----------
        image_name : str
            Base image name to check

        Returns
        -------
        bool
            True if local image digest is the same as the remote image digest
        """
        return self.get_local_image_digest(image_name) == self.get_remote_image_digest(image_name)

    def get_remote_image_digest(self, image_name: str) -> Optional[str]:
        """
        Get the digest of the remote version of an image

        Parameters
        ----------
        image_name : str
            Name of the image to get the digest

        Returns
        -------
        str
            Image digest, including `sha256:` prefix
        """
        remote_info = self.docker_client.images.get_registry_data(image_name)
        digest: Optional[str] = remote_info.attrs.get("Descriptor", {}).get("digest")
        return digest

    def get_local_image_digest(self, image_name: str) -> Optional[str]:
        """
        Get the digest of the local version of an image

        Parameters
        ----------
        image_name : str
            Name of the image to get the digest

        Returns
        -------
        str
            Image digest, including `sha256:` prefix
        """
        image_info = self.docker_client.images.get(image_name)
        try:
            full_digest: str = image_info.attrs.get("RepoDigests", [None])[0]
            return full_digest.split("@")[1]
        except (AttributeError, IndexError):
            return None