blob: 93ac4bc3e3fa5da1bd2fbc460deef89266d9f0ed [file] [log] [blame] [edit]
# 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.
"""Field handler classes.
The field handlers are meant to parse information from or do some other generic
action for a specific field type for the build_api script.
"""
import contextlib
import functools
import logging
import os
import shutil
from typing import Iterator, List, Optional
from chromite.third_party.google.protobuf import message as protobuf_message
from chromite.api.controller import controller_util
from chromite.api.gen.chromiumos import common_pb2
from chromite.lib import chroot_lib
from chromite.lib import osutils
class Error(Exception):
"""Base error class for the module."""
class InvalidResultPathError(Error):
"""Result path is invalid."""
class MissingChrootMessage(Error):
"""Message is missing Chroot field."""
class InvalidPathHandlerError(Error):
"""PathHandler params are invalid."""
class ChrootHandler:
"""Translate a Chroot message to chroot enter arguments and env."""
def __init__(self, clear_field) -> None:
self.clear_field = clear_field
def handle(self, message, recurse=True) -> Optional["chroot_lib.Chroot"]:
"""Parse a message for a chroot field."""
# Find the Chroot field. Search for the field by type to prevent it
# being tied to a naming convention.
for descriptor in message.DESCRIPTOR.fields:
field = getattr(message, descriptor.name)
if isinstance(field, common_pb2.Chroot):
chroot = field
if self.clear_field:
message.ClearField(descriptor.name)
return self.parse_chroot(chroot)
# Recurse down one level. This is handy for meta-endpoints that use
# another endpoint's request to produce data for or about the second
# endpoint. e.g. PackageService/NeedsChromeSource.
if recurse:
for descriptor in message.DESCRIPTOR.fields:
field = getattr(message, descriptor.name)
if isinstance(field, protobuf_message.Message):
chroot = self.handle(field, recurse=False)
if chroot:
return chroot
# Complain loudly if a message is used without a Chroot field.
raise MissingChrootMessage("No chroot message found.")
def parse_chroot(
self, chroot_message: common_pb2.Chroot
) -> "chroot_lib.Chroot":
"""Parse a Chroot message instance."""
return controller_util.ParseChroot(chroot_message)
def handle_chroot(
message: protobuf_message.Message, clear_field: bool = True
) -> "chroot_lib.Chroot":
"""Find and parse the chroot field, returning the Chroot instance."""
handler = ChrootHandler(clear_field)
return handler.handle(message)
class PathHandler:
"""Handles transferring a file or directory into or out of the chroot."""
INSIDE = common_pb2.Path.INSIDE
OUTSIDE = common_pb2.Path.OUTSIDE
def __init__(
self,
field: common_pb2.Path,
destination: str,
delete: bool,
chroot: Optional[chroot_lib.Chroot] = None,
reset: Optional[bool] = True,
) -> None:
"""Path handler initialization.
Args:
field: The Path message.
destination: The destination base path. If not set, paths are
translated only.
delete: Whether the copied file(s) should be deleted on cleanup.
chroot: Chroot object to use for translating the paths in/out of
the chroot as necessary -- modifying the destination path when
moving files into the chroot, or modifying the source path when
moving files outside.
reset: Whether to reset the state on cleanup.
"""
assert isinstance(field, common_pb2.Path)
assert field.path
assert field.location
if delete and not destination:
raise InvalidPathHandlerError(
"`delete` cannot be set with no destination."
)
self.field = field
self.destination = destination
self.chroot = chroot
self.delete = delete
self.tempdir = None
self.reset = reset
# For resetting the state.
self._transferred = False
self._original_message = common_pb2.Path()
self._original_message.CopyFrom(self.field)
def transfer(self, direction: int) -> None:
"""Copy the file or directory to its destination.
Args:
direction: The direction files are being copied (into or out of the
chroot). Specifying the direction allows avoiding performing
unnecessary copies.
"""
if self._transferred:
return
assert direction in [self.INSIDE, self.OUTSIDE]
if self.field.location == direction:
# Already in the correct location, nothing to do.
return
# Create a tempdir for the copied file if we're cleaning it up
# afterwards.
if self.delete:
self.tempdir = osutils.TempDir(base_dir=self.destination)
destination = self.tempdir.tempdir
else:
destination = self.destination
source = self.field.path
if direction == self.OUTSIDE and self.chroot:
source = self.chroot.full_path(source)
if not destination:
# Handle TRANSFER_TRANSLATE. Either `Chroot.full_path` or
# `Chroot.chroot_path` will do the actual translation.
dest_path = source
else:
if os.path.isfile(source):
# File - use old file name, just copy it into the destination.
dest_path = os.path.join(destination, os.path.basename(source))
copy_fn = shutil.copy
else:
# Directory - just copy everything into the new location.
dest_path = destination
copy_fn = functools.partial(
osutils.CopyDirContents, allow_nonempty=True
)
logging.debug("Copying %s to %s", source, dest_path)
copy_fn(source, dest_path)
# Clean up the destination path for returning, if applicable.
return_path = dest_path
if direction == self.INSIDE and self.chroot:
return_path = self.chroot.chroot_path(return_path)
if not destination:
logging.debug("Translated %s to %s", self.field.path, return_path)
self.field.path = return_path
self.field.location = direction
self._transferred = True
def cleanup(self) -> None:
"""Post-execution cleanup."""
if self.tempdir:
self.tempdir.Cleanup()
self.tempdir = None
if self.reset:
self.field.CopyFrom(self._original_message)
class SyncedDirHandler:
"""Handler for syncing directories across the chroot boundary."""
def __init__(
self,
field: common_pb2.SyncedDir,
destination: str,
chroot: chroot_lib.Chroot,
) -> None:
self.field = field
self.chroot = chroot
self.source = self.field.dir
if not self.source.endswith(os.sep):
self.source += os.sep
self.destination = destination
if not self.destination.endswith(os.sep):
self.destination += os.sep
# For resetting the message later.
self._original_message = common_pb2.SyncedDir()
self._original_message.CopyFrom(self.field)
def _sync(self, src, dest) -> None:
logging.info("Syncing %s to %s", src, dest)
# TODO: This would probably be more efficient with rsync.
osutils.EmptyDir(dest)
osutils.CopyDirContents(src, dest)
def sync_in(self) -> None:
"""Sync files from the source directory to the destination directory."""
self._sync(self.source, self.destination)
self.field.dir = self.chroot.chroot_path(self.destination)
def sync_out(self) -> None:
"""Sync files from the destination directory to the source directory."""
self._sync(self.destination, self.source)
self.field.CopyFrom(self._original_message)
@contextlib.contextmanager
def copy_paths_in(
message: protobuf_message.Message,
destination: str,
delete: Optional[bool] = True,
chroot: Optional[chroot_lib.Chroot] = None,
) -> Iterator[List[PathHandler]]:
"""Context manager function to transfer and cleanup all Path messages.
Args:
message: A message whose Path messages should be transferred.
destination: The base destination path.
delete: Whether the file(s) should be deleted.
chroot: Chroot object to use for translating the final destination path
into the chroot.
Yields:
list[PathHandler]: The path handlers.
"""
assert destination
handlers = _extract_handlers(
message, destination, chroot, delete=delete, reset=True
)
for handler in handlers:
handler.transfer(PathHandler.INSIDE)
try:
yield handlers
finally:
for handler in handlers:
handler.cleanup()
@contextlib.contextmanager
def sync_dirs(
message: protobuf_message.Message,
destination: str,
chroot: chroot_lib.Chroot,
) -> Iterator[SyncedDirHandler]:
"""Context manager function to handle SyncedDir messages.
The sync semantics are effectively:
rsync -r --del source/ destination/
* The endpoint runs. *
rsync -r --del destination/ source/
Args:
message: A message whose SyncedPath messages should be synced.
destination: The destination path.
chroot: Chroot object to use for translating the final destination path
into the chroot.
Yields:
The handlers.
"""
assert destination
handlers = _extract_handlers(
message,
destination,
chroot=chroot,
delete=False,
reset=True,
message_type=common_pb2.SyncedDir,
)
for handler in handlers:
handler.sync_in()
try:
yield handlers
finally:
for handler in handlers:
handler.sync_out()
def extract_results(
request_message: protobuf_message.Message,
response_message: protobuf_message.Message,
chroot: "chroot_lib.Chroot",
) -> None:
"""Transfer all response Path messages to the request's ResultPath.
Args:
request_message: The request message containing a ResultPath message.
response_message: The response message whose Path message(s) are to be
transferred.
chroot: The chroot the files are being copied out of.
"""
# Find the ResultPath.
for descriptor in request_message.DESCRIPTOR.fields:
field = getattr(request_message, descriptor.name)
if isinstance(field, common_pb2.ResultPath):
result_path_message = field
break
else:
# No ResultPath to handle.
return
destination = result_path_message.path.path
if result_path_message.transfer == common_pb2.ResultPath.TRANSFER_TRANSLATE:
if destination:
raise InvalidResultPathError(
"ResultPath.path must be empty for TRANSFER_TRANSLATE."
f" Value=`{destination}`."
)
elif not destination:
# ResultPath wasn't filled; don't copy to undefined location.
return
handlers = _extract_handlers(
response_message, destination, chroot=chroot, delete=False, reset=False
)
for handler in handlers:
handler.transfer(PathHandler.OUTSIDE)
handler.cleanup()
def _extract_handlers(
message,
destination,
chroot,
delete=False,
reset=False,
field_name=None,
message_type=None,
):
"""Recursive helper for handle_paths to extract Path messages."""
message_type = message_type or common_pb2.Path
is_path_target = message_type is common_pb2.Path
is_synced_target = message_type is common_pb2.SyncedDir
is_message = isinstance(message, protobuf_message.Message)
is_result_path = isinstance(message, common_pb2.ResultPath)
if not is_message or is_result_path:
# Base case: Nothing to handle.
# There's nothing we can do with scalar values.
# Skip ResultPath instances to avoid unnecessary file copying.
return []
elif is_path_target and isinstance(message, common_pb2.Path):
# Base case: Create handler for this message.
if not message.path or not message.location:
logging.debug("Skipping %s; incomplete.", field_name or "message")
return []
handler = PathHandler(
message, destination, delete=delete, chroot=chroot, reset=reset
)
return [handler]
elif is_synced_target and isinstance(message, common_pb2.SyncedDir):
if not message.dir or not destination:
logging.debug(
"Skipping %s; no directory given or missing destination.",
field_name or "message",
)
return []
handler = SyncedDirHandler(message, destination, chroot)
return [handler]
# Iterate through each field and recurse.
handlers = []
for descriptor in message.DESCRIPTOR.fields:
field = getattr(message, descriptor.name)
if field_name:
new_field_name = "%s.%s" % (field_name, descriptor.name)
else:
new_field_name = descriptor.name
if isinstance(field, protobuf_message.Message):
# Recurse for nested Paths.
handlers.extend(
_extract_handlers(
field,
destination,
chroot,
delete,
reset,
field_name=new_field_name,
message_type=message_type,
)
)
else:
# If it's iterable it may be a repeated field, try each element.
try:
iterator = iter(field)
except TypeError:
# Definitely not a repeated field, just move on.
continue
for element in iterator:
handlers.extend(
_extract_handlers(
element,
destination,
chroot,
delete,
reset,
field_name=new_field_name,
message_type=message_type,
)
)
return handlers