Add Bitness enum to NativePoc.

This allows shorter, more readable code when poc bitness is set at runtime.
This is generally called as:

NativePoc.builder()
    .pocName(....)
    .bitness(is64bit ? ONLY64 : ONLY32)
    .build()
    .run(this);

Bug: 267681397
Test: run test for CVE_2022_20416 modified with this API.
Change-Id: If081e2202dbcaf2a43926f96575d8dfd84e9903f

1 files changed, 17 insertions, 0 deletions

diff --git a/libraries/sts-common-util/host-side/src/com/android/sts/common/NativePoc.java b/libraries/sts-common-util/host-side/src/com/android/sts/common/NativePoc.java
index 1aabaa990..1acbdd00b 100644
--- a/libraries/sts-common-util/host-side/src/com/android/sts/common/NativePoc.java
+++ b/libraries/sts-common-util/host-side/src/com/android/sts/common/NativePoc.java
@@ -85,6 +85,12 @@ public abstract class NativePoc {
                 .assumePocExitSuccess(true);
     }
 
+    public static enum Bitness {
+        AUTO, // push 32 or 64 bit version of PoC depending on device arch
+        ONLY32, // push only 32bit version of PoC
+        ONLY64 // push only 64bit version of PoC; raises error when running on 32bit-only device
+    }
+
     @AutoValue.Builder
     public abstract static class Builder {
         /** Name of executable to be uploaded and run. Do not include "_sts??" suffix. */
@@ -146,6 +152,17 @@ public abstract class NativePoc {
 
         abstract Builder only64(boolean value);
 
+        /** Force using 32bit or 64bit version of the native poc */
+        public Builder bitness(Bitness bitness) {
+            if (bitness == Bitness.ONLY32) {
+                return only32(true).only64(false);
+            }
+            if (bitness == Bitness.ONLY64) {
+                return only32(false).only64(true);
+            }
+            return only32(false).only64(false);
+        }
+
         /**
          * Function to run after the PoC finishes executing but before assertion or cleanups.
          *