• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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.util.image;
17 
18 import static org.junit.Assert.assertTrue;
19 
20 import com.android.ddmlib.Log.LogLevel;
21 import com.android.tradefed.build.IBuildInfo;
22 import com.android.tradefed.build.IDeviceBuildInfo;
23 import com.android.tradefed.device.DeviceNotAvailableException;
24 import com.android.tradefed.device.IManagedTestDevice;
25 import com.android.tradefed.device.ITestDevice;
26 import com.android.tradefed.device.ITestDevice.RecoveryMode;
27 import com.android.tradefed.device.NativeDevice;
28 import com.android.tradefed.device.SnapuserdWaitPhase;
29 import com.android.tradefed.device.TestDevice;
30 import com.android.tradefed.device.TestDeviceState;
31 import com.android.tradefed.invoker.TestInformation;
32 import com.android.tradefed.invoker.logger.CurrentInvocation;
33 import com.android.tradefed.invoker.logger.InvocationMetricLogger;
34 import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationGroupMetricKey;
35 import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
36 import com.android.tradefed.invoker.tracing.CloseableTraceScope;
37 import com.android.tradefed.invoker.tracing.TracePropagatingExecutorService;
38 import com.android.tradefed.log.LogUtil.CLog;
39 import com.android.tradefed.result.error.DeviceErrorIdentifier;
40 import com.android.tradefed.result.error.InfraErrorIdentifier;
41 import com.android.tradefed.targetprep.FastbootDeviceFlasher;
42 import com.android.tradefed.targetprep.FlashingResourcesParser;
43 import com.android.tradefed.targetprep.TargetSetupError;
44 import com.android.tradefed.util.CommandResult;
45 import com.android.tradefed.util.CommandStatus;
46 import com.android.tradefed.util.FileUtil;
47 import com.android.tradefed.util.IRunUtil;
48 import com.android.tradefed.util.RunUtil;
49 import com.android.tradefed.util.ZipUtil;
50 import com.android.tradefed.util.ZipUtil2;
51 import com.android.tradefed.util.executor.ParallelDeviceExecutor;
52 import com.android.tradefed.util.image.DeviceImageTracker.FileCacheTracker;
53 
54 import com.google.common.collect.ImmutableSet;
55 
56 import java.io.File;
57 import java.io.IOException;
58 import java.util.ArrayList;
59 import java.util.Arrays;
60 import java.util.HashMap;
61 import java.util.LinkedHashSet;
62 import java.util.List;
63 import java.util.Map;
64 import java.util.Set;
65 import java.util.concurrent.Callable;
66 import java.util.concurrent.CompletableFuture;
67 import java.util.concurrent.ExecutionException;
68 import java.util.concurrent.Executors;
69 import java.util.concurrent.Future;
70 import java.util.concurrent.ThreadFactory;
71 import java.util.concurrent.TimeUnit;
72 import java.util.concurrent.atomic.AtomicInteger;
73 import java.util.stream.Collectors;
74 
75 /** A utility to leverage the incremental image and device update. */
76 public class IncrementalImageUtil {
77 
78     private static final AtomicInteger poolNumber = new AtomicInteger(1);
79 
80     public static final Set<String> DYNAMIC_PARTITIONS_TO_DIFF =
81             ImmutableSet.of(
82                     "product.img",
83                     "system.img",
84                     "system_dlkm.img",
85                     "system_ext.img",
86                     "vendor.img",
87                     "vendor_dlkm.img");
88 
89     private final File mTargetImage;
90     private final ITestDevice mDevice;
91     private final File mCreateSnapshotBinary;
92     private final boolean mApplySnapshot;
93     private final boolean mWipeAfterApplySnapshot;
94     private final boolean mUpdateBootloaderFromUserspace;
95     private boolean mNewFlow;
96     private final SnapuserdWaitPhase mWaitPhase;
97 
98     private boolean mAllowSameBuildFlashing = false;
99     private boolean mAllowUnzipBaseline = false;
100     private boolean mBootloaderNeedsFlashing = false;
101     private boolean mBasebandNeedsFlashing = false;
102     private boolean mUpdateWasCompleted = false;
103     private File mSourceDirectory;
104     private File mTargetDirectory;
105 
106     private ParallelPreparation mParallelSetup;
107     private final IRunUtil mRunUtil;
108 
initialize( ITestDevice device, IDeviceBuildInfo build, File createSnapshot, boolean isIsolatedSetup, boolean allowTrackerlessUpdate, Set<String> allowedTransition, boolean newFlow, boolean updateBootloaderFromUserspace, SnapuserdWaitPhase waitPhase, boolean useMerkleTree)109     public static IncrementalImageUtil initialize(
110             ITestDevice device,
111             IDeviceBuildInfo build,
112             File createSnapshot,
113             boolean isIsolatedSetup,
114             boolean allowTrackerlessUpdate,
115             Set<String> allowedTransition,
116             boolean newFlow,
117             boolean updateBootloaderFromUserspace,
118             SnapuserdWaitPhase waitPhase,
119             boolean useMerkleTree)
120             throws DeviceNotAvailableException {
121         String serialNumber = device.getSerialNumber();
122         FileCacheTracker tracker =
123                 DeviceImageTracker.getDefaultCache().getBaselineDeviceImage(serialNumber);
124         boolean crossRelease = false;
125         if (tracker == null) {
126             CLog.d("Not tracking current baseline image for %s", serialNumber);
127         } else {
128             if (tracker.branch.contains("release")) {
129                 CLog.d("Skipping incremental flashing for release builds origin.");
130                 return null;
131             }
132             if (!tracker.branch.equals(build.getBuildBranch())) {
133                 if (allowedTransition.contains(tracker.branch)
134                         && allowedTransition.contains(build.getBuildBranch())) {
135                     CLog.d(
136                             "Allowing transition from %s => %s",
137                             tracker.branch, build.getBuildBranch());
138                 } else {
139                     CLog.d("Newer build is not on the same branch.");
140                     return null;
141                 }
142             }
143             if (!tracker.flavor.equals(build.getBuildFlavor())) {
144                 CLog.d(
145                         "Allowing cross-flavor update from '%s' to '%s'",
146                         tracker.flavor, build.getBuildFlavor());
147                 crossRelease = true;
148             }
149         }
150         if (!isSnapshotSupported(device, useMerkleTree)) {
151             CLog.d("Incremental flashing not supported.");
152             return null;
153         }
154         if (crossRelease) {
155             InvocationMetricLogger.addInvocationMetrics(
156                     InvocationMetricKey.INCREMENTAL_ACROSS_RELEASE_COUNT, 1);
157         }
158 
159         if (tracker != null) {
160             InvocationMetricLogger.addInvocationMetrics(
161                     InvocationMetricKey.DEVICE_IMAGE_CACHE_ORIGIN,
162                     String.format("%s:%s:%s", tracker.branch, tracker.buildId, tracker.flavor));
163         } else {
164             InvocationMetricLogger.addInvocationMetrics(
165                     InvocationMetricKey.DEVICE_IMAGE_CACHE_ORIGIN, "no baseline");
166         }
167         File merkleTreeDir = null;
168         if (useMerkleTree) {
169             device.executeShellV2Command("mkdir -p /data/verity-hash/");
170             CommandResult merkleDump =
171                     device.executeShellV2Command("snapshotctl dump-verity-hash /data/verity-hash/");
172             if (CommandStatus.SUCCESS.equals(merkleDump.getStatus())) {
173                 try {
174                     merkleTreeDir =
175                             FileUtil.createTempDir(
176                                     "device-merkle-tree", CurrentInvocation.getWorkFolder());
177                     boolean res = device.pullDir("/data/verity-hash/", merkleTreeDir);
178                     if (!res) {
179                         CLog.w("Failed to pull merkle tree");
180                         FileUtil.recursiveDelete(merkleTreeDir);
181                         merkleTreeDir = null;
182                     }
183                 } catch (IOException e) {
184                     CLog.e(e);
185                     FileUtil.recursiveDelete(merkleTreeDir);
186                     merkleTreeDir = null;
187                 }
188             }
189         }
190         return new IncrementalImageUtil(
191                 device,
192                 build.getDeviceImageFile(),
193                 createSnapshot,
194                 newFlow,
195                 updateBootloaderFromUserspace,
196                 waitPhase,
197                 merkleTreeDir);
198     }
199 
IncrementalImageUtil( ITestDevice device, File targetImage, File createSnapshot, boolean newFlow, boolean updateBootloaderFromUserspace, SnapuserdWaitPhase waitPhase, File deviceMerkleTree)200     public IncrementalImageUtil(
201             ITestDevice device,
202             File targetImage,
203             File createSnapshot,
204             boolean newFlow,
205             boolean updateBootloaderFromUserspace,
206             SnapuserdWaitPhase waitPhase,
207             File deviceMerkleTree) {
208         mDevice = device;
209         mApplySnapshot = true;
210         mWipeAfterApplySnapshot = true;
211         mNewFlow = newFlow;
212         mUpdateBootloaderFromUserspace = updateBootloaderFromUserspace;
213         mWaitPhase = waitPhase;
214 
215         mTargetImage = targetImage;
216         mRunUtil = new RunUtil();
217         // TODO: clean up when docker image is updated
218         mRunUtil.setEnvVariable("LD_LIBRARY_PATH", "/tradefed/lib64");
219         if (createSnapshot != null) {
220             File snapshot = createSnapshot;
221             try {
222                 if (createSnapshot.getName().endsWith(".zip")
223                         && ZipUtil.isZipFileValid(createSnapshot, false)) {
224                     File destDir = ZipUtil2.extractZipToTemp(createSnapshot, "create_snapshot");
225                     snapshot = FileUtil.findFile(destDir, "create_snapshot");
226                 }
227             } catch (IOException e) {
228                 CLog.e(e);
229             }
230             mCreateSnapshotBinary = snapshot;
231             FileUtil.chmodGroupRWX(snapshot);
232         } else {
233             mCreateSnapshotBinary = null;
234         }
235         mParallelSetup =
236                 new ParallelPreparation(
237                         Thread.currentThread().getThreadGroup(),
238                         mTargetImage,
239                         deviceMerkleTree);
240         mParallelSetup.start();
241     }
242 
copyImage(File originalImage)243     private static File copyImage(File originalImage) throws IOException {
244         if (originalImage.isDirectory()) {
245             CLog.d("Baseline was already unzipped for %s", originalImage);
246             File copy =
247                     FileUtil.createTempDir(
248                             FileUtil.getBaseName(originalImage.getName()),
249                             CurrentInvocation.getWorkFolder());
250             FileUtil.recursiveHardlink(originalImage, copy);
251             return copy;
252         } else {
253             File copy =
254                     FileUtil.createTempFile(
255                             FileUtil.getBaseName(originalImage.getName()),
256                             ".img",
257                             CurrentInvocation.getWorkFolder());
258             copy.delete();
259             FileUtil.hardlinkFile(originalImage, copy);
260             return copy;
261         }
262     }
263 
264     /** Returns whether or not we can use the snapshot logic to update the device */
isSnapshotSupported(ITestDevice device, boolean useMerkle)265     public static boolean isSnapshotSupported(ITestDevice device, boolean useMerkle)
266             throws DeviceNotAvailableException {
267         // Ensure snapshotctl exists
268         CommandResult whichOutput = device.executeShellV2Command("which snapshotctl");
269         CLog.d("stdout: %s, stderr: %s", whichOutput.getStdout(), whichOutput.getStderr());
270         if (!whichOutput.getStdout().contains("/system/bin/snapshotctl")) {
271             return false;
272         }
273         CommandResult helpOutput = device.executeShellV2Command("snapshotctl");
274         CLog.d("stdout: %s, stderr: %s", helpOutput.getStdout(), helpOutput.getStderr());
275         if (useMerkle) {
276             if (!helpOutput.getStdout().contains("dump-verity-hash")
277                     && !helpOutput.getStderr().contains("dump-verity-hash")) {
278                 return false;
279             }
280         }
281         if (helpOutput.getStdout().contains("apply-update")
282                 || helpOutput.getStderr().contains("apply-update")) {
283             return true;
284         }
285         return false;
286     }
287 
notifyBootloaderNeedsRevert()288     public void notifyBootloaderNeedsRevert() {
289         mBootloaderNeedsFlashing = true;
290     }
291 
notifyBasebadNeedsRevert()292     public void notifyBasebadNeedsRevert() {
293         mBasebandNeedsFlashing = true;
294     }
295 
allowSameBuildFlashing()296     public void allowSameBuildFlashing() {
297         mAllowSameBuildFlashing = true;
298     }
299 
isSameBuildFlashingAllowed()300     public boolean isSameBuildFlashingAllowed() {
301         return mAllowSameBuildFlashing;
302     }
303 
allowUnzipBaseline()304     public void allowUnzipBaseline() {
305         mAllowUnzipBaseline = true;
306     }
307 
useUpdatedFlow()308     public boolean useUpdatedFlow() {
309         return mNewFlow;
310     }
311 
312     /** Returns whether device is currently using snapshots or not. */
isSnapshotInUse(ITestDevice device)313     public static boolean isSnapshotInUse(ITestDevice device) throws DeviceNotAvailableException {
314         CommandResult dumpOutput = device.executeShellV2Command("snapshotctl dump");
315         CLog.d("stdout: %s, stderr: %s", dumpOutput.getStdout(), dumpOutput.getStderr());
316         if (dumpOutput.getStdout().contains("Using snapuserd: 0")) {
317             return false;
318         }
319         return true;
320     }
321 
updateDeviceWithNewFlow(File currentBootloader, File currentRadio)322     public void updateDeviceWithNewFlow(File currentBootloader, File currentRadio)
323             throws DeviceNotAvailableException, TargetSetupError {
324         if (!mNewFlow || !mApplySnapshot || !mWipeAfterApplySnapshot) {
325             mNewFlow = false;
326             return;
327         }
328         // If device isn't online, we can't use the new flow
329         if (!TestDeviceState.ONLINE.equals(mDevice.getDeviceState())) {
330             mNewFlow = false;
331             return;
332         }
333         InvocationMetricLogger.addInvocationMetrics(InvocationMetricKey.INCREMENTAL_NEW_FLOW, 1);
334         // If enable, push the bootloader from userspace like OTA
335         if (mUpdateBootloaderFromUserspace) {
336             updateBootloaderFromUserspace(currentBootloader);
337         }
338         CLog.d("Updating with new incremental flow.");
339         updateDevice(currentBootloader, currentRadio);
340     }
341 
342     /** Updates the device using the snapshot logic. */
updateDevice(File currentBootloader, File currentRadio)343     public void updateDevice(File currentBootloader, File currentRadio)
344             throws DeviceNotAvailableException, TargetSetupError {
345         if (mDevice.isStateBootloaderOrFastbootd()) {
346             mDevice.rebootUntilOnline();
347         }
348         if (!mDevice.enableAdbRoot()) {
349             throw new TargetSetupError(
350                     "Failed to obtain root, this is required for incremental update.",
351                     InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
352         }
353         try {
354             internalUpdateDevice(currentBootloader, currentRadio);
355         } catch (DeviceNotAvailableException | TargetSetupError | RuntimeException e) {
356             InvocationMetricLogger.addInvocationMetrics(
357                     InvocationMetricKey.INCREMENTAL_FLASHING_UPDATE_FAILURE, 1);
358             throw e;
359         }
360     }
361 
updateBootloaderFromUserspace(File currentBootloader)362     private void updateBootloaderFromUserspace(File currentBootloader)
363             throws DeviceNotAvailableException, TargetSetupError {
364         File bootloaderDir = null;
365         try (CloseableTraceScope ignored = new CloseableTraceScope("update_bootloader_userspace")) {
366             String listAbPartitions = mDevice.getProperty("ro.product.ab_ota_partitions");
367             if (listAbPartitions == null) {
368                 throw new TargetSetupError(
369                         "Couldn't query ab_ota_partitions",
370                         InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
371             }
372             String bootSuffix = mDevice.getProperty("ro.boot.slot_suffix");
373             if (bootSuffix == null) {
374                 throw new TargetSetupError(
375                         "Couldn't query ro.boot.slot_suffix",
376                         InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
377             }
378             if (bootSuffix.equals("_a")) {
379                 bootSuffix = "_b";
380             } else if (bootSuffix.equals("_b")) {
381                 bootSuffix = "_a";
382             } else {
383                 throw new TargetSetupError(
384                         String.format("unexpected ro.boot.slot_suffix: %s", bootSuffix),
385                         InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
386             }
387 
388             Set<String> partitions =
389                     Arrays.asList(listAbPartitions.split(",")).stream()
390                             .map(p -> p + ".img")
391                             .collect(Collectors.toSet());
392             CLog.d("Bootloader partitions to be considered: %s", partitions);
393             try {
394                 bootloaderDir =
395                         FileUtil.createTempDir("bootloader", CurrentInvocation.getWorkFolder());
396                 FastbootPack.unpack(currentBootloader, bootloaderDir, null, false);
397             } catch (IOException e) {
398                 throw new TargetSetupError(
399                         e.getMessage(), e, InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
400             }
401             Set<File> toBePushed = new LinkedHashSet<File>();
402             for (File f : bootloaderDir.listFiles()) {
403                 if (partitions.contains(f.getName())) {
404                     toBePushed.add(f);
405                 }
406             }
407             CLog.d("Bootloader partitions to be updated: %s", toBePushed);
408             mDevice.executeShellV2Command("mkdir -p /data/bootloader");
409             for (File push : toBePushed) {
410                 boolean success = mDevice.pushFile(push, "/data/bootloader/" + push.getName());
411                 if (!success) {
412                     throw new TargetSetupError(
413                             "Failed to push bootloader partition.",
414                             InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
415                 }
416             }
417             for (File write : toBePushed) {
418                 CommandResult writeRes =
419                         mDevice.executeShellV2Command(
420                                 String.format(
421                                         "dd if=/data/bootloader/%s of=/dev/block/by-name/%s%s",
422                                         write.getName(),
423                                         FileUtil.getBaseName(write.getName()),
424                                         bootSuffix));
425                 if (!CommandStatus.SUCCESS.equals(writeRes.getStatus())) {
426                     throw new TargetSetupError(
427                             String.format(
428                                     "Failed to write bootloader partition: %s",
429                                     writeRes.getStderr()),
430                             InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
431                 }
432             }
433         } finally {
434             FileUtil.recursiveDelete(bootloaderDir);
435         }
436     }
437 
internalUpdateDevice(File currentBootloader, File currentRadio)438     private void internalUpdateDevice(File currentBootloader, File currentRadio)
439             throws DeviceNotAvailableException, TargetSetupError {
440         InvocationMetricLogger.addInvocationMetrics(
441                 InvocationMetricKey.INCREMENTAL_FLASHING_ATTEMPT_COUNT, 1);
442         // Join the unzip thread
443         long startWait = System.currentTimeMillis();
444         try {
445             mParallelSetup.join();
446         } catch (InterruptedException e) {
447             mParallelSetup.cleanUpFiles();
448             throw new RuntimeException(e);
449         } finally {
450             InvocationMetricLogger.addInvocationMetrics(
451                     InvocationMetricKey.INCREMENTAL_FLASHING_WAIT_PARALLEL_SETUP,
452                     System.currentTimeMillis() - startWait);
453         }
454         if (mParallelSetup.getError() != null) {
455             mParallelSetup.cleanUpFiles();
456             InvocationMetricLogger.addInvocationMetrics(
457                     InvocationMetricKey.INCREMENTAL_FALLBACK_REASON,
458                     mParallelSetup.getError().getMessage());
459             throw mParallelSetup.getError();
460         }
461         boolean bootComplete =
462                 mDevice.waitForBootComplete(mDevice.getOptions().getAvailableTimeout());
463         if (!bootComplete) {
464             mParallelSetup.cleanUpFiles();
465             throw new TargetSetupError(
466                     "Failed to boot within timeout.",
467                     InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
468         }
469         // We need a few seconds after boot complete for update_engine to finish
470         // TODO: we could improve by listening to some update_engine messages.
471         if (!mNewFlow) {
472             RunUtil.getDefault().sleep(5000L);
473         }
474         File srcDirectory = mParallelSetup.getSrcDirectory();
475         File targetDirectory = mParallelSetup.getTargetDirectory();
476         File workDir = mParallelSetup.getWorkDir();
477         try (CloseableTraceScope ignored = new CloseableTraceScope("update_device")) {
478             // Once block comparison is successful, log the information
479             logTargetInformation(targetDirectory);
480             long totalPatchSizes = logPatchesInformation(workDir);
481             // if we have more than 2.5GB we will overflow super partition size to /data and we
482             // can't use the feature
483             if (totalPatchSizes > 2300000000L) {
484                 InvocationMetricLogger.addInvocationMetrics(
485                         InvocationMetricKey.INCREMENTAL_FALLBACK_REASON, "Patches too large.");
486                 throw new TargetSetupError(
487                         String.format(
488                                 "Total patch size is %s bytes. Too large to use the feature."
489                                         + " falling back",
490                                 totalPatchSizes),
491                         InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
492             }
493 
494             mDevice.executeShellV2Command("mkdir -p /data/ndb");
495             mDevice.executeShellV2Command("rm -rf /data/ndb/*.patch");
496 
497             mDevice.executeShellV2Command("snapshotctl unmap-snapshots");
498             mDevice.executeShellV2Command("snapshotctl delete-snapshots");
499 
500             RecoveryMode mode = mDevice.getRecoveryMode();
501             mDevice.setRecoveryMode(RecoveryMode.NONE);
502             try {
503                 List<Callable<Boolean>> pushTasks = new ArrayList<>();
504                 for (File f : workDir.listFiles()) {
505                     try (CloseableTraceScope push =
506                             new CloseableTraceScope("push:" + f.getName())) {
507                         pushTasks.add(
508                                 () -> {
509                                     boolean success;
510                                     if (f.isDirectory()) {
511                                         success = mDevice.pushDir(f, "/data/ndb/");
512                                     } else {
513                                         success = mDevice.pushFile(f, "/data/ndb/" + f.getName());
514                                     }
515                                     CLog.d(
516                                             "Push status: %s. %s->%s",
517                                             success, f, "/data/ndb/" + f.getName());
518                                     assertTrue(success);
519                                     return true;
520                                 });
521                     }
522                 }
523                 ParallelDeviceExecutor<Boolean> pushExec =
524                         new ParallelDeviceExecutor<Boolean>(pushTasks.size());
525                 pushExec.invokeAll(pushTasks, 0, TimeUnit.MINUTES);
526                 if (pushExec.hasErrors()) {
527                     for (Throwable err : pushExec.getErrors()) {
528                         InvocationMetricLogger.addInvocationMetrics(
529                                 InvocationMetricKey.INCREMENTAL_FALLBACK_REASON, err.getMessage());
530                         if (err instanceof DeviceNotAvailableException) {
531                             throw (DeviceNotAvailableException) err;
532                         }
533                     }
534                     throw new TargetSetupError(
535                             String.format("Failed to push patches."),
536                             pushExec.getErrors().get(0),
537                             InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
538                 }
539             } finally {
540                 mDevice.setRecoveryMode(mode);
541             }
542 
543             CommandResult listSnapshots = mDevice.executeShellV2Command("ls -l /data/ndb/");
544             CLog.d("stdout: %s, stderr: %s", listSnapshots.getStdout(), listSnapshots.getStderr());
545             mDevice.logOnDevice("Tradefed", LogLevel.DEBUG, "Running snapshotctl apply-update");
546             String applyCommand = "snapshotctl apply-update /data/ndb/";
547             if (mWipeAfterApplySnapshot) {
548                 applyCommand += " -w";
549             }
550             CommandResult mapOutput = mDevice.executeShellV2Command(applyCommand);
551             CLog.d("stdout: %s, stderr: %s", mapOutput.getStdout(), mapOutput.getStderr());
552             if (!CommandStatus.SUCCESS.equals(mapOutput.getStatus())) {
553                 InvocationMetricLogger.addInvocationMetrics(
554                         InvocationMetricKey.INCREMENTAL_FALLBACK_REASON, "Failed apply-update");
555                 // Clean state if apply-update fails
556                 mDevice.executeShellV2Command("snapshotctl unmap-snapshots");
557                 mDevice.executeShellV2Command("snapshotctl delete-snapshots");
558                 throw new TargetSetupError(
559                         String.format(
560                                 "Failed to apply-update.\nstdout:%s\nstderr:%s",
561                                 mapOutput.getStdout(), mapOutput.getStderr()),
562                         InfraErrorIdentifier.INCREMENTAL_FLASHING_ERROR);
563             }
564             try {
565                 if (mNewFlow && mDevice instanceof TestDevice) {
566                     ((TestDevice) mDevice).setFirstBootloaderReboot();
567                 }
568                 mDevice.rebootIntoBootloader();
569             } catch (DeviceNotAvailableException e) {
570                 if (mNewFlow) {
571                     InvocationMetricLogger.addInvocationMetrics(
572                             InvocationMetricKey.INCREMENTAL_FIRST_BOOTLOADER_REBOOT_FAIL, 1);
573                 }
574                 throw e;
575             }
576 
577             if (mApplySnapshot) {
578                 if (mWipeAfterApplySnapshot) {
579                     CommandResult cancelResults =
580                             mDevice.executeFastbootCommand("snapshot-update", "cancel");
581                     CLog.d("Cancel status: %s", cancelResults.getStatus());
582                     CLog.d("Cancel stdout: %s", cancelResults.getStdout());
583                     CLog.d("Cancel stderr: %s", cancelResults.getStderr());
584                     CommandResult wipeResults = mDevice.executeFastbootCommand("-w");
585                     CLog.d("wipe status: %s", wipeResults.getStatus());
586                     CLog.d("wipe stdout: %s", wipeResults.getStdout());
587                     CLog.d("wipe stderr: %s", wipeResults.getStderr());
588                 }
589                 updateBootloaderAndBasebandIfNeeded(
590                         targetDirectory, currentBootloader, currentRadio);
591             }
592             flashStaticPartition(targetDirectory);
593             mSourceDirectory = srcDirectory;
594 
595             mDevice.enableAdbRoot();
596 
597             if (mApplySnapshot) {
598                 mDevice.notifySnapuserd(mWaitPhase);
599                 mDevice.waitForSnapuserd(SnapuserdWaitPhase.BLOCK_AFTER_UPDATE);
600             } else {
601                 // If patches are mounted, just print snapuserd once
602                 CommandResult psOutput = mDevice.executeShellV2Command("ps -ef | grep snapuserd");
603                 CLog.d("stdout: %s, stderr: %s", psOutput.getStdout(), psOutput.getStderr());
604             }
605             mTargetDirectory = targetDirectory;
606             mUpdateWasCompleted = true;
607         } catch (DeviceNotAvailableException | RuntimeException e) {
608             if (mSourceDirectory == null) {
609                 FileUtil.recursiveDelete(srcDirectory);
610             }
611             throw e;
612         } finally {
613             FileUtil.recursiveDelete(workDir);
614         }
615     }
616 
617     /** Returns whether update was completed or not. */
updateCompleted()618     public boolean updateCompleted() {
619         return mUpdateWasCompleted;
620     }
621 
getExtractedTargetDirectory()622     public File getExtractedTargetDirectory() {
623         return mTargetDirectory;
624     }
625 
626     /** When doing some of the apply logic, we can clean up files right after setup. */
cleanAfterSetup()627     public void cleanAfterSetup() {
628         if (!mApplySnapshot) {
629             return;
630         }
631         // Delete the copy we made to use the incremental update
632         FileUtil.recursiveDelete(mSourceDirectory);
633         FileUtil.recursiveDelete(mTargetDirectory);
634         // In case of same build flashing, we should clean the setup operation
635         if (mParallelSetup != null) {
636             try {
637                 mParallelSetup.join();
638             } catch (InterruptedException e) {
639                 CLog.e(e);
640             }
641             mParallelSetup.cleanUpFiles();
642         }
643     }
644 
645     /*
646      * Returns the device to its original state.
647      */
teardownDevice(TestInformation testInfo)648     public void teardownDevice(TestInformation testInfo) throws DeviceNotAvailableException {
649         // Delete the copy we made to use the incremental update
650         FileUtil.recursiveDelete(mSourceDirectory);
651         FileUtil.recursiveDelete(mTargetDirectory);
652         // In case of same build flashing, we should clean the setup operation
653         if (mParallelSetup != null) {
654             try {
655                 mParallelSetup.join();
656             } catch (InterruptedException e) {
657                 CLog.e(e);
658             }
659             mParallelSetup.cleanUpFiles();
660         }
661     }
662 
updateBootloaderAndBasebandIfNeeded( File deviceImageUnzipped, File bootloader, File baseband)663     private void updateBootloaderAndBasebandIfNeeded(
664             File deviceImageUnzipped, File bootloader, File baseband)
665             throws DeviceNotAvailableException, TargetSetupError {
666         FlashingResourcesParser parser = new FlashingResourcesParser(deviceImageUnzipped);
667         if (bootloader == null) {
668             CLog.w("No bootloader file to flash.");
669         } else {
670             if (shouldFlashBootloader(mDevice, parser.getRequiredBootloaderVersion())) {
671                 CommandResult bootloaderFlashTarget =
672                         mDevice.executeFastbootCommand(
673                                 "flash", "bootloader", bootloader.getAbsolutePath());
674                 CLog.d("Status: %s", bootloaderFlashTarget.getStatus());
675                 CLog.d("stdout: %s", bootloaderFlashTarget.getStdout());
676                 CLog.d("stderr: %s", bootloaderFlashTarget.getStderr());
677                 mDevice.rebootIntoBootloader();
678             }
679         }
680         if (baseband == null) {
681             CLog.w("No baseband file to flash");
682         } else {
683             if (shouldFlashBaseband(mDevice, parser.getRequiredBasebandVersion())) {
684                 CommandResult radioFlashTarget =
685                         mDevice.executeFastbootCommand(
686                                 "flash", "radio", baseband.getAbsolutePath());
687                 CLog.d("Status: %s", radioFlashTarget.getStatus());
688                 CLog.d("stdout: %s", radioFlashTarget.getStdout());
689                 CLog.d("stderr: %s", radioFlashTarget.getStderr());
690                 mDevice.rebootIntoBootloader();
691             }
692         }
693     }
694 
revertBootloaderAndBasebandifNeeded(File bootloader, File baseband)695     private void revertBootloaderAndBasebandifNeeded(File bootloader, File baseband)
696             throws DeviceNotAvailableException {
697         if (mBootloaderNeedsFlashing) {
698             if (bootloader == null) {
699                 CLog.w("No bootloader file to flash.");
700             } else {
701                 mDevice.rebootIntoBootloader();
702 
703                 CommandResult bootloaderFlashTarget =
704                         mDevice.executeFastbootCommand(
705                                 "flash", "bootloader", bootloader.getAbsolutePath());
706                 CLog.d("Status: %s", bootloaderFlashTarget.getStatus());
707                 CLog.d("stdout: %s", bootloaderFlashTarget.getStdout());
708                 CLog.d("stderr: %s", bootloaderFlashTarget.getStderr());
709             }
710         }
711         if (mBasebandNeedsFlashing) {
712             if (baseband == null) {
713                 CLog.w("No baseband file to flash");
714             } else {
715                 mDevice.rebootIntoBootloader();
716 
717                 CommandResult radioFlashTarget =
718                         mDevice.executeFastbootCommand(
719                                 "flash", "radio", baseband.getAbsolutePath());
720                 CLog.d("Status: %s", radioFlashTarget.getStatus());
721                 CLog.d("stdout: %s", radioFlashTarget.getStdout());
722                 CLog.d("stderr: %s", radioFlashTarget.getStderr());
723             }
724         }
725     }
726 
blockCompare(File srcImage, File srcMerkleTree, File targetImage, File workDir)727     private void blockCompare(File srcImage, File srcMerkleTree, File targetImage, File workDir) {
728         try (CloseableTraceScope ignored =
729                 new CloseableTraceScope("block_compare:" + srcImage.getName())) {
730             mRunUtil.setWorkingDir(workDir);
731 
732             String createSnapshot = "create_snapshot"; // Expected to be on PATH
733             if (mCreateSnapshotBinary != null && mCreateSnapshotBinary.exists()) {
734                 createSnapshot = mCreateSnapshotBinary.getAbsolutePath();
735             }
736             String[] command = null;
737             if (srcMerkleTree.exists()) {
738                 command =
739                         new String[] {
740                             createSnapshot,
741                             "--source=" + srcMerkleTree.getAbsolutePath(),
742                             "--target=" + targetImage.getAbsolutePath(),
743                             "--merkel_tree"
744                         };
745             } else {
746                 command =
747                         new String[] {
748                             createSnapshot,
749                             "--source=" + srcImage.getAbsolutePath(),
750                             "--target=" + targetImage.getAbsolutePath()
751                         };
752             }
753 
754             CommandResult result = mRunUtil.runTimedCmd(0L, command);
755             if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
756                 throw new RuntimeException(
757                         String.format("%s\n%s", result.getStdout(), result.getStderr()));
758             }
759             File[] listFiles = workDir.listFiles();
760             CLog.d("%s", Arrays.asList(listFiles));
761         }
762     }
763 
flashStaticPartition(File imageDirectory)764     private boolean flashStaticPartition(File imageDirectory)
765             throws DeviceNotAvailableException, TargetSetupError {
766         // Invalidate properties to be updated after reboot into the new
767         // image.
768         if (mDevice instanceof NativeDevice) {
769             ((NativeDevice) mDevice).invalidatePropertyCache();
770         }
771         Map<String, String> envMap = new HashMap<>();
772         envMap.put("ANDROID_PRODUCT_OUT", imageDirectory.getAbsolutePath());
773         CommandResult fastbootResult =
774                 mDevice.executeLongFastbootCommand(
775                         envMap,
776                         "flashall",
777                         "--exclude-dynamic-partitions",
778                         "--disable-super-optimization");
779         CLog.d("Status: %s", fastbootResult.getStatus());
780         CLog.d("stdout: %s", fastbootResult.getStdout());
781         CLog.d("stderr: %s", fastbootResult.getStderr());
782         if (!CommandStatus.SUCCESS.equals(fastbootResult.getStatus())) {
783             return false;
784         }
785         RecoveryMode recoveryMode = mDevice.getRecoveryMode();
786         try {
787             mDevice.setRecoveryMode(RecoveryMode.NONE);
788             ((IManagedTestDevice) mDevice).getMonitor().attachFinalState(TestDeviceState.RECOVERY);
789             boolean available = mDevice.waitForDeviceAvailable(5 * 60 * 1000L);
790             if (!available) {
791                 if (TestDeviceState.RECOVERY.equals(mDevice.getDeviceState())) {
792                     InvocationMetricLogger.addInvocationMetrics(
793                             InvocationMetricKey.INCREMENTAL_RECOVERY_FALLBACK, 1);
794                     // Go back to bootloader for fallback flashing
795                     mDevice.rebootIntoBootloader();
796                     CommandResult result = mDevice.executeFastbootCommand("-w");
797                     CLog.d("wipe status: %s", result.getStatus());
798                     CLog.d("wipe stdout: %s", result.getStdout());
799                     CLog.d("wipe stderr: %s", result.getStderr());
800                     throw new TargetSetupError(
801                             "Device went to recovery unexpectedly",
802                             DeviceErrorIdentifier.DEVICE_UNEXPECTED_RESPONSE);
803                 }
804             }
805         } finally {
806             mDevice.setRecoveryMode(recoveryMode);
807         }
808         return true;
809     }
810 
logPatchesInformation(File patchesDirectory)811     private long logPatchesInformation(File patchesDirectory) {
812         long totalPatchesSize = 0L;
813         for (File patch : patchesDirectory.listFiles()) {
814             if (patch == null) {
815                 CLog.w("Something went wrong listing %s", patchesDirectory);
816                 return 0L;
817             }
818             totalPatchesSize += patch.length();
819             InvocationMetricLogger.addInvocationMetrics(
820                     InvocationGroupMetricKey.INCREMENTAL_FLASHING_PATCHES_SIZE,
821                     patch.getName(),
822                     patch.length());
823         }
824         return totalPatchesSize;
825     }
826 
logTargetInformation(File targetDirectory)827     private void logTargetInformation(File targetDirectory) {
828         for (File patch : targetDirectory.listFiles()) {
829             if (patch == null) {
830                 CLog.w("Something went wrong listing target %s", targetDirectory);
831                 return;
832             }
833             if (DYNAMIC_PARTITIONS_TO_DIFF.contains(patch.getName())) {
834                 InvocationMetricLogger.addInvocationMetrics(
835                         InvocationGroupMetricKey.INCREMENTAL_FLASHING_TARGET_SIZE,
836                         patch.getName(),
837                         patch.length());
838             }
839         }
840     }
841 
shouldFlashBootloader(ITestDevice device, String bootloaderVersion)842     private boolean shouldFlashBootloader(ITestDevice device, String bootloaderVersion)
843             throws DeviceNotAvailableException, TargetSetupError {
844         String currentBootloaderVersion =
845                 FastbootDeviceFlasher.fetchImageVersion(mRunUtil, device, "bootloader");
846         if (bootloaderVersion != null && !bootloaderVersion.equals(currentBootloaderVersion)) {
847             CLog.i("Flashing bootloader %s", bootloaderVersion);
848             return true;
849         } else {
850             CLog.i("Bootloader is already version %s, skipping flashing", currentBootloaderVersion);
851             return false;
852         }
853     }
854 
shouldFlashBaseband(ITestDevice device, String basebandVersion)855     private boolean shouldFlashBaseband(ITestDevice device, String basebandVersion)
856             throws DeviceNotAvailableException, TargetSetupError {
857         String currentBaseBandVersion =
858                 FastbootDeviceFlasher.fetchImageVersion(mRunUtil, device, "baseband");
859         if (basebandVersion != null && !basebandVersion.equals(currentBaseBandVersion)) {
860             CLog.i("Flashing bootloader %s", basebandVersion);
861             return true;
862         } else {
863             CLog.i("Bootloader is already version %s, skipping flashing", currentBaseBandVersion);
864             return false;
865         }
866     }
867 
getSplVersion(IBuildInfo build)868     private static String getSplVersion(IBuildInfo build) {
869         File buildProp = build.getFile("build.prop");
870         if (buildProp == null) {
871             CLog.d("No target build.prop found for comparison.");
872             return null;
873         }
874         try {
875             String props = FileUtil.readStringFromFile(buildProp);
876             for (String line : props.split("\n")) {
877                 if (line.startsWith("ro.build.version.security_patch=")) {
878                     return line.split("=")[1];
879                 }
880             }
881         } catch (IOException e) {
882             CLog.e(e);
883         }
884         return null;
885     }
886 
887     private class ParallelPreparation extends Thread {
888 
889         private final File mDeviceOriginMerkleTree;
890         private final File mSetupTargetImage;
891 
892         private File mSrcDirectory;
893         private File mTargetDirectory;
894         private File mWorkDir;
895         private TargetSetupError mError;
896 
ParallelPreparation( ThreadGroup currentGroup, File targetImage, File deviceMerkleTree)897         public ParallelPreparation(
898                 ThreadGroup currentGroup, File targetImage, File deviceMerkleTree) {
899             super(currentGroup, "incremental-flashing-preparation");
900             setDaemon(true);
901             this.mDeviceOriginMerkleTree = deviceMerkleTree;
902             this.mSetupTargetImage = targetImage;
903         }
904 
905         @Override
run()906         public void run() {
907             ThreadGroup currentGroup = Thread.currentThread().getThreadGroup();
908             ThreadFactory factory =
909                     new ThreadFactory() {
910                         @Override
911                         public Thread newThread(Runnable r) {
912                             Thread t =
913                                     new Thread(
914                                             currentGroup,
915                                             r,
916                                             "unzip-pool-task-" + poolNumber.getAndIncrement());
917                             t.setDaemon(true);
918                             return t;
919                         }
920                     };
921             try (CloseableTraceScope ignored = new CloseableTraceScope("unzip_device_images")) {
922                 mSrcDirectory = mDeviceOriginMerkleTree;
923                 mTargetDirectory = FileUtil.createTempDir("incremental_target");
924                 Future<Boolean> futureTargetDir =
925                         CompletableFuture.supplyAsync(
926                                 () -> {
927                                     if (mSetupTargetImage.isDirectory()) {
928                                         try (CloseableTraceScope unzipTarget =
929                                                 new CloseableTraceScope("hardlink_target")) {
930                                             FileUtil.recursiveHardlink(
931                                                     mSetupTargetImage, mTargetDirectory);
932                                             return true;
933                                         } catch (IOException ioe) {
934                                             throw new RuntimeException(ioe);
935                                         }
936                                     }
937                                     try (CloseableTraceScope unzipTarget =
938                                             new CloseableTraceScope("unzip_target")) {
939                                         ZipUtil2.extractZip(mSetupTargetImage, mTargetDirectory);
940                                         return true;
941                                     } catch (IOException ioe) {
942                                         throw new RuntimeException(ioe);
943                                     }
944                                 },
945                                 TracePropagatingExecutorService.create(
946                                         Executors.newFixedThreadPool(1, factory)));
947                 // Join the unzipping
948                 futureTargetDir.get();
949             } catch (InterruptedException | IOException | ExecutionException e) {
950                 FileUtil.recursiveDelete(mSrcDirectory);
951                 FileUtil.recursiveDelete(mTargetDirectory);
952                 mSrcDirectory = null;
953                 mTargetDirectory = null;
954                 mError =
955                         new TargetSetupError(
956                                 e.getMessage(), e, InfraErrorIdentifier.FAIL_TO_CREATE_FILE);
957                 return;
958             }
959 
960             try {
961                 mWorkDir = FileUtil.createTempDir("block_compare_workdir");
962             } catch (IOException e) {
963                 FileUtil.recursiveDelete(mWorkDir);
964                 FileUtil.recursiveDelete(mSrcDirectory);
965                 FileUtil.recursiveDelete(mTargetDirectory);
966                 mSrcDirectory = null;
967                 mTargetDirectory = null;
968                 mError =
969                         new TargetSetupError(
970                                 e.getMessage(), e, InfraErrorIdentifier.FAIL_TO_CREATE_FILE);
971                 return;
972             }
973 
974             List<Callable<Boolean>> callableTasks = new ArrayList<>();
975             for (String partition : mSrcDirectory.list()) {
976                 String merklePartition = partition.replaceAll(".img", ".pb");
977                 partition = partition.replaceAll(".pb", ".img");
978 
979                 File possibleSrc = new File(mSrcDirectory, partition);
980                 File sourceMerkleTree = new File(mDeviceOriginMerkleTree, merklePartition);
981                 File possibleTarget = new File(mTargetDirectory, partition);
982                 File workDirectory = mWorkDir;
983                 if ((possibleSrc.exists() || sourceMerkleTree.exists())
984                         && possibleTarget.exists()) {
985                     if (DYNAMIC_PARTITIONS_TO_DIFF.contains(partition)) {
986                         callableTasks.add(
987                                 () -> {
988                                     blockCompare(
989                                             possibleSrc,
990                                             sourceMerkleTree,
991                                             possibleTarget,
992                                             workDirectory);
993                                     return true;
994                                 });
995                     }
996                 } else {
997                     CLog.e("Skipping %s no src or target", partition);
998                 }
999             }
1000             ParallelDeviceExecutor<Boolean> executor =
1001                     new ParallelDeviceExecutor<Boolean>(callableTasks.size());
1002             executor.invokeAll(callableTasks, 0, TimeUnit.MINUTES);
1003             if (executor.hasErrors()) {
1004                 mError =
1005                         new TargetSetupError(
1006                                 executor.getErrors().get(0).getMessage(),
1007                                 executor.getErrors().get(0),
1008                                 InfraErrorIdentifier.BLOCK_COMPARE_ERROR);
1009             }
1010         }
1011 
getSrcDirectory()1012         public File getSrcDirectory() {
1013             return mSrcDirectory;
1014         }
1015 
getTargetDirectory()1016         public File getTargetDirectory() {
1017             return mTargetDirectory;
1018         }
1019 
getWorkDir()1020         public File getWorkDir() {
1021             return mWorkDir;
1022         }
1023 
getError()1024         public TargetSetupError getError() {
1025             return mError;
1026         }
1027 
cleanUpFiles()1028         public void cleanUpFiles() {
1029             FileUtil.recursiveDelete(mSrcDirectory);
1030             FileUtil.recursiveDelete(mTargetDirectory);
1031             FileUtil.recursiveDelete(mWorkDir);
1032         }
1033     }
1034 }
1035