rust_tools: Provide a big hammer to do everything to uprev Rust

Create a new subcommand that can call both `create` and `remove`,
as well as preparing the repo and uploading the CLs.

Also update the steps to match the latest changes in UPGRADE.md.

BUG=chromium:1112551
TEST=unittest; create an example CL

Change-Id: I225d07d3e765daabd6ce8fc29309a5f11ef9cbae
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2355193
Commit-Queue: Tiancong Wang <tcwang@google.com>
Tested-by: Tiancong Wang <tcwang@google.com>
Reviewed-by: George Burgess <gbiv@chromium.org>

2 files changed, 271 insertions, 48 deletions

diff --git a/rust_tools/rust_uprev.py b/rust_tools/rust_uprev.py
index 50f85eb9..4ff92109 100755
--- a/rust_tools/rust_uprev.py
+++ b/rust_tools/rust_uprev.py
@@ -26,6 +26,10 @@ Replace `create --rust_version 1.45.0` with `remove --rust_version 1.43.0`
 if you want to remove all 1.43.0 related stuff in the same CL. Remember to
 use a different state file if you choose to run different subcommands.
 
+If you want a hammer that can do everything for you, use the subcommand
+`roll`. It can create a Rust uprev CL with `create` and `remove` and upload
+the CL to chromium code review.
+
 See `--help` for all available options.
 """
 
@@ -43,12 +47,13 @@ import sys
 import tempfile
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, T, Tuple
 
-from llvm_tools import chroot
+from llvm_tools import chroot, git
 RUST_PATH = '/mnt/host/source/src/third_party/chromiumos-overlay/dev-lang/rust'
 
 
-def get_command_output(command: List[str]) -> str:
-  return subprocess.check_output(command, encoding='utf-8').strip()
+def get_command_output(command: List[str], *args, **kwargs) -> str:
+  return subprocess.check_output(
+      command, encoding='utf-8', *args, **kwargs).strip()
 
 
 class RustVersion(NamedTuple):
@@ -109,18 +114,8 @@ def parse_commandline_args() -> argparse.Namespace:
       help='Continue the steps from the state file',
   )
 
-  subparsers = parser.add_subparsers(dest='subparser_name')
-  subparser_names = []
-
-  create_parser = subparsers.add_parser('create')
-  subparser_names.append('create')
-  create_parser.add_argument(
-      '--rust_version',
-      type=RustVersion.parse,
-      required=True,
-      help='Rust version to upgrade to, in the form a.b.c',
-  )
-  create_parser.add_argument(
+  create_parser_template = argparse.ArgumentParser(add_help=False)
+  create_parser_template.add_argument(
       '--template',
       type=RustVersion.parse,
       default=None,
@@ -128,19 +123,70 @@ def parse_commandline_args() -> argparse.Namespace:
       'a.b.c The ebuild has to exist in the chroot. If not specified, the '
       'tool will use the current Rust version in the chroot as template.',
   )
-  create_parser.add_argument(
+  create_parser_template.add_argument(
       '--skip_compile',
       action='store_true',
       help='Skip compiling rust to test the tool. Only for testing',
   )
 
+  subparsers = parser.add_subparsers(dest='subparser_name')
+  subparser_names = []
+  subparser_names.append('create')
+  create_parser = subparsers.add_parser(
+      'create',
+      parents=[create_parser_template],
+      help='Create changes uprevs Rust to a new version',
+  )
+  create_parser.add_argument(
+      '--rust_version',
+      type=RustVersion.parse,
+      required=True,
+      help='Rust version to uprev to, in the form a.b.c',
+  )
+
   subparser_names.append('remove')
-  remove_parser = subparsers.add_parser('remove')
+  remove_parser = subparsers.add_parser(
+      'remove',
+      help='Clean up old Rust version from chroot',
+  )
   remove_parser.add_argument(
       '--rust_version',
       type=RustVersion.parse,
+      default=None,
+      help='Rust version to remove, in the form a.b.c If not '
+      'specified, the tool will remove the oldest version in the chroot',
+  )
+
+  subparser_names.append('roll')
+  roll_parser = subparsers.add_parser(
+      'roll',
+      parents=[create_parser_template],
+      help='A command can create and upload a Rust uprev CL, including '
+      'preparing the repo, creating new Rust uprev, deleting old uprev, '
+      'and upload a CL to crrev.',
+  )
+  roll_parser.add_argument(
+      '--uprev',
+      type=RustVersion.parse,
       required=True,
-      help='Rust version to upgrade to, in the form a.b.c',
+      help='Rust version to uprev to, in the form a.b.c',
+  )
+  roll_parser.add_argument(
+      '--remove',
+      type=RustVersion.parse,
+      default=None,
+      help='Rust version to remove, in the form a.b.c If not '
+      'specified, the tool will remove the oldest version in the chroot',
+  )
+  roll_parser.add_argument(
+      '--skip_cross_compiler',
+      action='store_true',
+      help='Skip updating cross-compiler in the chroot',
+  )
+  roll_parser.add_argument(
+      '--no_upload',
+      action='store_true',
+      help='If specified, the tool will not upload the CL for review',
   )
 
   args = parser.parse_args()
@@ -341,12 +387,34 @@ def upload_to_localmirror(tempdir: str, rust_version: RustVersion) -> None:
   logging.info('Downloading Rust from %s', rust_src)
   gsutil_location = f'gs://chromeos-localmirror/distfiles/{tarfile_name}'
 
-  local_file = os.path.join(tempdir, tarfile_name)
-  subprocess.check_call(['curl', '-f', '-o', local_file, rust_src])
+  # Download Rust's source
+  rust_file = os.path.join(tempdir, tarfile_name)
+  subprocess.check_call(['curl', '-f', '-o', rust_file, rust_src])
+
+  # Verify the signature of the source
+  sig_file = os.path.join(tempdir, 'rustc_sig.asc')
+  subprocess.check_call(['curl', '-f', '-o', sig_file, f'{rust_src}.asc'])
+  try:
+    subprocess.check_output(['gpg', '--verify', sig_file, rust_file],
+                            encoding='utf-8',
+                            stderr=subprocess.STDOUT)
+  except subprocess.CalledProcessError as e:
+    if "gpg: Can't check signature" not in e.output:
+      raise RuntimeError(f'Failed to execute `gpg --verify`, {e.output}')
+
+    # If it fails to verify the signature, try import rustc key, and retry.
+    keys = get_command_output(
+        ['curl', '-f', 'https://keybase.io/rust/pgp_keys.asc'])
+    subprocess.run(['gpg', '--import'],
+                   input=keys,
+                   encoding='utf-8',
+                   check=True)
+    subprocess.check_call(['gpg', '--verify', sig_file, rust_file])
+
   # Since we are using `-n` to skip an item if it already exists, there's no
   # need to check if the file exists on GS bucket or not.
   subprocess.check_call(
-      ['gsutil', 'cp', '-n', '-a', 'public-read', local_file, gsutil_location])
+      ['gsutil', 'cp', '-n', '-a', 'public-read', rust_file, gsutil_location])
 
 
 def perform_step(state_file: pathlib.Path,
@@ -378,10 +446,7 @@ def perform_step(state_file: pathlib.Path,
 
 def create_rust_uprev(rust_version: RustVersion,
                       template: Optional[RustVersion], skip_compile: bool,
-                      run_step: Callable[[
-                          str, Callable[[], T], Optional[Callable[[Any], T]],
-                          Optional[Callable[[T], Any]]
-                      ], T]) -> None:
+                      run_step: Callable[[], T]) -> None:
   stage0_info = run_step(
       'parse stage0 file', lambda: parse_stage0_file(rust_version))
   template_version = run_step(
@@ -409,29 +474,82 @@ def create_rust_uprev(rust_version: RustVersion,
       template_version, rust_version))
 
 
+def find_oldest_rust_version_inchroot() -> RustVersion:
+  rust_versions = [
+      RustVersion.parse(x) for x in os.listdir(RUST_PATH) if '.ebuild' in x
+  ]
+
+  if len(rust_versions) <= 1:
+    raise RuntimeError('Expect to find more than one Rust versions')
+  return min(rust_versions)
+
+
 def remove_files(filename: str, path: str) -> None:
   subprocess.check_call(['git', 'rm', filename], cwd=path)
 
 
-def remove_rust_uprev(rust_version: RustVersion, run_step: Callable[[
-    str, Callable[[], T], Optional[Callable[[Any], T]], Optional[
-        Callable[[T], Any]]
-], T]) -> None:
+def remove_rust_uprev(rust_version: Optional[RustVersion],
+                      run_step: Callable[[], T]) -> None:
+  delete_version = run_step(
+      'find rust version to delete',
+      lambda: rust_version or find_oldest_rust_version_inchroot(),
+      result_from_json=prepare_uprev_from_json,
+  )
   run_step(
       'remove patches', lambda: remove_files(
-          f'files/rust-{rust_version}-*.patch', RUST_PATH))
-  run_step('remove ebuild', lambda: remove_files(f'rust-{rust_version}.ebuild',
-                                                 RUST_PATH))
+          f'files/rust-{delete_version}-*.patch', RUST_PATH))
+  run_step('remove ebuild', lambda: remove_files(
+      f'rust-{delete_version}.ebuild', RUST_PATH))
   ebuild_file = get_command_output(['equery', 'w', 'rust'])
   run_step('update manifest', lambda: update_manifest(ebuild_file))
   run_step('remove version from rust packages', lambda: update_rust_packages(
-      rust_version, add=False))
+      delete_version, add=False))
   run_step(
       'remove virtual/rust', lambda: remove_files(
-          f'rust-{rust_version}.ebuild',
+          f'rust-{delete_version}.ebuild',
           os.path.join(RUST_PATH, '../../virtual/rust')))
 
 
+def create_new_repo(rust_version: RustVersion) -> None:
+  output = get_command_output(['git', 'status', '--porcelain'], cwd=RUST_PATH)
+  if output:
+    raise RuntimeError(
+        f'{RUST_PATH} has uncommitted changes, please either discard them '
+        'or commit them.')
+  git.CreateBranch(RUST_PATH, f'rust-to-{rust_version}')
+
+
+def build_cross_compiler() -> None:
+  # Get target triples in ebuild
+  rust_ebuild = get_command_output(['equery', 'w', 'rust'])
+  with open(rust_ebuild, encoding='utf-8') as f:
+    contents = f.read()
+
+  target_triples_re = re.compile(r'RUSTC_TARGET_TRIPLES=\(([^)]+)\)')
+  m = target_triples_re.search(contents)
+  assert m, 'RUST_TARGET_TRIPLES not found in rust ebuild'
+  target_triples = m.group(1).strip().split('\n')
+  for target in target_triples:
+    if 'cros-' not in target:
+      continue
+    target = target.strip()
+    logging.info('Emerging cross compiler %s', target)
+    subprocess.check_call(['sudo', 'emerge', '-G', f'cross-{target}/gcc'])
+
+
+def create_new_commit(rust_version: RustVersion) -> None:
+  subprocess.check_call(['git', 'add', '-A'], cwd=RUST_PATH)
+  messages = [
+      f'[DO NOT SUBMIT] dev-lang/rust: upgrade to Rust {rust_version}',
+      '',
+      'This CL is created by rust_uprev tool automatically.'
+      '',
+      'BUG=None',
+      'TEST=Use CQ to test the new Rust version',
+  ]
+  git.UploadChanges(RUST_PATH, f'rust-to-{rust_version}', messages)
+
+
 def main() -> None:
   if not chroot.InChroot():
     raise RuntimeError('This script must be executed inside chroot')
@@ -461,8 +579,18 @@ def main() -> None:
   if args.subparser_name == 'create':
     create_rust_uprev(args.rust_version, args.template, args.skip_compile,
                       run_step)
-  else:
+  elif args.subparser_name == 'remove':
     remove_rust_uprev(args.rust_version, run_step)
+  else:
+    # If you have added more subparser_name, please also add the handlers above
+    assert args.subparser_name == 'roll'
+    run_step('create new repo', lambda: create_new_repo(args.uprev))
+    if not args.skip_cross_compiler:
+      run_step('build cross compiler', build_cross_compiler)
+    create_rust_uprev(args.uprev, args.template, args.skip_compile, run_step)
+    remove_rust_uprev(args.remove, run_step)
+    if not args.no_upload:
+      run_step('create rust uprev CL', lambda: create_new_commit(args.uprev))
 
 
 if __name__ == '__main__':
diff --git a/rust_tools/rust_uprev_test.py b/rust_tools/rust_uprev_test.py
index e007b822..a28c551e 100755
--- a/rust_tools/rust_uprev_test.py
+++ b/rust_tools/rust_uprev_test.py
@@ -13,6 +13,8 @@ import subprocess
 import unittest
 from unittest import mock
 
+from llvm_tools import git
+
 import rust_uprev
 
 
@@ -223,6 +225,67 @@ class UpdateRustPackagesTests(unittest.TestCase):
         package_after)
 
 
+class UploadToLocalmirrorTests(unittest.TestCase):
+  """Tests for upload_to_localmirror"""
+
+  def setUp(self):
+    self.tempdir = '/tmp/any/dir'
+    self.new_version = rust_uprev.RustVersion(1, 3, 5)
+    self.tarfile_name = f'rustc-{self.new_version}-src.tar.gz'
+    self.rust_src = f'https://static.rust-lang.org/dist/{self.tarfile_name}'
+    self.gsurl = f'gs://chromeos-localmirror/distfiles/{self.tarfile_name}'
+    self.rust_file = os.path.join(self.tempdir, self.tarfile_name)
+    self.sig_file = os.path.join(self.tempdir, 'rustc_sig.asc')
+
+  @mock.patch.object(subprocess, 'check_call')
+  @mock.patch.object(subprocess, 'check_output')
+  @mock.patch.object(subprocess, 'run')
+  def test_pass_without_retry(self, mock_run, mock_output, mock_call):
+    rust_uprev.upload_to_localmirror(self.tempdir, self.new_version)
+    mock_output.assert_called_once_with(
+        ['gpg', '--verify', self.sig_file, self.rust_file],
+        encoding='utf-8',
+        stderr=subprocess.STDOUT)
+    mock_call.assert_has_calls([
+        mock.call(['curl', '-f', '-o', self.rust_file, self.rust_src]),
+        mock.call(['curl', '-f', '-o', self.sig_file, f'{self.rust_src}.asc']),
+        mock.call([
+            'gsutil', 'cp', '-n', '-a', 'public-read', self.rust_file,
+            self.gsurl
+        ])
+    ])
+    mock_run.assert_not_called()
+
+  @mock.patch.object(subprocess, 'check_call')
+  @mock.patch.object(subprocess, 'check_output')
+  @mock.patch.object(subprocess, 'run')
+  @mock.patch.object(rust_uprev, 'get_command_output')
+  def test_pass_with_retry(self, mock_output, mock_run, mock_check, mock_call):
+    mock_check.side_effect = subprocess.CalledProcessError(
+        returncode=2, cmd=None, output="gpg: Can't check signature")
+    mock_output.return_value = 'some_gpg_keys'
+    rust_uprev.upload_to_localmirror(self.tempdir, self.new_version)
+    mock_check.assert_called_once_with(
+        ['gpg', '--verify', self.sig_file, self.rust_file],
+        encoding='utf-8',
+        stderr=subprocess.STDOUT)
+    mock_output.assert_called_once_with(
+        ['curl', '-f', 'https://keybase.io/rust/pgp_keys.asc'])
+    mock_run.assert_called_once_with(['gpg', '--import'],
+                                     input='some_gpg_keys',
+                                     encoding='utf-8',
+                                     check=True)
+    mock_call.assert_has_calls([
+        mock.call(['curl', '-f', '-o', self.rust_file, self.rust_src]),
+        mock.call(['curl', '-f', '-o', self.sig_file, f'{self.rust_src}.asc']),
+        mock.call(['gpg', '--verify', self.sig_file, self.rust_file]),
+        mock.call([
+            'gsutil', 'cp', '-n', '-a', 'public-read', self.rust_file,
+            self.gsurl
+        ])
+    ])
+
+
 class RustUprevOtherStagesTests(unittest.TestCase):
   """Tests for other steps in rust_uprev"""
 
@@ -305,20 +368,52 @@ class RustUprevOtherStagesTests(unittest.TestCase):
         os.path.join(virtual_rust_dir, f'rust-{self.new_version}.ebuild'))
     mock_exists.assert_called_once_with(virtual_rust_dir)
 
+  @mock.patch.object(os, 'listdir')
+  def test_find_oldest_rust_version_inchroot_pass(self, mock_ls):
+    mock_ls.return_value = [
+        f'rust-{self.old_version}.ebuild',
+        f'rust-{self.current_version}.ebuild', f'rust-{self.new_version}.ebuild'
+    ]
+    actual = rust_uprev.find_oldest_rust_version_inchroot()
+    expected = self.old_version
+    self.assertEqual(expected, actual)
+
+  @mock.patch.object(os, 'listdir')
+  def test_find_oldest_rust_version_inchroot_fail_with_only_one_ebuild(
+      self, mock_ls):
+    mock_ls.return_value = [f'rust-{self.new_version}.ebuild']
+    with self.assertRaises(RuntimeError) as context:
+      rust_uprev.find_oldest_rust_version_inchroot()
+    self.assertEqual('Expect to find more than one Rust versions',
+                     str(context.exception))
+
+  @mock.patch.object(rust_uprev, 'get_command_output')
+  @mock.patch.object(git, 'CreateBranch')
+  def test_create_new_repo(self, mock_branch, mock_output):
+    mock_output.return_value = ''
+    rust_uprev.create_new_repo(self.new_version)
+    mock_branch.assert_called_once_with(rust_uprev.RUST_PATH,
+                                        f'rust-to-{self.new_version}')
+
+  @mock.patch.object(rust_uprev, 'get_command_output')
   @mock.patch.object(subprocess, 'check_call')
-  def test_upload_to_localmirror(self, mock_call):
-    tempdir = '/tmp/any/dir'
-    rust_uprev.upload_to_localmirror(tempdir, self.new_version)
-
-    tarfile_name = f'rustc-{self.new_version}-src.tar.gz'
-    rust_src = f'https://static.rust-lang.org/dist/{tarfile_name}'
-    gsurl = f'gs://chromeos-localmirror/distfiles/{tarfile_name}'
-    local_file = os.path.join(tempdir, tarfile_name)
-    mock_call.assert_has_calls([
-        mock.call(['curl', '-f', '-o', local_file, rust_src]),
-        mock.call(
-            ['gsutil', 'cp', '-n', '-a', 'public-read', local_file, gsurl])
-    ])
+  def test_build_cross_compiler(self, mock_call, mock_output):
+    mock_output.return_value = f'rust-{self.new_version}.ebuild'
+    cros_targets = [
+        'x86_64-cros-linux-gnu', 'armv7a-cros-linux-gnueabihf',
+        'aarch64-cros-linux-gnu'
+    ]
+    all_triples = ['x86_64-pc-linux-gnu'] + cros_targets
+    rust_ebuild = 'RUSTC_TARGET_TRIPLES=(' + '\n\t'.join(all_triples) + ')'
+    mock_open = mock.mock_open(read_data=rust_ebuild)
+    with mock.patch('builtins.open', mock_open):
+      rust_uprev.build_cross_compiler()
+
+    emerge_calls = [
+        mock.call(['sudo', 'emerge', '-G', f'cross-{x}/gcc'])
+        for x in cros_targets
+    ]
+    mock_call.assert_has_calls(emerge_calls)
 
 
 if __name__ == '__main__':