xref: /aosp_15_r20/external/autotest/server/cros/chrome_sideloader.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1# Copyright 2021 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import base64
6import logging
7import os
8import json
9import random
10import stat
11import string
12
13# Shell command to force unmount a mount point if it is mounted
14FORCED_UMOUNT_DIR_IF_MOUNTPOINT_CMD = (
15        'if mountpoint -q %(dir)s; then umount -l %(dir)s; fi')
16# Shell command to set exec and suid flags
17SET_MOUNT_FLAGS_CMD = 'mount -o remount,exec,suid %s'
18# Shell command to send SIGHUP to dbus daemon
19DBUS_RELOAD_COMMAND = 'killall -HUP dbus-daemon'
20
21
22def extract_from_image(host, image_name, dest_dir):
23    """
24    Extract contents of an image to a directory.
25
26    @param host: The DUT to execute the command on
27    @param image_name: Name of image
28    @param dest_dir: directory where contents of image will be placed.
29
30    """
31
32    if not host.path_exists('/var/lib/imageloader/%s' % image_name):
33        raise Exception('Image %s not found on host %s' % (image_name, host))
34
35    def gen_random_str(length):
36        """
37        Generate random string
38
39        @param length: Length of the string
40
41        @return random string of specified length
42
43        """
44        return ''.join(
45                [random.choice(string.hexdigits) for _ in range(length)])
46
47    image_mount_point = '/tmp/image_%s' % gen_random_str(8)
48
49    # Create directories from scratch
50    host.run(['rm', '-rf', dest_dir])
51    host.run(['mkdir', '-p', '--mode', '0755', dest_dir, image_mount_point])
52
53    try:
54        # Mount image and copy content to the destination directory
55        host.run([
56                'imageloader', '--mount',
57                '--mount_component=%s' % image_name,
58                '--mount_point=%s' % image_mount_point
59        ])
60
61        host.run(['cp', '-r', '%s/*' % image_mount_point, '%s/' % dest_dir])
62    except Exception as e:
63        raise Exception(
64                'Error extracting content from image %s on host %s ' %
65                (image_name, host), e)
66    finally:
67        # Unmount image and remove the temporary directory
68        host.run([
69                'imageloader', '--unmount',
70                '--mount_point=%s' % image_mount_point
71        ])
72        host.run(['rm', '-rf', image_mount_point])
73
74
75def _stop_chrome_if_necessary(host):
76    """
77    Stop chrome if it is running.
78
79    @param host: The DUT to execute the command on
80
81    @return True if chrome was stopped. False otherwise.
82
83    """
84    status = host.run_output('status ui')
85    if 'start' in status:
86        return host.run('stop ui', ignore_status=True).exit_status == 0
87
88    return False
89
90
91def _mount_chrome(host, chrome_dir, chrome_mount_point):
92    """
93    Mount chrome to a mount point
94
95    @param host: The DUT to execute the command on
96    @param chrome_dir: directory where the chrome binary and artifacts
97                       will be placed.
98    @param chrome_mount_point: Chrome mount point
99
100    """
101    chrome_stopped = _stop_chrome_if_necessary(host)
102    _umount_chrome(host, chrome_mount_point)
103
104    # Mount chrome to the desired chrome directory
105    # Upon restart, this version of chrome will be used instead.
106    host.run(['mount', '--rbind', chrome_dir, chrome_mount_point])
107
108    # Chrome needs partition to have exec and suid flags set
109    host.run(SET_MOUNT_FLAGS_CMD % chrome_mount_point)
110
111    # Send SIGHUP to dbus-daemon to tell it to reload its configs. This won't
112    # pick up major changes (bus type, logging, etc.), but all we care about is
113    # getting the latest policy from /opt/google/chrome/dbus so that Chrome will
114    # be authorized to take ownership of its service names.
115    host.run(DBUS_RELOAD_COMMAND, ignore_status=True)
116
117    if chrome_stopped:
118        host.run('start ui', ignore_status=True)
119
120
121def _umount_chrome(host, chrome_mount_point):
122    """
123    Unmount chrome
124
125    @param host: The DUT to execute the command on
126    @param chrome_mount_point: Chrome mount point
127
128    """
129    chrome_stopped = _stop_chrome_if_necessary(host)
130    # Unmount chrome. Upon restart, the default version of chrome
131    # under the root partition will be used.
132    try:
133        host.run(FORCED_UMOUNT_DIR_IF_MOUNTPOINT_CMD %
134                 {'dir': chrome_mount_point})
135    except Exception as e:
136        raise Exception('Exception during cleanup on host %s' % host, e)
137
138    if chrome_stopped:
139        host.run('start ui', ignore_status=True)
140
141
142def setup_host(host, chrome_dir, chrome_mount_point):
143    """
144    Perform setup on host.
145
146    Mount chrome to point to the version provisioned by TLS.
147    The provisioning mechanism of chrome from the chrome builder is
148    based on Lacros Tast Test on Skylab (go/lacros-tast-on-skylab).
149
150    The lacros image provisioned by TLS contains the chrome binary
151    and artifacts.
152
153    @param host: The DUT to execute the command on
154    @param chrome_dir: directory where the chrome binary and artifacts
155                       will be placed.
156    @param chrome_mount_point: Chrome mount point
157
158    """
159    logging.info("Setting up host:%s", host)
160    try:
161        extract_from_image(host, 'lacros', chrome_dir)
162        if chrome_mount_point:
163            _mount_chrome(host, '%s/out/Release' % chrome_dir,
164                          chrome_mount_point)
165    except Exception as e:
166        raise Exception(
167                'Exception while mounting %s on host %s' %
168                (chrome_mount_point, host), e)
169
170
171def cleanup_host(host, chrome_dir, chrome_mount_point):
172    """
173    Umount chrome and perform cleanup.
174
175    @param host: The DUT to execute the command on
176    @param chrome_dir: directory where the chrome binary and artifacts
177                       is placed.
178    @param chrome_mount_point: Chrome mount point
179
180    """
181    logging.info("Unmounting chrome on host: %s", host)
182    try:
183        if chrome_mount_point:
184            _umount_chrome(host, chrome_mount_point)
185        host.run(['rm', '-rf', chrome_dir])
186    except Exception as e:
187        raise Exception('Exception during cleanup on host %s' % host, e)
188
189
190def get_tast_expr_from_file(host, args_dict, results_dir, base_path=None):
191    """
192    Get Tast expression from argument dictionary using a file.
193    If the tast_expr_file and tast_expr_key are in the dictionary returns the
194    tast expression from the file. If either/both args are not in the dict,
195    None is returned.
196    tast_expr_file expects a file containing a json dictionary which it will
197    then use tast_expr_key to pull the tast_expr.
198
199    The tast_expr_file is a json file containing a dictionary of names to tast
200    expressions like:
201
202    {
203    "default": "(\"group:mainline\" && \"dep:lacros\" && !informational)",
204    "tast_disabled_tests_from_lacros_example": "(\"group:mainline\" && \"dep:lacros\" && !informational && !\"name:lacros.Basic\")"
205    }
206
207    @param host: Host having the provisioned lacros image with the file
208    @param args_dict: Argument dictionary
209    @param results_dir: Where to store the tast_expr_file from the dut
210    @param base_path: Base path of the provisioned folder
211
212    """
213    tast_expr_file_name = args_dict.get('tast_expr_file')
214    tast_expr_key = args_dict.get('tast_expr_key')
215    if tast_expr_file_name and tast_expr_key:
216        if base_path:
217            tast_expr_file_name = os.path.join(base_path, tast_expr_file_name)
218
219        # Get the tast expr file from the provisioned lacros folder
220        if not host.path_exists(tast_expr_file_name):
221            raise Exception(
222                    'tast_expr_file: %s could not be found on the dut' %
223                    tast_expr_file_name)
224        local_file_name = os.path.join(results_dir,
225                                       os.path.basename(tast_expr_file_name))
226        st = os.stat(results_dir)
227        os.chmod(results_dir, st.st_mode | stat.S_IWRITE)
228        host.get_file(tast_expr_file_name, local_file_name, delete_dest=True)
229
230        with open(local_file_name) as tast_expr_file:
231            expr_dict = json.load(tast_expr_file)
232            expr = expr_dict.get(tast_expr_key)
233            # If both args were provided, the entry is expected in the file
234            if not expr:
235                raise Exception('tast_expr_key: %s could not be found' %
236                                tast_expr_key)
237            logging.info("tast_expr retreived from:%s", tast_expr_file)
238            return expr
239    return None
240
241
242def get_tast_expr(args_dict):
243    """
244    Get Tast expression from argument dictionary.
245    Users have options of using tast_expr or tast_expr_b64 in dictionary.
246    tast_expr_b64 expects a base64 encoded tast_expr, for instance:
247      tast_expr = '("group:mainline" && "dep:lacros")'
248      tast_expr_b64 = base64.b64encode(s.encode('utf-8')).decode('ascii')
249
250    @param args_dict: Argument dictionary
251
252    """
253    expr = args_dict.get('tast_expr')
254    if expr:
255        return expr
256
257    expr_b64 = args_dict.get('tast_expr_b64')
258    if expr_b64:
259        try:
260            expr = base64.b64decode(expr_b64).decode()
261            return expr
262        except Exception as e:
263            raise Exception('Failed to decode tast_expr_b64: %s' %
264                            expr_b64) from e
265
266    raise Exception(
267            '''Tast expression is unspecified: set tast_expr or tast_expr_b64 in --args.\n'''
268            '''  Example: test_that --args="tast_expr=lacros.Basic"\n'''
269            '''  If the expression contains spaces, consider transforming it to\n'''
270            '''  base64 and passing it via tast_expr_b64 flag.\n'''
271            '''  Example:\n'''
272            '''    In Python:\n'''
273            '''      tast_expr = '("group:mainline" && "dep:lacros")'\n'''
274            '''      # Yields 'KCJncm91cDptYWlubGluZSIgJiYgImRlcDpsYWNyb3MiKQ=='\n'''
275            '''      tast_expr_b64 = base64.b64encode(s.encode('utf-8')).decode('ascii')\n'''
276            '''    Then in Autotest CLI:\n'''
277            '''      test_that --args="tast_expr_b64=KCJncm91cDptYWlubGluZSIgJiYgImRlcDpsYWNyb3MiKQ=="\n'''
278            '''  More details at go/lacros-on-skylab.''')
279