From 4912a244470b0afb98e4fe75b6661c47cecb7991 Mon Sep 17 00:00:00 2001
From: Inseob Kim <inseob@google.com>
Date: Mon, 25 Jul 2022 11:30:02 +0900
Subject: [PATCH] Use embedded launcher for python binaries

Bug: 239386651
Test: m selinux_policy
Change-Id: Ic267fcfe4c38b51f8cf2469157b7cb57b84ad779
---
 tests/Android.bp               | 10 ++++++++++
 tests/sepolicy_tests.py        | 28 ++++++++++++++++++++++------
 tests/treble_sepolicy_tests.py | 30 ++++++++++++++++++++++++------
 3 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/tests/Android.bp b/tests/Android.bp
index 8ca952dcd..e2713461c 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -43,6 +43,11 @@ python_binary_host {
     srcs: [
         "treble_sepolicy_tests.py",
     ],
+    version: {
+        py3: {
+            embedded_launcher: true,
+        },
+    },
     libs: [
         "mini_cil_parser",
         "pysepolwrap",
@@ -55,6 +60,11 @@ python_binary_host {
     srcs: [
         "sepolicy_tests.py",
     ],
+    version: {
+        py3: {
+            embedded_launcher: true,
+        },
+    },
     libs: ["pysepolwrap"],
     data: [":libsepolwrap"],
 }
diff --git a/tests/sepolicy_tests.py b/tests/sepolicy_tests.py
index e940681ad..63144dda4 100644
--- a/tests/sepolicy_tests.py
+++ b/tests/sepolicy_tests.py
@@ -15,9 +15,12 @@
 from optparse import OptionParser
 from optparse import Option, OptionValueError
 import os
+import pkgutil
 import policy
 import re
+import shutil
 import sys
+import tempfile
 
 SHARED_LIB_EXTENSION = '.dylib' if sys.platform == 'darwin' else '.so'
 
@@ -146,7 +149,11 @@ Tests = [
     "TestDmaHeapDevTypeViolations",
 ]
 
-if __name__ == '__main__':
+def do_main(libpath):
+    """
+    Args:
+        libpath: string, path to libsepolwrap.so
+    """
     usage = "sepolicy_tests -f vendor_file_contexts -f "
     usage +="plat_file_contexts -p policy [--test test] [--help]"
     parser = OptionParser(option_class=MultipleOption, usage=usage)
@@ -158,11 +165,6 @@ if __name__ == '__main__':
 
     (options, args) = parser.parse_args()
 
-    libpath = os.path.join(os.path.dirname(os.path.realpath(__file__)),
-                           "libsepolwrap" + SHARED_LIB_EXTENSION)
-    if not os.path.exists(libpath):
-        sys.exit("Error: libsepolwrap does not exist. Is this binary corrupted?\n")
-
     if not options.policy:
         sys.exit("Must specify monolithic policy file\n" + parser.usage)
     if not os.path.exists(options.policy):
@@ -207,3 +209,17 @@ if __name__ == '__main__':
 
     if len(results) > 0:
         sys.exit(results)
+
+if __name__ == '__main__':
+    temp_dir = tempfile.mkdtemp()
+    try:
+        libname = "libsepolwrap" + SHARED_LIB_EXTENSION
+        libpath = os.path.join(temp_dir, libname)
+        with open(libpath, "wb") as f:
+            blob = pkgutil.get_data("sepolicy_tests", libname)
+            if not blob:
+                sys.exit("Error: libsepolwrap does not exist. Is this binary corrupted?\n")
+            f.write(blob)
+        do_main(libpath)
+    finally:
+        shutil.rmtree(temp_dir)
diff --git a/tests/treble_sepolicy_tests.py b/tests/treble_sepolicy_tests.py
index 64a9e950a..b49f138a7 100644
--- a/tests/treble_sepolicy_tests.py
+++ b/tests/treble_sepolicy_tests.py
@@ -16,10 +16,13 @@ from optparse import OptionParser
 from optparse import Option, OptionValueError
 import os
 import mini_parser
+import pkgutil
 import policy
 from policy import MatchPathPrefix
 import re
+import shutil
 import sys
+import tempfile
 
 DEBUG=False
 SHARED_LIB_EXTENSION = '.dylib' if sys.platform == 'darwin' else '.so'
@@ -341,7 +344,13 @@ Tests = {"CoredomainViolations": TestCoredomainViolations,
          "TrebleCompatMapping": TestTrebleCompatMapping,
          "ViolatorAttributes": TestViolatorAttributes}
 
-if __name__ == '__main__':
+def do_main(libpath):
+    """
+    Args:
+        libpath: string, path to libsepolwrap.so
+    """
+    global pol, FakeTreble
+
     usage = "treble_sepolicy_tests "
     usage += "-f nonplat_file_contexts -f plat_file_contexts "
     usage += "-p curr_policy -b base_policy -o old_policy "
@@ -374,11 +383,6 @@ if __name__ == '__main__':
             sys.exit("Error: File_contexts file " + f + " does not exist\n" +
                     parser.usage)
 
-    libpath = os.path.join(os.path.dirname(os.path.realpath(__file__)),
-                           "libsepolwrap" + SHARED_LIB_EXTENSION)
-    if not os.path.exists(libpath):
-        sys.exit("Error: libsepolwrap does not exist. Is this binary corrupted?\n")
-
     # Mapping files and public platform policy are only necessary for the
     # TrebleCompatMapping test.
     if options.tests is None or options.tests == "TrebleCompatMapping":
@@ -428,3 +432,17 @@ if __name__ == '__main__':
 
     if len(results) > 0:
         sys.exit(results)
+
+if __name__ == '__main__':
+    temp_dir = tempfile.mkdtemp()
+    try:
+        libname = "libsepolwrap" + SHARED_LIB_EXTENSION
+        libpath = os.path.join(temp_dir, libname)
+        with open(libpath, "wb") as f:
+            blob = pkgutil.get_data("treble_sepolicy_tests", libname)
+            if not blob:
+                sys.exit("Error: libsepolwrap does not exist. Is this binary corrupted?\n")
+            f.write(blob)
+        do_main(libpath)
+    finally:
+        shutil.rmtree(temp_dir)