lib/remoteexec_util.py

# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Module to use remoteexec from builders."""

import getpass
import os

from chromite.lib import constants


class Remoteexec(object):
  """Interface to use remoteexec on bots."""

  def __init__(self, reclient_dir, reproxy_cfg_file):
    """Initializes a Remoteexec instance.

    Args:
      reclient_dir: Path to the re-client directory that contains
        the reproxy, bootstrap, rewrapper binaries.
      reproxy_cfg_file: Path to the config file for starting reproxy.
    """

    if not os.path.isdir(reclient_dir):
      raise ValueError(
          f'reclient_dir does not point to a directory: {reclient_dir}')

    self.reclient_dir = reclient_dir

    # TODO(crbug.com/1256966): `reproxy_cfg_file` is either the full path
    # of a config generated by recipes or the name of a committed file
    # in `chromite/sdk/reclient_cfgs. When other builders, other than the
    # initial informational builder, are added, they should all be generated
    # by recipes rather than committed.
    if not os.path.exists(reproxy_cfg_file):
      reproxy_cfg_file = os.path.join(
          constants.CHROMITE_DIR, 'sdk', 'reclient_cfgs', reproxy_cfg_file)
      if not os.path.exists(reproxy_cfg_file):
        raise ValueError(
            f'reproxy_cfg_file does not exist: {reproxy_cfg_file}')
    self.reproxy_cfg_file = reproxy_cfg_file

  def __eq__(self, other):
    if self.__class__ is other.__class__:
      return (self.reclient_dir == other.reclient_dir
              and self.reproxy_cfg_file == other.reproxy_cfg_file)
    return NotImplemented

  def GetChrootExtraEnv(self):
    """Extra env vars set to do remoteexec inside chroot."""
    # These paths should match the paths in chroot that the
    # reclient directory and reproxy config file get mapped to
    # in sdk_lib/enter_chroot.sh
    reclient_dir = os.path.join('/home', getpass.getuser(), 'reclient')
    reproxy_cfg_file = os.path.join('/home', getpass.getuser(),
                                    'reclient_cfgs', 'reproxy_chroot.cfg')
    result = {'RECLIENT_DIR': reclient_dir,
              'REPROXY_CFG': reproxy_cfg_file}

    return result