1 /*
2  * Copyright (C) 2022 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.microdroid.test.device;
17 
18 import static android.content.pm.PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK;
19 
20 import static com.google.common.truth.Truth.assertThat;
21 import static com.google.common.truth.TruthJUnit.assume;
22 
23 import static org.junit.Assume.assumeFalse;
24 import static org.junit.Assume.assumeTrue;
25 
26 import android.app.Instrumentation;
27 import android.app.UiAutomation;
28 import android.content.Context;
29 import android.os.ParcelFileDescriptor;
30 import android.os.SystemProperties;
31 import android.system.Os;
32 import android.system.virtualmachine.VirtualMachine;
33 import android.system.virtualmachine.VirtualMachineCallback;
34 import android.system.virtualmachine.VirtualMachineConfig;
35 import android.system.virtualmachine.VirtualMachineException;
36 import android.system.virtualmachine.VirtualMachineManager;
37 import android.util.Log;
38 
39 import androidx.annotation.CallSuper;
40 import androidx.test.core.app.ApplicationProvider;
41 import androidx.test.platform.app.InstrumentationRegistry;
42 
43 import com.android.microdroid.test.common.DeviceProperties;
44 import com.android.microdroid.test.common.MetricsProcessor;
45 import com.android.microdroid.testservice.ITestService;
46 import com.android.virt.vm_attestation.testservice.IAttestationService;
47 import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
48 
49 import java.io.BufferedReader;
50 import java.io.ByteArrayOutputStream;
51 import java.io.File;
52 import java.io.IOException;
53 import java.io.InputStream;
54 import java.io.InputStreamReader;
55 import java.util.Arrays;
56 import java.util.Collections;
57 import java.util.HashSet;
58 import java.util.OptionalLong;
59 import java.util.Set;
60 import java.util.concurrent.CompletableFuture;
61 import java.util.concurrent.ExecutorService;
62 import java.util.concurrent.Executors;
63 import java.util.concurrent.TimeUnit;
64 
65 public abstract class MicrodroidDeviceTestBase {
66     private static final String TAG = "MicrodroidDeviceTestBase";
67     private final String MAX_PERFORMANCE_TASK_PROFILE = "CPUSET_SP_TOP_APP";
68 
69     protected static final String KERNEL_VERSION = SystemProperties.get("ro.kernel.version");
70     protected static final Set<String> SUPPORTED_GKI_VERSIONS =
71             Collections.unmodifiableSet(
72                     new HashSet(Arrays.asList("android14-6.1-pkvm_experimental")));
73 
isCuttlefish()74     public static boolean isCuttlefish() {
75         return getDeviceProperties().isCuttlefish();
76     }
77 
isCuttlefishArm64()78     private static boolean isCuttlefishArm64() {
79         return getDeviceProperties().isCuttlefishArm64();
80     }
81 
isHwasan()82     public static boolean isHwasan() {
83         return getDeviceProperties().isHwasan();
84     }
85 
isUserBuild()86     public static boolean isUserBuild() {
87         return getDeviceProperties().isUserBuild();
88     }
89 
getMetricPrefix()90     public static String getMetricPrefix() {
91         return MetricsProcessor.getMetricPrefix(getDeviceProperties().getMetricsTag());
92     }
93 
getDeviceProperties()94     private static DeviceProperties getDeviceProperties() {
95         return DeviceProperties.create(SystemProperties::get);
96     }
97 
grantPermission(String permission)98     protected final void grantPermission(String permission) {
99         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
100         UiAutomation uiAutomation = instrumentation.getUiAutomation();
101         uiAutomation.grantRuntimePermission(instrumentation.getContext().getPackageName(),
102                 permission);
103     }
104 
revokePermission(String permission)105     protected final void revokePermission(String permission) {
106         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
107         UiAutomation uiAutomation = instrumentation.getUiAutomation();
108         uiAutomation.revokeRuntimePermission(instrumentation.getContext().getPackageName(),
109                 permission);
110     }
111 
setMaxPerformanceTaskProfile()112     protected final void setMaxPerformanceTaskProfile() throws IOException {
113         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
114         UiAutomation uiAutomation = instrumentation.getUiAutomation();
115         String cmd = "settaskprofile " + Os.gettid() + " " + MAX_PERFORMANCE_TASK_PROFILE;
116         String out = runInShell(TAG, uiAutomation, cmd).trim();
117         String expect = "Profile " + MAX_PERFORMANCE_TASK_PROFILE + " is applied successfully!";
118         if (!expect.equals(out)) {
119             throw new IOException("Could not apply max performance task profile: " + out);
120         }
121     }
122 
123     private final Context mCtx = ApplicationProvider.getApplicationContext();
124     private boolean mProtectedVm;
125     private String mGki;
126 
getContext()127     protected Context getContext() {
128         return mCtx;
129     }
130 
getVirtualMachineManager()131     public VirtualMachineManager getVirtualMachineManager() {
132         return mCtx.getSystemService(VirtualMachineManager.class);
133     }
134 
newVmConfigBuilderWithPayloadConfig(String configPath)135     public VirtualMachineConfig.Builder newVmConfigBuilderWithPayloadConfig(String configPath) {
136         return new VirtualMachineConfig.Builder(mCtx)
137                 .setProtectedVm(mProtectedVm)
138                 .setOs(os())
139                 .setPayloadConfigPath(configPath);
140     }
141 
newVmConfigBuilderWithPayloadBinary(String binaryPath)142     public VirtualMachineConfig.Builder newVmConfigBuilderWithPayloadBinary(String binaryPath) {
143         return new VirtualMachineConfig.Builder(mCtx)
144                 .setProtectedVm(mProtectedVm)
145                 .setOs(os())
146                 .setPayloadBinaryName(binaryPath);
147     }
148 
isProtectedVm()149     protected final boolean isProtectedVm() {
150         return mProtectedVm;
151     }
152 
os()153     protected final String os() {
154         return mGki != null ? "microdroid_gki-" + mGki : "microdroid";
155     }
156 
157     /**
158      * Creates a new virtual machine, potentially removing an existing virtual machine with given
159      * name.
160      */
forceCreateNewVirtualMachine(String name, VirtualMachineConfig config)161     public VirtualMachine forceCreateNewVirtualMachine(String name, VirtualMachineConfig config)
162             throws VirtualMachineException {
163         final VirtualMachineManager vmm = getVirtualMachineManager();
164         deleteVirtualMachineIfExists(name);
165         return vmm.create(name, config);
166     }
167 
deleteVirtualMachineIfExists(String name)168     protected void deleteVirtualMachineIfExists(String name) throws VirtualMachineException {
169         VirtualMachineManager vmm = getVirtualMachineManager();
170         boolean deleteExisting;
171         try {
172             deleteExisting = vmm.get(name) != null;
173         } catch (VirtualMachineException e) {
174             // VM exists, i.e. there are some files for it, but they could not be successfully
175             // loaded.
176             deleteExisting = true;
177         }
178         if (deleteExisting) {
179             vmm.delete(name);
180         }
181     }
182 
prepareTestSetup(boolean protectedVm, String gki)183     public void prepareTestSetup(boolean protectedVm, String gki) {
184         assumeFeatureVirtualizationFramework();
185 
186         mProtectedVm = protectedVm;
187         mGki = gki;
188 
189         int capabilities = getVirtualMachineManager().getCapabilities();
190         if (protectedVm) {
191             assume().withMessage("Skip where protected VMs aren't supported")
192                     .that(capabilities & VirtualMachineManager.CAPABILITY_PROTECTED_VM)
193                     .isNotEqualTo(0);
194             assume().withMessage("Testing protected VMs on GSI isn't supported. b/272443823")
195                     .that(isGsi())
196                     .isFalse();
197         } else {
198             assume().withMessage("Skip where VMs aren't supported")
199                     .that(capabilities & VirtualMachineManager.CAPABILITY_NON_PROTECTED_VM)
200                     .isNotEqualTo(0);
201         }
202 
203         try {
204             assume().withMessage("Skip where requested OS \"" + os() + "\" isn't supported")
205                     .that(os())
206                     .isIn(getVirtualMachineManager().getSupportedOSList());
207         } catch (VirtualMachineException e) {
208             Log.e(TAG, "Error getting supported OS list", e);
209             throw new RuntimeException("Failed to get supported OS list.", e);
210         }
211     }
212 
assumeFeatureVirtualizationFramework()213     protected void assumeFeatureVirtualizationFramework() {
214         assume().withMessage("Device doesn't support AVF")
215                 .that(mCtx.getPackageManager().hasSystemFeature(FEATURE_VIRTUALIZATION_FRAMEWORK))
216                 .isTrue();
217         int vendorApiLevel = getVendorApiLevel();
218         boolean isGsi = isGsi();
219         Log.i(TAG, "isGsi = " + isGsi + ", vendor api level = " + vendorApiLevel);
220         assume().withMessage("GSI with vendor API level < 202404 may not support AVF")
221                 .that(isGsi && vendorApiLevel < 202404)
222                 .isFalse();
223     }
224 
225     protected boolean isGsi() {
226         return new File("/system/system_ext/etc/init/init.gsi.rc").exists();
227     }
228 
229     protected static int getVendorApiLevel() {
230         return SystemProperties.getInt("ro.board.api_level", 0);
231     }
232 
233     protected void assumeSupportedDevice() {
234         assume().withMessage("Skip on 5.4 kernel. b/218303240")
235                 .that(KERNEL_VERSION)
236                 .isNotEqualTo("5.4");
237 
238         // Cuttlefish on Arm 64 doesn't and cannot support any form of virtualization, so there's
239         // no point running any of these tests.
240         assume().withMessage("Virtualization not supported on Arm64 Cuttlefish. b/341889915")
241                 .that(isCuttlefishArm64())
242                 .isFalse();
243     }
244 
245     protected void assumeNoUpdatableVmSupport() throws VirtualMachineException {
246         assume().withMessage("Secretkeeper not supported")
247                 .that(getVirtualMachineManager().isUpdatableVmSupported())
248                 .isFalse();
249     }
250 
251     public abstract static class VmEventListener implements VirtualMachineCallback {
252         private ExecutorService mExecutorService = Executors.newSingleThreadExecutor();
253         private OptionalLong mVcpuStartedNanoTime = OptionalLong.empty();
254         private OptionalLong mKernelStartedNanoTime = OptionalLong.empty();
255         private OptionalLong mInitStartedNanoTime = OptionalLong.empty();
256         private OptionalLong mPayloadStartedNanoTime = OptionalLong.empty();
257         private StringBuilder mConsoleOutput = new StringBuilder();
258         private StringBuilder mLogOutput = new StringBuilder();
259         private boolean mProcessedBootTimeMetrics = false;
260 
261         private synchronized void processBootTimeMetrics(String log) {
262             if (!mVcpuStartedNanoTime.isPresent()) {
263                 mVcpuStartedNanoTime = OptionalLong.of(System.nanoTime());
264             }
265             if (log.contains("Starting payload...") && !mKernelStartedNanoTime.isPresent()) {
266                 mKernelStartedNanoTime = OptionalLong.of(System.nanoTime());
267             }
268             if (log.contains("Run /init as init process") && !mInitStartedNanoTime.isPresent()) {
269                 mInitStartedNanoTime = OptionalLong.of(System.nanoTime());
270             }
271             if (log.contains("microdroid_manager") && log.contains("executing main task")
272                     && !mPayloadStartedNanoTime.isPresent()) {
273                 mPayloadStartedNanoTime = OptionalLong.of(System.nanoTime());
274             }
275         }
276 
277         private void logVmOutputAndMonitorBootTimeMetrics(
278                 String tag, InputStream vmOutputStream, String name, StringBuilder result) {
279             mProcessedBootTimeMetrics = true;
280             new Thread(
281                             () -> {
282                                 try {
283                                     BufferedReader reader =
284                                             new BufferedReader(
285                                                     new InputStreamReader(vmOutputStream));
286                                     String line;
287                                     while ((line = reader.readLine()) != null
288                                             && !Thread.interrupted()) {
289                                         processBootTimeMetrics(line);
290                                         Log.i(tag, name + ": " + line);
291                                         result.append(line + "\n");
292                                     }
293                                 } catch (Exception e) {
294                                     Log.w(tag, name, e);
295                                 }
296                             })
297                     .start();
298         }
299 
runToFinish(String logTag, VirtualMachine vm)300         public void runToFinish(String logTag, VirtualMachine vm)
301                 throws VirtualMachineException, InterruptedException {
302             vm.setCallback(mExecutorService, this);
303             vm.run();
304             if (vm.getConfig().isVmOutputCaptured()) {
305                 logVmOutputAndMonitorBootTimeMetrics(
306                         logTag, vm.getConsoleOutput(), "Console", mConsoleOutput);
307                 logVmOutputAndMonitorBootTimeMetrics(logTag, vm.getLogOutput(), "Log", mLogOutput);
308             }
309             mExecutorService.awaitTermination(300, TimeUnit.SECONDS);
310         }
311 
getVcpuStartedNanoTime()312         public OptionalLong getVcpuStartedNanoTime() {
313             return mVcpuStartedNanoTime;
314         }
315 
getKernelStartedNanoTime()316         public OptionalLong getKernelStartedNanoTime() {
317             return mKernelStartedNanoTime;
318         }
319 
getInitStartedNanoTime()320         public OptionalLong getInitStartedNanoTime() {
321             return mInitStartedNanoTime;
322         }
323 
getPayloadStartedNanoTime()324         public OptionalLong getPayloadStartedNanoTime() {
325             return mPayloadStartedNanoTime;
326         }
327 
getConsoleOutput()328         public String getConsoleOutput() {
329             return mConsoleOutput.toString();
330         }
331 
getLogOutput()332         public String getLogOutput() {
333             return mLogOutput.toString();
334         }
335 
hasProcessedBootTimeMetrics()336         public boolean hasProcessedBootTimeMetrics() {
337             return mProcessedBootTimeMetrics;
338         }
339 
forceStop(VirtualMachine vm)340         protected void forceStop(VirtualMachine vm) {
341             try {
342                 vm.stop();
343             } catch (VirtualMachineException e) {
344                 throw new RuntimeException(e);
345             }
346         }
347 
348         @Override
onPayloadStarted(VirtualMachine vm)349         public void onPayloadStarted(VirtualMachine vm) {}
350 
351         @Override
onPayloadReady(VirtualMachine vm)352         public void onPayloadReady(VirtualMachine vm) {}
353 
354         @Override
onPayloadFinished(VirtualMachine vm, int exitCode)355         public void onPayloadFinished(VirtualMachine vm, int exitCode) {}
356 
357         @Override
onError(VirtualMachine vm, int errorCode, String message)358         public void onError(VirtualMachine vm, int errorCode, String message) {}
359 
360         @Override
361         @CallSuper
onStopped(VirtualMachine vm, int reason)362         public void onStopped(VirtualMachine vm, int reason) {
363             vm.clearCallback();
364             mExecutorService.shutdown();
365         }
366     }
367 
368     public enum BootTimeMetric {
369         TOTAL,
370         VM_START,
371         BOOTLOADER,
372         KERNEL,
373         USERSPACE,
374     }
375 
376     public static class BootResult {
377         public final boolean payloadStarted;
378         public final int deathReason;
379         public final long apiCallNanoTime;
380         public final long endToEndNanoTime;
381 
382         public final boolean processedBootTimeMetrics;
383         public final OptionalLong vcpuStartedNanoTime;
384         public final OptionalLong kernelStartedNanoTime;
385         public final OptionalLong initStartedNanoTime;
386         public final OptionalLong payloadStartedNanoTime;
387 
388         public final String consoleOutput;
389         public final String logOutput;
390 
BootResult( boolean payloadStarted, int deathReason, long apiCallNanoTime, long endToEndNanoTime, boolean processedBootTimeMetrics, OptionalLong vcpuStartedNanoTime, OptionalLong kernelStartedNanoTime, OptionalLong initStartedNanoTime, OptionalLong payloadStartedNanoTime, String consoleOutput, String logOutput)391         BootResult(
392                 boolean payloadStarted,
393                 int deathReason,
394                 long apiCallNanoTime,
395                 long endToEndNanoTime,
396                 boolean processedBootTimeMetrics,
397                 OptionalLong vcpuStartedNanoTime,
398                 OptionalLong kernelStartedNanoTime,
399                 OptionalLong initStartedNanoTime,
400                 OptionalLong payloadStartedNanoTime,
401                 String consoleOutput,
402                 String logOutput) {
403             this.apiCallNanoTime = apiCallNanoTime;
404             this.payloadStarted = payloadStarted;
405             this.deathReason = deathReason;
406             this.endToEndNanoTime = endToEndNanoTime;
407             this.processedBootTimeMetrics = processedBootTimeMetrics;
408             this.vcpuStartedNanoTime = vcpuStartedNanoTime;
409             this.kernelStartedNanoTime = kernelStartedNanoTime;
410             this.initStartedNanoTime = initStartedNanoTime;
411             this.payloadStartedNanoTime = payloadStartedNanoTime;
412             this.consoleOutput = consoleOutput;
413             this.logOutput = logOutput;
414         }
415 
getVcpuStartedNanoTime()416         private long getVcpuStartedNanoTime() {
417             return vcpuStartedNanoTime.getAsLong();
418         }
419 
getKernelStartedNanoTime()420         private long getKernelStartedNanoTime() {
421             // pvmfw emits log at the end which is used to estimate the kernelStart time.
422             // In case of no pvmfw run(non-protected mode), use vCPU started time instead.
423             return kernelStartedNanoTime.orElse(vcpuStartedNanoTime.getAsLong());
424         }
425 
getInitStartedNanoTime()426         private long getInitStartedNanoTime() {
427             return initStartedNanoTime.getAsLong();
428         }
429 
getPayloadStartedNanoTime()430         private long getPayloadStartedNanoTime() {
431             return payloadStartedNanoTime.getAsLong();
432         }
433 
getVMStartingElapsedNanoTime()434         public long getVMStartingElapsedNanoTime() {
435             return getVcpuStartedNanoTime() - apiCallNanoTime;
436         }
437 
getBootloaderElapsedNanoTime()438         public long getBootloaderElapsedNanoTime() {
439             return getKernelStartedNanoTime() - getVcpuStartedNanoTime();
440         }
441 
getKernelElapsedNanoTime()442         public long getKernelElapsedNanoTime() {
443             return getInitStartedNanoTime() - getKernelStartedNanoTime();
444         }
445 
getUserspaceElapsedNanoTime()446         public long getUserspaceElapsedNanoTime() {
447             return getPayloadStartedNanoTime() - getInitStartedNanoTime();
448         }
449 
getBootTimeMetricNanoTime(BootTimeMetric metric)450         public OptionalLong getBootTimeMetricNanoTime(BootTimeMetric metric) {
451             if (metric == BootTimeMetric.TOTAL) {
452                 return OptionalLong.of(endToEndNanoTime);
453             }
454 
455             if (processedBootTimeMetrics) {
456                 switch (metric) {
457                     case VM_START:
458                         return OptionalLong.of(getVMStartingElapsedNanoTime());
459                     case BOOTLOADER:
460                         return OptionalLong.of(getBootloaderElapsedNanoTime());
461                     case KERNEL:
462                         return OptionalLong.of(getKernelElapsedNanoTime());
463                     case USERSPACE:
464                         return OptionalLong.of(getUserspaceElapsedNanoTime());
465                 }
466             }
467 
468             return OptionalLong.empty();
469         }
470     }
471 
tryBootVm(String logTag, String vmName)472     public BootResult tryBootVm(String logTag, String vmName)
473             throws VirtualMachineException, InterruptedException {
474         VirtualMachine vm = getVirtualMachineManager().get(vmName);
475         final CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
476         final CompletableFuture<Integer> deathReason = new CompletableFuture<>();
477         final CompletableFuture<Long> endTime = new CompletableFuture<>();
478         VmEventListener listener =
479                 new VmEventListener() {
480                     @Override
481                     public void onPayloadStarted(VirtualMachine vm) {
482                         endTime.complete(System.nanoTime());
483                         payloadStarted.complete(true);
484                         forceStop(vm);
485                     }
486 
487                     @Override
488                     public void onStopped(VirtualMachine vm, int reason) {
489                         deathReason.complete(reason);
490                         super.onStopped(vm, reason);
491                     }
492                 };
493         long apiCallNanoTime = System.nanoTime();
494         listener.runToFinish(logTag, vm);
495         return new BootResult(
496                 payloadStarted.getNow(false),
497                 deathReason.getNow(VmEventListener.STOP_REASON_INFRASTRUCTURE_ERROR),
498                 apiCallNanoTime,
499                 endTime.getNow(apiCallNanoTime) - apiCallNanoTime,
500                 listener.hasProcessedBootTimeMetrics(),
501                 listener.getVcpuStartedNanoTime(),
502                 listener.getKernelStartedNanoTime(),
503                 listener.getInitStartedNanoTime(),
504                 listener.getPayloadStartedNanoTime(),
505                 listener.getConsoleOutput(),
506                 listener.getLogOutput());
507     }
508 
509     /** Execute a command. Returns stdout. */
runInShell(String tag, UiAutomation uiAutomation, String command)510     protected String runInShell(String tag, UiAutomation uiAutomation, String command) {
511         try (InputStream is =
512                         new ParcelFileDescriptor.AutoCloseInputStream(
513                                 uiAutomation.executeShellCommand(command));
514                 ByteArrayOutputStream out = new ByteArrayOutputStream()) {
515             is.transferTo(out);
516             String stdout = out.toString("UTF-8");
517             Log.i(tag, "Got stdout : " + stdout);
518             return stdout;
519         } catch (IOException e) {
520             Log.e(tag, "Error executing: " + command, e);
521             throw new RuntimeException("Failed to run the command.", e);
522         }
523     }
524 
525     /** Execute a command. Returns the concatenation of stdout and stderr. */
runInShellWithStderr(String tag, UiAutomation uiAutomation, String command)526     protected String runInShellWithStderr(String tag, UiAutomation uiAutomation, String command) {
527         ParcelFileDescriptor[] files = uiAutomation.executeShellCommandRwe(command);
528         try (InputStream stdout = new ParcelFileDescriptor.AutoCloseInputStream(files[0]);
529                 InputStream stderr = new ParcelFileDescriptor.AutoCloseInputStream(files[2]);
530                 ByteArrayOutputStream out = new ByteArrayOutputStream()) {
531             files[1].close(); // The command's stdin
532             stdout.transferTo(out);
533             stderr.transferTo(out);
534             String output = out.toString("UTF-8");
535             Log.i(tag, "Got output : " + stdout);
536             return output;
537         } catch (IOException e) {
538             Log.e(tag, "Error executing: " + command, e);
539             throw new RuntimeException("Failed to run the command.", e);
540         }
541     }
542 
543     protected static class TestResults {
544         public Exception mException;
545         public Integer mAddInteger;
546         public String mAppRunProp;
547         public String mSublibRunProp;
548         public String mExtraApkTestProp;
549         public String mApkContentsPath;
550         public String mEncryptedStoragePath;
551         public String[] mEffectiveCapabilities;
552         public int mUid;
553         public String mFileContent;
554         public byte[] mBcc;
555         public long[] mTimings;
556         public int mFileMode;
557         public int mMountFlags;
558         public String mConsoleInput;
559 
assertNoException()560         public void assertNoException() {
561             if (mException != null) {
562                 // Rethrow, wrapped in a new exception, so we get stack traces of the original
563                 // failure as well as the body of the test.
564                 throw new RuntimeException(mException);
565             }
566         }
567     }
568 
runVmAttestationService( String logTag, VirtualMachine vm, byte[] challenge, byte[] messageToSign)569     protected SigningResult runVmAttestationService(
570             String logTag, VirtualMachine vm, byte[] challenge, byte[] messageToSign)
571             throws Exception {
572 
573         CompletableFuture<Exception> exception = new CompletableFuture<>();
574         CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
575         CompletableFuture<SigningResult> signingResultFuture = new CompletableFuture<>();
576         VmEventListener listener =
577                 new VmEventListener() {
578                     @Override
579                     public void onPayloadReady(VirtualMachine vm) {
580                         payloadReady.complete(true);
581                         try {
582                             IAttestationService service =
583                                     IAttestationService.Stub.asInterface(
584                                             vm.connectToVsockServer(IAttestationService.PORT));
585                             signingResultFuture.complete(
586                                     service.signWithAttestationKey(challenge, messageToSign));
587                         } catch (Exception e) {
588                             exception.complete(e);
589                         } finally {
590                             forceStop(vm);
591                         }
592                     }
593                 };
594         listener.runToFinish(TAG, vm);
595 
596         assertThat(payloadReady.getNow(false)).isTrue();
597         assertThat(exception.getNow(null)).isNull();
598         SigningResult signingResult = signingResultFuture.getNow(null);
599         assertThat(signingResult).isNotNull();
600         return signingResult;
601     }
602 
runVmTestService( String logTag, VirtualMachine vm, RunTestsAgainstTestService testsToRun)603     protected TestResults runVmTestService(
604             String logTag, VirtualMachine vm, RunTestsAgainstTestService testsToRun)
605             throws Exception {
606         CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
607         CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
608         CompletableFuture<Boolean> payloadFinished = new CompletableFuture<>();
609         TestResults testResults = new TestResults();
610         VmEventListener listener =
611                 new VmEventListener() {
612                     ITestService mTestService = null;
613 
614                     private void initializeTestService(VirtualMachine vm) {
615                         try {
616                             mTestService =
617                                     ITestService.Stub.asInterface(
618                                             vm.connectToVsockServer(ITestService.PORT));
619                             // Make sure linkToDeath works, and include it in the log in case it's
620                             // helpful.
621                             mTestService
622                                     .asBinder()
623                                     .linkToDeath(
624                                             () -> Log.i(logTag, "ITestService binder died"), 0);
625                         } catch (Exception e) {
626                             testResults.mException = e;
627                         }
628                     }
629 
630                     private void testVMService(VirtualMachine vm) {
631                         try {
632                             if (mTestService == null) initializeTestService(vm);
633                             testsToRun.runTests(mTestService, testResults);
634                         } catch (Exception e) {
635                             testResults.mException = e;
636                         }
637                     }
638 
639                     private void quitVMService() {
640                         try {
641                             mTestService.quit();
642                         } catch (Exception e) {
643                             testResults.mException = e;
644                         }
645                     }
646 
647                     @Override
648                     public void onPayloadReady(VirtualMachine vm) {
649                         Log.i(logTag, "onPayloadReady");
650                         payloadReady.complete(true);
651                         testVMService(vm);
652                         quitVMService();
653                     }
654 
655                     @Override
656                     public void onPayloadStarted(VirtualMachine vm) {
657                         Log.i(logTag, "onPayloadStarted");
658                         payloadStarted.complete(true);
659                     }
660 
661                     @Override
662                     public void onPayloadFinished(VirtualMachine vm, int exitCode) {
663                         Log.i(logTag, "onPayloadFinished: " + exitCode);
664                         payloadFinished.complete(true);
665                         forceStop(vm);
666                     }
667                 };
668 
669         listener.runToFinish(logTag, vm);
670         assertThat(payloadStarted.getNow(false)).isTrue();
671         assertThat(payloadReady.getNow(false)).isTrue();
672         assertThat(payloadFinished.getNow(false)).isTrue();
673         return testResults;
674     }
675 
676     @FunctionalInterface
677     protected interface RunTestsAgainstTestService {
runTests(ITestService testService, TestResults testResults)678         void runTests(ITestService testService, TestResults testResults) throws Exception;
679     }
680 
assumeFeatureEnabled(String featureName)681     protected void assumeFeatureEnabled(String featureName) throws Exception {
682         assumeTrue(featureName + " not enabled", isFeatureEnabled(featureName));
683     }
684 
isFeatureEnabled(String featureName)685     protected boolean isFeatureEnabled(String featureName) throws Exception {
686         return getVirtualMachineManager().isFeatureEnabled(featureName);
687     }
688 
assumeProtectedVM()689     protected void assumeProtectedVM() {
690         assumeTrue("Skip on non-protected VM", mProtectedVm);
691     }
692 
assumeNonProtectedVM()693     protected void assumeNonProtectedVM() {
694         assumeFalse("Skip on protected VM", mProtectedVm);
695     }
696 }
697