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