1 /*
2  * Copyright (C) 2020 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 
17 package com.android.internal.util.test;
18 
19 import static org.junit.Assert.assertTrue;
20 
21 import com.android.tradefed.device.DeviceNotAvailableException;
22 import com.android.tradefed.device.ITestDevice;
23 import com.android.tradefed.log.LogUtil;
24 
25 import org.junit.Assert;
26 import org.junit.ClassRule;
27 import org.junit.rules.ExternalResource;
28 import org.junit.rules.TemporaryFolder;
29 import org.junit.rules.TestRule;
30 import org.junit.runner.Description;
31 import org.junit.runners.model.Statement;
32 
33 import java.io.File;
34 import java.io.FileOutputStream;
35 import java.io.IOException;
36 import java.io.InputStream;
37 import java.util.ArrayList;
38 
39 import javax.annotation.Nullable;
40 
41 /**
42  * Allows pushing files onto the device and various options for rebooting. Useful for installing
43  * APKs/files to system partitions which otherwise wouldn't be easily changed.
44  *
45  * It's strongly recommended to pass in a {@link ClassRule} annotated {@link TestRuleDelegate} to
46  * do a full reboot at the end of a test to ensure the device is in a valid state, assuming the
47  * default {@link RebootStrategy#FULL} isn't used.
48  */
49 public class SystemPreparer extends ExternalResource {
50     private static final long OVERLAY_ENABLE_TIMEOUT_MS = 30000;
51 
52     // The paths of the files pushed onto the device through this rule.
53     private ArrayList<String> mPushedFiles = new ArrayList<>();
54 
55     // The package names of packages installed through this rule.
56     private ArrayList<String> mInstalledPackages = new ArrayList<>();
57 
58     private final TemporaryFolder mHostTempFolder;
59     private final DeviceProvider mDeviceProvider;
60     private final RebootStrategy mRebootStrategy;
61     private final TearDownRule mTearDownRule;
62 
SystemPreparer(TemporaryFolder hostTempFolder, DeviceProvider deviceProvider)63     public SystemPreparer(TemporaryFolder hostTempFolder, DeviceProvider deviceProvider) {
64         this(hostTempFolder, RebootStrategy.FULL, null, deviceProvider);
65     }
66 
SystemPreparer(TemporaryFolder hostTempFolder, RebootStrategy rebootStrategy, @Nullable TestRuleDelegate testRuleDelegate, DeviceProvider deviceProvider)67     public SystemPreparer(TemporaryFolder hostTempFolder, RebootStrategy rebootStrategy,
68             @Nullable TestRuleDelegate testRuleDelegate, DeviceProvider deviceProvider) {
69         mHostTempFolder = hostTempFolder;
70         mDeviceProvider = deviceProvider;
71         mRebootStrategy = rebootStrategy;
72         mTearDownRule = new TearDownRule(mDeviceProvider);
73         if (testRuleDelegate != null) {
74             testRuleDelegate.setDelegate(mTearDownRule);
75         }
76     }
77 
78     /** Copies a file within the host test jar to a path on device. */
pushResourceFile(String filePath, String outputPath)79     public SystemPreparer pushResourceFile(String filePath, String outputPath)
80             throws DeviceNotAvailableException, IOException {
81         final ITestDevice device = mDeviceProvider.getDevice();
82         remount();
83         assertTrue(device.pushFile(copyResourceToTemp(filePath), outputPath));
84         mPushedFiles.add(outputPath);
85         return this;
86     }
87 
88     /** Copies a file directly from the host file system to a path on device. */
pushFile(File file, String outputPath)89     public SystemPreparer pushFile(File file, String outputPath)
90             throws DeviceNotAvailableException {
91         final ITestDevice device = mDeviceProvider.getDevice();
92         remount();
93         assertTrue(device.pushFile(file, outputPath));
94         mPushedFiles.add(outputPath);
95         return this;
96     }
97 
98     /** Deletes the given path from the device */
deleteFile(String file)99     public SystemPreparer deleteFile(String file) throws DeviceNotAvailableException {
100         final ITestDevice device = mDeviceProvider.getDevice();
101         remount();
102         device.deleteFile(file);
103         return this;
104     }
105 
106     /** Installs an APK within the host test jar onto the device. */
installResourceApk(String resourcePath, String packageName)107     public SystemPreparer installResourceApk(String resourcePath, String packageName)
108             throws DeviceNotAvailableException, IOException {
109         final ITestDevice device = mDeviceProvider.getDevice();
110         final File tmpFile = copyResourceToTemp(resourcePath);
111         final String result = device.installPackage(tmpFile, true /* reinstall */);
112         Assert.assertNull(result);
113         mInstalledPackages.add(packageName);
114         return this;
115     }
116 
117     /** Sets the enable state of an overlay package. */
setOverlayEnabled(String packageName, boolean enabled)118     public SystemPreparer setOverlayEnabled(String packageName, boolean enabled)
119             throws DeviceNotAvailableException {
120         final ITestDevice device = mDeviceProvider.getDevice();
121         final String enable = enabled ? "enable" : "disable";
122 
123         // Wait for the overlay to change its enabled state.
124         final long endMillis = System.currentTimeMillis() + OVERLAY_ENABLE_TIMEOUT_MS;
125         String result;
126         while (System.currentTimeMillis() <= endMillis) {
127             device.executeShellCommand(String.format("cmd overlay %s %s", enable, packageName));
128             result = device.executeShellCommand("cmd overlay dump isenabled "
129                     + packageName);
130             if (((enabled) ? "true\n" : "false\n").equals(result)) {
131                 return this;
132             }
133 
134             try {
135                 Thread.sleep(200);
136             } catch (InterruptedException ignore) {
137             }
138         }
139 
140         throw new IllegalStateException(String.format("Failed to %s overlay %s:\n%s", enable,
141                 packageName, device.executeShellCommand("cmd overlay list")));
142     }
143 
144     /** Restarts the device and waits until after boot is completed. */
reboot()145     public SystemPreparer reboot() throws DeviceNotAvailableException {
146         ITestDevice device = mDeviceProvider.getDevice();
147         switch (mRebootStrategy) {
148             case FULL:
149                 device.reboot();
150                 break;
151             case UNTIL_ONLINE:
152                 device.rebootUntilOnline();
153                 break;
154             case USERSPACE:
155                 device.rebootUserspace();
156                 break;
157             case USERSPACE_UNTIL_ONLINE:
158                 device.rebootUserspaceUntilOnline();
159                 break;
160             case START_STOP:
161                 device.executeShellCommand("stop");
162                 device.executeShellCommand("start");
163                 ITestDevice.RecoveryMode cachedRecoveryMode = device.getRecoveryMode();
164                 device.setRecoveryMode(ITestDevice.RecoveryMode.ONLINE);
165 
166                 if (device.isEncryptionSupported()) {
167                     if (device.isDeviceEncrypted()) {
168                         LogUtil.CLog.e("Device is encrypted after userspace reboot!");
169                         device.unlockDevice();
170                     }
171                 }
172 
173                 device.setRecoveryMode(cachedRecoveryMode);
174                 device.waitForDeviceAvailable();
175                 break;
176         }
177         return this;
178     }
179 
remount()180     public SystemPreparer remount() throws DeviceNotAvailableException {
181         mTearDownRule.remount();
182         return this;
183     }
184 
185     /** Copies a file within the host test jar to a temporary file on the host machine. */
copyResourceToTemp(String resourcePath)186     private File copyResourceToTemp(String resourcePath) throws IOException {
187         final File tempFile = mHostTempFolder.newFile();
188         final ClassLoader classLoader = getClass().getClassLoader();
189         try (InputStream assetIs = classLoader.getResource(resourcePath).openStream();
190              FileOutputStream assetOs = new FileOutputStream(tempFile)) {
191             if (assetIs == null) {
192                 throw new IllegalStateException("Failed to find resource " + resourcePath);
193             }
194 
195             int b;
196             while ((b = assetIs.read()) >= 0) {
197                 assetOs.write(b);
198             }
199         }
200 
201         return tempFile;
202     }
203 
204     /** Removes installed packages and files that were pushed to the device. */
205     @Override
after()206     protected void after() {
207         final ITestDevice device = mDeviceProvider.getDevice();
208         try {
209             remount();
210             for (final String file : mPushedFiles) {
211                 device.deleteFile(file);
212             }
213             for (final String packageName : mInstalledPackages) {
214                 device.uninstallPackage(packageName);
215             }
216             reboot();
217         } catch (DeviceNotAvailableException e) {
218             Assert.fail(e.toString());
219         }
220     }
221 
222     /**
223      * A hacky workaround since {@link org.junit.AfterClass} and {@link ClassRule} require static
224      * members. Will defer assignment of the actual {@link TestRule} to execute until after any
225      * test case has been run.
226      *
227      * In effect, this makes the {@link ITestDevice} to be accessible after all test cases have
228      * been executed, allowing {@link ITestDevice#reboot()} to be used to fully restore the device.
229      */
230     public static class TestRuleDelegate implements TestRule {
231 
232         private boolean mThrowOnNull;
233 
234         @Nullable
235         private TestRule mTestRule;
236 
TestRuleDelegate(boolean throwOnNull)237         public TestRuleDelegate(boolean throwOnNull) {
238             mThrowOnNull = throwOnNull;
239         }
240 
setDelegate(TestRule testRule)241         public void setDelegate(TestRule testRule) {
242             mTestRule = testRule;
243         }
244 
245         @Override
apply(Statement base, Description description)246         public Statement apply(Statement base, Description description) {
247             if (mTestRule == null) {
248                 if (mThrowOnNull) {
249                     throw new IllegalStateException("TestRule delegate was not set");
250                 } else {
251                     return new Statement() {
252                         @Override
253                         public void evaluate() throws Throwable {
254                             base.evaluate();
255                         }
256                     };
257                 }
258             }
259 
260             Statement statement = mTestRule.apply(base, description);
261             mTestRule = null;
262             return statement;
263         }
264     }
265 
266     /**
267      * Forces a full reboot at the end of the test class to restore any device state.
268      */
269     private static class TearDownRule extends ExternalResource {
270 
271         private DeviceProvider mDeviceProvider;
272         private boolean mInitialized;
273         private boolean mWasVerityEnabled;
274         private boolean mWasAdbRoot;
275         private boolean mIsVerityEnabled;
276 
277         TearDownRule(DeviceProvider deviceProvider) {
278             mDeviceProvider = deviceProvider;
279         }
280 
281         @Override
282         protected void before() {
283             // This method will never be run
284         }
285 
286         @Override
287         protected void after() {
288             try {
289                 initialize();
290                 ITestDevice device = mDeviceProvider.getDevice();
291                 if (mWasVerityEnabled != mIsVerityEnabled) {
292                     device.executeShellCommand(
293                             mWasVerityEnabled ? "enable-verity" : "disable-verity");
294                 }
295                 device.reboot();
296                 if (!mWasAdbRoot) {
297                     device.disableAdbRoot();
298                 }
299             } catch (DeviceNotAvailableException e) {
300                 Assert.fail(e.toString());
301             }
302         }
303 
304         /**
305          * Remount is done inside this class so that the verity state can be tracked.
306          */
307         public void remount() throws DeviceNotAvailableException {
308             initialize();
309             ITestDevice device = mDeviceProvider.getDevice();
310             device.enableAdbRoot();
311             if (mIsVerityEnabled) {
312                 mIsVerityEnabled = false;
313                 device.executeShellCommand("disable-verity");
314                 device.reboot();
315             }
316             device.executeShellCommand("remount");
317             device.waitForDeviceAvailable();
318         }
319 
320         private void initialize() throws DeviceNotAvailableException {
321             if (mInitialized) {
322                 return;
323             }
324             mInitialized = true;
325             ITestDevice device = mDeviceProvider.getDevice();
326             mWasAdbRoot = device.isAdbRoot();
327             device.enableAdbRoot();
328             String veritySystem = device.getProperty("partition.system.verified");
329             String verityVendor = device.getProperty("partition.vendor.verified");
330             mWasVerityEnabled = (veritySystem != null && !veritySystem.isEmpty())
331                     || (verityVendor != null && !verityVendor.isEmpty());
332             mIsVerityEnabled = mWasVerityEnabled;
333         }
334     }
335 
336     public interface DeviceProvider {
337         ITestDevice getDevice();
338     }
339 
340     /**
341      * How to reboot the device. Ordered from slowest to fastest.
342      */
343     public enum RebootStrategy {
344         /** @see ITestDevice#reboot() */
345         FULL,
346 
347         /** @see ITestDevice#rebootUntilOnline() () */
348         UNTIL_ONLINE,
349 
350         /** @see ITestDevice#rebootUserspace() */
351         USERSPACE,
352 
353         /** @see ITestDevice#rebootUserspaceUntilOnline() () */
354         USERSPACE_UNTIL_ONLINE,
355 
356         /**
357          * Uses shell stop && start to "reboot" the device. May leave invalid state after each test.
358          * Whether this matters or not depends on what's being tested.
359          *
360          * TODO(b/159540015): There's a bug with this causing unnecessary disk space usage, which
361          *  can eventually lead to an insufficient storage space error.
362          */
363         START_STOP
364     }
365 }
366