1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.android.tradefed.targetprep;
17 
18 import com.android.tradefed.build.BuildInfoKey.BuildInfoFileKey;
19 import com.android.tradefed.build.IBuildInfo;
20 import com.android.tradefed.config.ConfigurationException;
21 import com.android.tradefed.config.IConfiguration;
22 import com.android.tradefed.config.IConfigurationReceiver;
23 import com.android.tradefed.config.Option;
24 import com.android.tradefed.config.OptionClass;
25 import com.android.tradefed.device.DeviceNotAvailableException;
26 import com.android.tradefed.device.ITestDevice;
27 import com.android.tradefed.invoker.TestInformation;
28 import com.android.tradefed.log.LogUtil.CLog;
29 import com.android.tradefed.result.error.InfraErrorIdentifier;
30 import com.android.tradefed.util.ArrayUtil;
31 import com.android.tradefed.util.CommandResult;
32 import com.android.tradefed.util.CommandStatus;
33 import com.android.tradefed.util.FileUtil;
34 import com.android.tradefed.util.RunUtil;
35 
36 import com.google.common.annotations.VisibleForTesting;
37 
38 import java.io.File;
39 import java.io.IOException;
40 import java.time.Duration;
41 import java.util.ArrayList;
42 import java.util.Arrays;
43 import java.util.Collection;
44 import java.util.List;
45 import java.util.regex.Matcher;
46 import java.util.regex.Pattern;
47 
48 /** A {@link ITargetPreparer} that allows to mix a kernel image with the device image. */
49 @OptionClass(alias = "mix-kernel-target-preparer")
50 public class MixKernelTargetPreparer extends BaseTargetPreparer
51         implements ILabPreparer, IConfigurationReceiver {
52 
53     @Option(
54             name = "device-file-label",
55             description = "The label for the test device that stores device images.")
56     private String mDeviceLabel = "device";
57 
58     @Option(
59             name = "kernel-file-label",
60             description =
61                     "The file key prefix that is attached to the kernel images, "
62                             + " for example {kernel}Image.gz.")
63     private String mKernelLabel = "kernel";
64 
65     @Option(
66             name = "gki-file-label",
67             description =
68                     "The file key prefix that is attached the GKI images, "
69                             + "for example {gki}Image.gz.")
70     private String mGkiLabel = "gki";
71 
72     /*
73     To use mix kernel tool on your local file system, use flag --mix-kernel-tool-path
74     For example: --mix-kernel-tool-path
75     /the_path_to_internal_android_tree/vendor/google/tools/build_mixed_kernels
76 
77     A file directory path can also be used. The mixing tool will search the mix-kernel-tool-name in
78     the specified file directory
79     For example the tool path can be specified with flexible download as:
80     --mix-kernel-tool-path ab://git_master/flame-userdebug/LATEST/.*flame-tests-.*.zip?unzip=true
81 
82     If mix-kernel-tool-path is not specified, testsdir of the device build will be used as the
83     mix-kernel-tool-path.
84     */
85     @Option(
86             name = "mix-kernel-tool-path",
87             description =
88                     "The file path of mix kernel tool. It can be the absolute path of the tool"
89                         + " path, or the absolute directory path that contains the tool. It can be"
90                         + " used with flexible download feature for example --mix-kernel-tool-path"
91                         + " ab://git_master/flame-userdebug/LATEST/.*flame-tests-.*.zip?unzip=true")
92     private File mMixKernelToolPath = null;
93 
94     @Option(
95             name = "mix-kernel-tool-name",
96             description =
97                     "The mixing kernel tool file name, defaulted to build_mixed_kernels. "
98                             + "If mix-kernel-tool-path is a directory, the mix-kernel-tool-name "
99                             + "will be used to locate the tool in the directory of "
100                             + "mix-kernel-tool-path.")
101     private String mMixKernelToolName = "build_mixed_kernels_ramdisk";
102 
103     @Option(
104             name = "mix-kernel-script-wait-time",
105             description = "The maximum wait time for mix kernel script. By default is 20 minutes. ")
106     private Duration mMixingWaitTime = Duration.ofMinutes(20);
107 
108     @Option(
109             name = "mix-kernel-arg",
110             description =
111                     "Additional arguments to be passed to mix-kernel-script command, "
112                             + "including leading dash, e.g. \"--nocompress\"")
113     private Collection<String> mMixKernelArgs = new ArrayList<>();
114 
115     private IConfiguration mConfig;
116 
117     @Override
setConfiguration(IConfiguration configuration)118     public void setConfiguration(IConfiguration configuration) {
119         mConfig = configuration;
120     }
121 
122     @Override
setUp(TestInformation testInfo)123     public void setUp(TestInformation testInfo)
124             throws TargetSetupError, BuildError, DeviceNotAvailableException {
125         ITestDevice device = testInfo.getDevice();
126         IBuildInfo buildInfo = testInfo.getBuildInfo();
127 
128         File tmpDeviceDir = null;
129         File tmpKernelDir = null;
130         File tmpGkiDir = null;
131         File tmpNewDeviceDir = null;
132         try {
133             // Find the kernel mixing tool
134             findMixKernelTool(buildInfo);
135 
136             // Create temp dir for original device build, kernel build, and new device build
137             tmpDeviceDir = FileUtil.createTempDir("device_dir");
138             tmpKernelDir = FileUtil.createTempDir("kernel_dir");
139             tmpNewDeviceDir = FileUtil.createTempDir("new_device_dir");
140             tmpGkiDir = FileUtil.createTempDir("gki_dir");
141 
142             for (String fileKey : buildInfo.getVersionedFileKeys()) {
143                 CLog.i("Processing file %s", fileKey);
144                 File srcFile = buildInfo.getFile(fileKey);
145                 if (fileKey.contains("{" + mKernelLabel + "}")) {
146                     // Copy kernel image to tmpKernelDir
147                     copyLabelFileToDir(fileKey, srcFile, tmpKernelDir);
148                 } else if (fileKey.contains("{" + mGkiLabel + "}")) {
149                     // Copy GKI image to tmpGkilDir
150                     copyLabelFileToDir(fileKey, srcFile, tmpGkiDir);
151                 } else if (fileKey.contains("{" + mDeviceLabel + "}")) {
152                     // Copy device image to tmpDeviceDir
153                     copyLabelFileToDir(fileKey, srcFile, tmpDeviceDir);
154                 } else if (BuildInfoFileKey.DEVICE_IMAGE.getFileKey().equals(fileKey)) {
155                     copyDeviceImageToDir(srcFile, tmpDeviceDir);
156                 }
157             }
158 
159             if (tmpDeviceDir.listFiles().length == 0) {
160                 throw new TargetSetupError(
161                         "Could not find device images", device.getDeviceDescriptor());
162             }
163             if (tmpKernelDir.listFiles().length == 0) {
164                 throw new TargetSetupError(
165                         "Could not find kernel images", device.getDeviceDescriptor());
166             }
167             // Run the mix kernel tool and generate new device image into tmpNewDeviceDir
168             CLog.i("running mixkernel tool from %s", mMixKernelToolPath);
169             runMixKernelTool(device, tmpDeviceDir, tmpKernelDir, tmpGkiDir, tmpNewDeviceDir);
170             // Find the new device image and copy it to device build info's device image file
171             setNewDeviceImage(buildInfo, tmpNewDeviceDir);
172 
173         } catch (IOException e) {
174             throw new TargetSetupError(
175                     "Could not mix device and kernel images", e, device.getDeviceDescriptor());
176         } finally {
177             FileUtil.recursiveDelete(tmpDeviceDir);
178             FileUtil.recursiveDelete(tmpKernelDir);
179             FileUtil.recursiveDelete(tmpGkiDir);
180             FileUtil.recursiveDelete(tmpNewDeviceDir);
181         }
182     }
183 
184     /**
185      * Copy device image from source file to the specified destination directory.
186      *
187      * @param srcFile the device image source file will be copied from
188      * @param destDir the destination directory {@link File} that the files will be copied to
189      * @throws IOException if hit IOException
190      */
191     @VisibleForTesting
copyDeviceImageToDir(File srcFile, File destDir)192     void copyDeviceImageToDir(File srcFile, File destDir) throws IOException {
193         // Recover the original file name
194         String srcFileName = srcFile.getName();
195         String suffix = FileUtil.getExtension(srcFileName);
196         String base = FileUtil.getBaseName(srcFileName);
197         String newFileName;
198         Pattern pattern = Pattern.compile(".*-img-\\d+");
199         Matcher matcher = pattern.matcher(base);
200         if (matcher.find()) {
201             newFileName = matcher.group();
202         } else {
203             newFileName = new String("device-img-001");
204         }
205         if (!suffix.isEmpty()) {
206             newFileName = newFileName + suffix;
207         } else {
208             newFileName = newFileName + ".zip";
209         }
210         File dstFile = new File(destDir, newFileName);
211         CLog.i("Copy %s to %s", srcFile.toString(), dstFile.toString());
212         FileUtil.hardlinkFile(srcFile, dstFile);
213     }
214 
215     /**
216      * Copy files with labelled filekey to the specified destination directory, with {prefix}
217      * stripped. Also replaces {zip} with ".zip" due to MTT cannot handle .zip file natively. See
218      * test for examples.
219      *
220      * @param fileKey the fileKey of the source file will be copied from
221      * @param srcFile the source file will be copied from
222      * @param destDir the destination directory {@link File} that the files will be copied to
223      * @throws IOException if hit IOException
224      */
225     @VisibleForTesting
copyLabelFileToDir(String fileKey, File srcFile, File destDir)226     void copyLabelFileToDir(String fileKey, File srcFile, File destDir) throws IOException {
227         // Recover the original file name
228         String newFileName = fileKey.replace("{zip}", ".zip").replaceAll("\\{\\w+\\}", "");
229         File dstFile = new File(destDir, newFileName);
230         CLog.i("Copy %s %s to %s", fileKey, srcFile.toString(), dstFile.toString());
231         FileUtil.hardlinkFile(srcFile, dstFile);
232     }
233 
234     /**
235      * Find the mix kernel tool if mMixKernelToolPath does not exist. If mix-kernel-tool-path
236      * provides a valid mixing tool, use it. Otherwise, try to find the mixing tool from the build
237      * provider
238      *
239      * @param buildInfo the {@link IBuildInfo} where to search for mix kernel script
240      * @throws TargetSetupError if flashing script missing or fails
241      * @throws IOException if hit IOException
242      */
findMixKernelTool(IBuildInfo buildInfo)243     private void findMixKernelTool(IBuildInfo buildInfo) throws TargetSetupError, IOException {
244         if (mMixKernelToolPath == null) {
245             CLog.i(
246                     "File mix-kernel-tool-path is not configured. Use devices build's testsdir"
247                             + " as the mix kernel tool path.");
248             mMixKernelToolPath = buildInfo.getFile(BuildInfoFileKey.TESTDIR_IMAGE.getFileKey());
249             if (mMixKernelToolPath == null || !mMixKernelToolPath.isDirectory()) {
250                 throw new TargetSetupError(
251                         String.format(
252                                 "There is no %s to search for mix kernel tool. Please assign a test"
253                                         + " artifact with name %s and type TEST_PACKAGE",
254                                 BuildInfoFileKey.TESTDIR_IMAGE.getFileKey(),
255                                 BuildInfoFileKey.TESTDIR_IMAGE.getFileKey()),
256                         InfraErrorIdentifier.OPTION_CONFIGURATION_ERROR);
257             }
258         }
259         if (mMixKernelToolPath.isDirectory()) {
260             File mixKernelTool = FileUtil.findFile(mMixKernelToolPath, mMixKernelToolName);
261             if (mixKernelTool == null || !mixKernelTool.exists()) {
262                 throw new TargetSetupError(
263                         String.format(
264                                 "Could not find the mix kernel tool %s from %s",
265                                 mMixKernelToolName, mMixKernelToolPath),
266                         InfraErrorIdentifier.OPTION_CONFIGURATION_ERROR);
267             }
268             mMixKernelToolPath = mixKernelTool;
269         }
270         if (!mMixKernelToolPath.canExecute()) {
271             FileUtil.chmodGroupRWX(mMixKernelToolPath);
272         }
273     }
274 
275     /**
276      * Find the new device image and set it as device image for the device.
277      *
278      * @param buildInfo the build info
279      * @param newDeviceDir the directory {@link File} where to find the new device image
280      * @throws TargetSetupError if fail to get the new device image
281      * @throws IOException if hit IOException
282      */
283     @VisibleForTesting
setNewDeviceImage(IBuildInfo buildInfo, File newDeviceDir)284     void setNewDeviceImage(IBuildInfo buildInfo, File newDeviceDir)
285             throws TargetSetupError, IOException {
286         CLog.i(
287                 "Before mixing kernel, the device image %s is of size %d",
288                 buildInfo.getFile(BuildInfoFileKey.DEVICE_IMAGE.getFileKey()).toString(),
289                 buildInfo.getFile(BuildInfoFileKey.DEVICE_IMAGE.getFileKey()).length());
290 
291         File newDeviceImage = FileUtil.findFile(newDeviceDir, ".*-img-.*.zip$");
292         if (newDeviceImage == null || !newDeviceImage.exists()) {
293             throw new TargetSetupError(
294                     "Failed to get a new device image after mixing",
295                     InfraErrorIdentifier.OPTION_CONFIGURATION_ERROR);
296         }
297         CLog.i(
298                 "Successfully generated new device image %s of size %d",
299                 newDeviceImage, newDeviceImage.length());
300         String deviceImagePath = buildInfo.getFile(BuildInfoFileKey.DEVICE_IMAGE).getAbsolutePath();
301         buildInfo.getFile(BuildInfoFileKey.DEVICE_IMAGE.getFileKey()).delete();
302         FileUtil.hardlinkFile(newDeviceImage, new File(deviceImagePath));
303 
304         try {
305             // Modifying device image cannot work with incremental flashing so
306             // self disable it.
307             CLog.d("Disabling incremental flashing.");
308             mConfig.injectOptionValue("incremental-flashing", "false");
309             mConfig.injectOptionValue("force-disable-incremental-flashing", "true");
310         } catch (ConfigurationException ignore) {
311             CLog.e(ignore);
312         }
313 
314         CLog.i(
315                 "After mixing kernel, the device image %s is of size %d",
316                 buildInfo.getFile(BuildInfoFileKey.DEVICE_IMAGE.getFileKey()).toString(),
317                 buildInfo.getFile(BuildInfoFileKey.DEVICE_IMAGE.getFileKey()).length());
318     }
319 
320     /**
321      * Run mix kernel tool to generate the new device build
322      *
323      * <p>Mixing tool Usage: build_mixed_kernels device_dir out_dir target flavor kernel_dir
324      *
325      * @param device the test device
326      * @param oldDeviceDir the directory {@link File} contains old device images
327      * @param kernelDir the directory {@link File} contains kernel images destination
328      * @param gkiDir the directory {@link File} contains GKI kernel images destination
329      * @param newDeviceDir the directory {@link File} where new device images will be generated to
330      * @throws TargetSetupError if fails to run mix kernel tool
331      * @throws IOException
332      */
runMixKernelTool( ITestDevice device, File oldDeviceDir, File kernelDir, File gkiDir, File newDeviceDir)333     protected void runMixKernelTool(
334             ITestDevice device, File oldDeviceDir, File kernelDir, File gkiDir, File newDeviceDir)
335             throws TargetSetupError {
336         List<String> cmd = ArrayUtil.list(mMixKernelToolPath.getAbsolutePath());
337         // Tool command line: $0 [<options>] --gki_dir gkiDir oldDeviceDir kernelDir newDeviceDir
338         cmd.addAll(mMixKernelArgs);
339         if (gkiDir.listFiles().length > 0) {
340             cmd.add("--gki_dir");
341             cmd.add(gkiDir.toString());
342         }
343         cmd.add(oldDeviceDir.toString());
344         cmd.add(kernelDir.toString());
345         cmd.add(newDeviceDir.toString());
346         CLog.i("Run %s to mix kernel and device build", mMixKernelToolPath.toString());
347         CommandResult result =
348                 RunUtil.getDefault()
349                         .runTimedCmd(
350                                 mMixingWaitTime.toMillis(), cmd.toArray(new String[cmd.size()]));
351         if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
352             CLog.e(
353                     "Failed to run mix kernel tool. Exit code: %s, stdout: %s, stderr: %s",
354                     result.getStatus(), result.getStdout(), result.getStderr());
355             throw new TargetSetupError(
356                     "Failed to run mix kernel tool. Stderr: " + result.getStderr());
357         }
358         CLog.i(
359                 "Successfully mixed kernel to new device image in %s with files %s",
360                 newDeviceDir.toString(), Arrays.toString(newDeviceDir.list()));
361     }
362 }
363