api/controller/payload.py

# Copyright 2019 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Payload API Service."""

from typing import TYPE_CHECKING

from chromite.api import controller
from chromite.api import faux
from chromite.api import validate
from chromite.api.gen.chromite.api import payload_pb2
from chromite.lib import cros_build_lib
from chromite.lib.paygen import paygen_payload_lib
from chromite.service import payload


if TYPE_CHECKING:
    from chromite.api import api_config

_VALID_IMAGE_PAIRS = (
    ("src_signed_image", "tgt_signed_image"),
    ("src_unsigned_image", "tgt_unsigned_image"),
    ("src_dlc_image", "tgt_dlc_image"),
    ("full_update", "tgt_unsigned_image"),
    ("full_update", "tgt_signed_image"),
    ("full_update", "tgt_dlc_image"),
)
_VALID_MINIOS_PAIRS = (
    ("src_signed_image", "tgt_signed_image"),
    ("src_unsigned_image", "tgt_unsigned_image"),
    ("full_update", "tgt_unsigned_image"),
    ("full_update", "tgt_signed_image"),
)

# TODO: Remove to use the standard cache directory if possible, otherwise
#  document why it cannot be used and preferably move outside of the repo.
_DEFAULT_PAYGEN_CACHE_DIR = ".paygen_cache"

# We have more fields we might validate however, they're either
# 'oneof' or allowed to be the empty value by design. If @validate
# gets more complex in the future we can add more here.
@faux.empty_success
@faux.empty_completed_unsuccessfully_error
@validate.require("bucket")
def GeneratePayload(
    input_proto: payload_pb2.GenerationRequest,
    output_proto: payload_pb2.GenerationResponse,
    config: "api_config.ApiConfig",
) -> int:
    """Generate a update payload ('do paygen').

    Args:
        input_proto: Input proto.
        output_proto: Output proto.
        config: The API call config.

    Returns:
        A controller return code (e.g. controller.RETURN_CODE_SUCCESS).
    """

    # Resolve the tgt image oneof.
    tgt_name = input_proto.WhichOneof("tgt_image_oneof")
    try:
        tgt_image = getattr(input_proto, tgt_name)
    except AttributeError:
        cros_build_lib.Die("%s is not a known tgt image type", tgt_name)

    # Resolve the src image oneof.
    src_name = input_proto.WhichOneof("src_image_oneof")

    # If the source image is 'full_update' we lack a source entirely.
    if src_name == "full_update":
        src_image = None
    # Otherwise we have an image.
    else:
        try:
            src_image = getattr(input_proto, src_name)
        except AttributeError:
            cros_build_lib.Die("%s is not a known src image type", src_name)

    # Ensure they are compatible oneofs.
    if (src_name, tgt_name) not in _VALID_IMAGE_PAIRS:
        cros_build_lib.Die(
            "%s and %s are not valid image pairs", src_image, tgt_image
        )

    # Ensure that miniOS payloads are only requested for compatible image types.
    if input_proto.minios and (src_name, tgt_name) not in _VALID_MINIOS_PAIRS:
        cros_build_lib.Die(
            "%s and %s are not valid image pairs for miniOS",
            src_image,
            tgt_image,
        )

    # Find the value of bucket or default to 'chromeos-releases'.
    destination_bucket = input_proto.bucket or "chromeos-releases"

    # There's a potential that some paygen_lib library might raise here, but since
    # we're still involved in config we'll keep it before the validate_only.
    payload_config = payload.PayloadConfig(
        tgt_image,
        src_image,
        destination_bucket,
        input_proto.minios,
        input_proto.verify,
        upload=not input_proto.dryrun,
        cache_dir=_DEFAULT_PAYGEN_CACHE_DIR,
    )

    # If configured for validation only we're done here.
    if config.validate_only:
        return controller.RETURN_CODE_VALID_INPUT

    # Do payload generation.
    local_path, remote_uri = "", ""
    try:
        local_path, remote_uri = payload_config.GeneratePayload()
    except paygen_payload_lib.PayloadGenerationSkippedException as e:
        # If paygen was skipped, provide a reason if possible.
        if isinstance(e, paygen_payload_lib.NoMiniOSPartitionException):
            reason = payload_pb2.GenerationResponse.NOT_MINIOS_COMPATIBLE
            output_proto.failure_reason = reason

    _SetGeneratePayloadOutputProto(output_proto, local_path, remote_uri)
    if remote_uri or input_proto.dryrun and local_path:
        return controller.RETURN_CODE_SUCCESS
    elif output_proto.failure_reason:
        return controller.RETURN_CODE_UNSUCCESSFUL_RESPONSE_AVAILABLE
    else:
        return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY


def _SetGeneratePayloadOutputProto(
    output_proto: payload_pb2.GenerationResponse,
    local_path: str,
    remote_uri: str,
):
    """Set the output proto with the results from the service class.

    Args:
        output_proto: The output proto.
        local_path: set output_proto with the local path, or ''.
        remote_uri: set output_proto with the remote uri, or ''.
    """
    output_proto.local_path = local_path or ""
    output_proto.remote_uri = remote_uri or ""