Tests to ensure native and java WorkSource are in sync

Bug: b/237719490
Test: atest WorkSourceParcelTest, atest PowerManagerTest
Change-Id: Ia4bf8ad292e067e4e39b0b8627ef0d31105ba5c1

7 files changed, 425 insertions, 228 deletions

diff --git a/core/tests/coretests/Android.bp b/core/tests/coretests/Android.bp
index 34b4c5177852..67842f05165c 100644
--- a/core/tests/coretests/Android.bp
+++ b/core/tests/coretests/Android.bp
@@ -73,6 +73,7 @@ android_test {
     ],
     jni_libs: [
         "libpowermanagertest_jni",
+        "libworksourceparceltest_jni",
     ],
 
     sdk_version: "core_platform",
diff --git a/core/tests/coretests/jni/Android.bp b/core/tests/coretests/jni/Android.bp
index edac8ef25fe7..7cc844ae9fc6 100644
--- a/core/tests/coretests/jni/Android.bp
+++ b/core/tests/coretests/jni/Android.bp
@@ -43,3 +43,27 @@ cc_test_library {
     ],
     gtest: false,
 }
+
+cc_test_library {
+    name: "libworksourceparceltest_jni",
+    srcs: [
+        "NativeWorkSourceParcelTest.cpp",
+    ],
+    shared_libs: [
+        "libandroid",
+        "libandroid_runtime_lazy",
+        "libbase",
+        "libbinder",
+        "liblog",
+        "libnativehelper",
+        "libpowermanager",
+        "libutils",
+    ],
+    header_libs: ["jni_headers"],
+    stl: "libc++_static",
+    cflags: [
+        "-Werror",
+        "-Wall",
+    ],
+    gtest: false,
+}
diff --git a/core/tests/coretests/jni/NativePowerManagerTest.cpp b/core/tests/coretests/jni/NativePowerManagerTest.cpp
index 5f20e4f9c407..c15c0cefcb08 100644
--- a/core/tests/coretests/jni/NativePowerManagerTest.cpp
+++ b/core/tests/coretests/jni/NativePowerManagerTest.cpp
@@ -18,6 +18,7 @@
 #define LOG_TAG "NativePowerManagerTest"
 
 #include "jni.h"
+#include "ParcelHelper.h"
 
 #include <android_util_Binder.h>
 #include <binder/IServiceManager.h>
@@ -36,21 +37,6 @@ using android::base::StringPrintf;
 
 namespace android {
 
-#define FIND_CLASS(var, className) \
-    var = env->FindClass(className); \
-    LOG_FATAL_IF(!(var), "Unable to find class %s", className);
-
-#define GET_FIELD_ID(var, clazz, fieldName, fieldDescriptor) \
-    var = env->GetFieldID(clazz, fieldName, fieldDescriptor); \
-    LOG_FATAL_IF(!(var), "Unable to find field %s", fieldName);
-
-#define GET_STATIC_METHOD_ID(var, clazz, fieldName, fieldDescriptor) \
-    var = env->GetStaticMethodID(clazz, fieldName, fieldDescriptor); \
-    LOG_FATAL_IF(!(var), "Unable to find method %s", fieldName);
-
-static jclass gParcelClazz;
-static jfieldID gParcelDataFieldID;
-static jmethodID gParcelObtainMethodID;
 static struct BatterySaverPolicyConfigFieldId {
     jfieldID adjustBrightnessFactor;
     jfieldID advertiseIsEnabled;
@@ -73,102 +59,6 @@ static struct BatterySaverPolicyConfigFieldId {
     jfieldID soundTriggerMode;
 } gBSPCFieldIds;
 
-static jobject nativeObtainParcel(JNIEnv* env) {
-    jobject parcel = env->CallStaticObjectMethod(gParcelClazz, gParcelObtainMethodID);
-    if (parcel == nullptr) {
-        jniThrowException(env, "java/lang/IllegalArgumentException", "Obtain parcel failed.");
-    }
-    return parcel;
-}
-
-static Parcel* nativeGetParcelData(JNIEnv* env, jobject obj) {
-    Parcel* parcel = reinterpret_cast<Parcel*>(env->GetLongField(obj, gParcelDataFieldID));
-    if (parcel && parcel->objectsCount() != 0) {
-        jniThrowException(env, "java/lang/IllegalArgumentException", "Invalid parcel object.");
-    }
-    parcel->setDataPosition(0);
-    return parcel;
-}
-
-static jobject nativeObtainWorkSourceParcel(JNIEnv* env, jobject /* obj */, jintArray uidArray,
-            jobjectArray nameArray) {
-    std::vector<int32_t> uids;
-    std::optional<std::vector<std::optional<String16>>> names = std::nullopt;
-
-    if (uidArray != nullptr) {
-        jint *ptr = env->GetIntArrayElements(uidArray, 0);
-        for (jint i = 0; i < env->GetArrayLength(uidArray); i++) {
-            uids.push_back(static_cast<int32_t>(ptr[i]));
-        }
-    }
-
-    if (nameArray != nullptr) {
-        std::vector<std::optional<String16>> namesVec;
-        for (jint i = 0; i < env->GetArrayLength(nameArray); i++) {
-            jstring string = (jstring) (env->GetObjectArrayElement(nameArray, i));
-            const char *rawString = env->GetStringUTFChars(string, 0);
-            namesVec.push_back(std::make_optional<String16>(String16(rawString)));
-        }
-        names = std::make_optional(std::move(namesVec));
-    }
-
-    WorkSource ws = WorkSource(uids, names);
-    jobject wsParcel = nativeObtainParcel(env);
-    Parcel* parcel = nativeGetParcelData(env, wsParcel);
-    status_t err = ws.writeToParcel(parcel);
-    if (err != OK) {
-        jniThrowException(env, "java/lang/IllegalArgumentException",
-                            StringPrintf("WorkSource writeToParcel failed %d", err).c_str());
-    }
-    parcel->setDataPosition(0);
-    return wsParcel;
-}
-
-static void nativeUnparcelAndVerifyWorkSource(JNIEnv* env, jobject /* obj */, jobject wsParcel,
-        jintArray uidArray, jobjectArray nameArray) {
-    WorkSource ws = {};
-    Parcel* parcel = nativeGetParcelData(env, wsParcel);
-
-    status_t err = ws.readFromParcel(parcel);
-    if (err != OK) {
-        ALOGE("WorkSource writeToParcel failed %d", err);
-    }
-
-    // Now we have a native WorkSource object, verify it.
-    if (uidArray != nullptr) {
-        jint *ptr = env->GetIntArrayElements(uidArray, 0);
-        for (jint i = 0; i < env->GetArrayLength(uidArray); i++) {
-            if (ws.getUids().at(i) != static_cast<int32_t>(ptr[i])) {
-                jniThrowException(env, "java/lang/IllegalArgumentException",
-                            StringPrintf("WorkSource uid not equal %d %d",
-                            ws.getUids().at(i), static_cast<int32_t>(ptr[i])).c_str());
-            }
-        }
-    } else {
-        if (ws.getUids().size() != 0) {
-            jniThrowException(env, "java/lang/IllegalArgumentException",
-                    StringPrintf("WorkSource parcel size not 0").c_str());
-        }
-    }
-
-    if (nameArray != nullptr) {
-        std::vector<std::optional<String16>> namesVec;
-        for (jint i = 0; i < env->GetArrayLength(nameArray); i++) {
-            jstring string = (jstring) (env->GetObjectArrayElement(nameArray, i));
-            const char *rawString = env->GetStringUTFChars(string, 0);
-            if (String16(rawString) != ws.getNames()->at(i)) {
-                jniThrowException(env, "java/lang/IllegalArgumentException",
-                            StringPrintf("WorkSource uid not equal %s", rawString).c_str());
-            }
-        }
-    } else {
-        if (ws.getNames() != std::nullopt) {
-            jniThrowException(env, "java/lang/IllegalArgumentException",
-                    StringPrintf("WorkSource parcel name not empty").c_str());
-        }
-    }
-}
-
 static jobject nativeObtainPowerSaveStateParcel(JNIEnv* env, jobject /* obj */,
         jboolean batterySaverEnabled, jboolean globalBatterySaverEnabled,
         jint locationMode, jint soundTriggerMode, jfloat brightnessFactor) {
@@ -305,10 +195,6 @@ extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */)
     JNIEnv* env;
     const JNINativeMethod methodTable[] = {
         /* name, signature, funcPtr */
-        { "nativeObtainWorkSourceParcel", "([I[Ljava/lang/String;)Landroid/os/Parcel;",
-                (void*) nativeObtainWorkSourceParcel },
-        { "nativeUnparcelAndVerifyWorkSource", "(Landroid/os/Parcel;[I[Ljava/lang/String;)V",
-                (void*) nativeUnparcelAndVerifyWorkSource },
         { "nativeObtainPowerSaveStateParcel", "(ZZIIF)Landroid/os/Parcel;",
                 (void*) nativeObtainPowerSaveStateParcel },
         { "nativeUnparcelAndVerifyPowerSaveState", "(Landroid/os/Parcel;ZZIIF)V",
@@ -327,34 +213,40 @@ extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */)
         return JNI_ERR;
     }
 
-    jclass bspcClazz;
-    FIND_CLASS(gParcelClazz, "android/os/Parcel");
-    GET_FIELD_ID(gParcelDataFieldID, gParcelClazz, "mNativePtr", "J");
-    GET_STATIC_METHOD_ID(gParcelObtainMethodID, gParcelClazz, "obtain", "()Landroid/os/Parcel;");
-    FIND_CLASS(bspcClazz, "android/os/BatterySaverPolicyConfig");
-    GET_FIELD_ID(gBSPCFieldIds.adjustBrightnessFactor, bspcClazz, "mAdjustBrightnessFactor", "F");
-    GET_FIELD_ID(gBSPCFieldIds.advertiseIsEnabled, bspcClazz, "mAdvertiseIsEnabled", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.deferFullBackup, bspcClazz, "mDeferFullBackup", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.deferKeyValueBackup, bspcClazz, "mDeferKeyValueBackup", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.deviceSpecificSettings, bspcClazz, "mDeviceSpecificSettings",
-                    "Ljava/util/Map;");
-    GET_FIELD_ID(gBSPCFieldIds.disableAnimation, bspcClazz, "mDisableAnimation", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.disableAod, bspcClazz, "mDisableAod", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.disableLaunchBoost, bspcClazz, "mDisableLaunchBoost", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.disableOptionalSensors, bspcClazz, "mDisableOptionalSensors", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.disableVibration, bspcClazz, "mDisableVibration", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.enableAdjustBrightness, bspcClazz, "mEnableAdjustBrightness", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.enableDataSaver, bspcClazz, "mEnableDataSaver", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.enableFirewall, bspcClazz, "mEnableFirewall", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.enableNightMode, bspcClazz, "mEnableNightMode", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.enableQuickDoze, bspcClazz, "mEnableQuickDoze", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.forceAllAppsStandby, bspcClazz, "mForceAllAppsStandby", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.forceBackgroundCheck, bspcClazz, "mForceBackgroundCheck", "Z");
-    GET_FIELD_ID(gBSPCFieldIds.locationMode, bspcClazz, "mLocationMode", "I");
-    GET_FIELD_ID(gBSPCFieldIds.soundTriggerMode, bspcClazz, "mSoundTriggerMode", "I");
+    loadParcelClass(env);
+
+    jclass bspcClazz = FindClassOrDie(env, "android/os/BatterySaverPolicyConfig");
+
+    gBSPCFieldIds.adjustBrightnessFactor =
+            GetFieldIDOrDie(env, bspcClazz, "mAdjustBrightnessFactor", "F");
+    gBSPCFieldIds.advertiseIsEnabled = GetFieldIDOrDie(env, bspcClazz, "mAdvertiseIsEnabled", "Z");
+    gBSPCFieldIds.deferFullBackup = GetFieldIDOrDie(env, bspcClazz, "mDeferFullBackup", "Z");
+    gBSPCFieldIds.deferKeyValueBackup =
+            GetFieldIDOrDie(env, bspcClazz, "mDeferKeyValueBackup", "Z");
+    gBSPCFieldIds.deviceSpecificSettings =
+            GetFieldIDOrDie(env, bspcClazz, "mDeviceSpecificSettings", "Ljava/util/Map;");
+    gBSPCFieldIds.disableAnimation = GetFieldIDOrDie(env, bspcClazz, "mDisableAnimation", "Z");
+    gBSPCFieldIds.disableAod = GetFieldIDOrDie(env, bspcClazz, "mDisableAod", "Z");
+    gBSPCFieldIds.disableLaunchBoost = GetFieldIDOrDie(env, bspcClazz, "mDisableLaunchBoost", "Z");
+    gBSPCFieldIds.disableOptionalSensors =
+            GetFieldIDOrDie(env, bspcClazz, "mDisableOptionalSensors", "Z");
+    gBSPCFieldIds.disableVibration = GetFieldIDOrDie(env, bspcClazz, "mDisableVibration", "Z");
+    gBSPCFieldIds.enableAdjustBrightness =
+            GetFieldIDOrDie(env, bspcClazz, "mEnableAdjustBrightness", "Z");
+    gBSPCFieldIds.enableDataSaver = GetFieldIDOrDie(env, bspcClazz, "mEnableDataSaver", "Z");
+    gBSPCFieldIds.enableFirewall = GetFieldIDOrDie(env, bspcClazz, "mEnableFirewall", "Z");
+    gBSPCFieldIds.enableNightMode = GetFieldIDOrDie(env, bspcClazz, "mEnableNightMode", "Z");
+    gBSPCFieldIds.enableQuickDoze = GetFieldIDOrDie(env, bspcClazz, "mEnableQuickDoze", "Z");
+    gBSPCFieldIds.forceAllAppsStandby =
+            GetFieldIDOrDie(env, bspcClazz, "mForceAllAppsStandby", "Z");
+    gBSPCFieldIds.forceBackgroundCheck =
+            GetFieldIDOrDie(env, bspcClazz, "mForceBackgroundCheck", "Z");
+    gBSPCFieldIds.locationMode = GetFieldIDOrDie(env, bspcClazz, "mLocationMode", "I");
+    gBSPCFieldIds.soundTriggerMode = GetFieldIDOrDie(env, bspcClazz, "mSoundTriggerMode", "I");
 
     jniRegisterNativeMethods(env, "android/os/PowerManagerTest", methodTable,
                 sizeof(methodTable) / sizeof(JNINativeMethod));
+
     return JNI_VERSION_1_6;
 }
 
diff --git a/core/tests/coretests/jni/NativeWorkSourceParcelTest.cpp b/core/tests/coretests/jni/NativeWorkSourceParcelTest.cpp
new file mode 100644
index 000000000000..db1f7bde4b4f
--- /dev/null
+++ b/core/tests/coretests/jni/NativeWorkSourceParcelTest.cpp
@@ -0,0 +1,156 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//#define LOG_NDEBUG 0
+#define LOG_TAG "NativeWorkSourceParcelTest"
+
+#include "jni.h"
+#include "ParcelHelper.h"
+
+#include <android_util_Binder.h>
+#include <nativehelper/JNIHelp.h>
+#include <nativehelper/ScopedPrimitiveArray.h>
+#include <utils/Log.h>
+
+#include <android/WorkSource.h>
+#include <android-base/stringprintf.h>
+#include <android-base/strings.h>
+
+using namespace android::os;
+using android::base::StringPrintf;
+
+namespace android {
+
+static jobject nativeObtainWorkSourceParcel(JNIEnv* env, jobject /* obj */, jintArray uidArray,
+            jobjectArray nameArray, jint parcelEndMarker) {
+    std::vector<int32_t> uids;
+    std::optional<std::vector<std::optional<String16>>> names = std::nullopt;
+
+    if (uidArray != nullptr) {
+        ScopedIntArrayRO workSourceUids(env, uidArray);
+        for (int i = 0; i < workSourceUids.size(); i++) {
+            uids.push_back(static_cast<int32_t>(workSourceUids[i]));
+        }
+    }
+
+    if (nameArray != nullptr) {
+        std::vector<std::optional<String16>> namesVec;
+        for (jint i = 0; i < env->GetArrayLength(nameArray); i++) {
+            jstring string = static_cast<jstring>(env->GetObjectArrayElement(nameArray, i));
+            const char *rawString = env->GetStringUTFChars(string, 0);
+            namesVec.push_back(std::make_optional<String16>(String16(rawString)));
+        }
+        names = std::make_optional(std::move(namesVec));
+    }
+
+    WorkSource ws = WorkSource(uids, names);
+    jobject wsParcel = nativeObtainParcel(env);
+    Parcel* parcel = nativeGetParcelData(env, wsParcel);
+
+    // write WorkSource and if no error write end marker
+    status_t err = ws.writeToParcel(parcel) ?: parcel->writeInt32(parcelEndMarker);
+
+    if (err != OK) {
+        jniThrowException(env, "java/lang/IllegalArgumentException",
+                            StringPrintf("WorkSource writeToParcel failed %d", err).c_str());
+    }
+    parcel->setDataPosition(0);
+    return wsParcel;
+}
+
+static void nativeUnparcelAndVerifyWorkSource(JNIEnv* env, jobject /* obj */, jobject wsParcel,
+        jintArray uidArray, jobjectArray nameArray, jint parcelEndMarker) {
+    WorkSource ws = {};
+    Parcel* parcel = nativeGetParcelData(env, wsParcel);
+    int32_t endMarker;
+
+    // read WorkSource and if no error read end marker
+    status_t err = ws.readFromParcel(parcel) ?: parcel->readInt32(&endMarker);
+    int32_t dataAvailable = parcel->dataAvail();
+
+    if (err != OK) {
+        ALOGE("WorkSource readFromParcel failed %d", err);
+    }
+
+    // Now we have a native WorkSource object, verify it.
+    if (dataAvailable > 0) { // not all data read from the parcel
+        jniThrowException(env, "java/lang/IllegalArgumentException",
+                StringPrintf("WorkSource contains more data than native read (%d)",
+                dataAvailable).c_str());
+    } else if (endMarker != parcelEndMarker) { // more date than available read from parcel
+        jniThrowException(env, "java/lang/IllegalArgumentException",
+                StringPrintf("WorkSource contains less data than native read").c_str());
+    }
+
+    if (uidArray != nullptr) {
+        ScopedIntArrayRO workSourceUids(env, uidArray);
+        for (int i = 0; i < workSourceUids.size(); i++) {
+            if (ws.getUids().at(i) != static_cast<int32_t>(workSourceUids[i])) {
+                jniThrowException(env, "java/lang/IllegalArgumentException",
+                            StringPrintf("WorkSource uid not equal %d %d",
+                            ws.getUids().at(i), static_cast<int32_t>(workSourceUids[i])).c_str());
+            }
+        }
+    } else {
+        if (ws.getUids().size() != 0) {
+            jniThrowException(env, "java/lang/IllegalArgumentException",
+                    StringPrintf("WorkSource parcel size not 0").c_str());
+        }
+    }
+
+    if (nameArray != nullptr) {
+        std::vector<std::optional<String16>> namesVec;
+        for (jint i = 0; i < env->GetArrayLength(nameArray); i++) {
+            jstring string = (jstring) (env->GetObjectArrayElement(nameArray, i));
+            const char *rawString = env->GetStringUTFChars(string, 0);
+            if (String16(rawString) != ws.getNames()->at(i)) {
+                jniThrowException(env, "java/lang/IllegalArgumentException",
+                            StringPrintf("WorkSource uid not equal %s", rawString).c_str());
+            }
+        }
+    } else {
+        if (ws.getNames() != std::nullopt) {
+            jniThrowException(env, "java/lang/IllegalArgumentException",
+                    StringPrintf("WorkSource parcel name not empty").c_str());
+        }
+    }
+}
+
+extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */)
+{
+    JNIEnv* env;
+
+    const JNINativeMethod workSourceMethodTable[] = {
+        /* name, signature, funcPtr */
+        { "nativeObtainWorkSourceParcel", "([I[Ljava/lang/String;I)Landroid/os/Parcel;",
+                (void*) nativeObtainWorkSourceParcel },
+        { "nativeUnparcelAndVerifyWorkSource", "(Landroid/os/Parcel;[I[Ljava/lang/String;I)V",
+                (void*) nativeUnparcelAndVerifyWorkSource },
+    };
+
+    if (vm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
+        return JNI_ERR;
+    }
+
+    loadParcelClass(env);
+
+    jniRegisterNativeMethods(env, "android/os/WorkSourceParcelTest", workSourceMethodTable,
+                sizeof(workSourceMethodTable) / sizeof(JNINativeMethod));
+
+    return JNI_VERSION_1_6;
+}
+
+} /* namespace android */
diff --git a/core/tests/coretests/jni/ParcelHelper.h b/core/tests/coretests/jni/ParcelHelper.h
new file mode 100644
index 000000000000..1485a2a3be17
--- /dev/null
+++ b/core/tests/coretests/jni/ParcelHelper.h
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <nativehelper/JNIHelp.h>
+#include <binder/Parcel.h>
+
+namespace android {
+    static jclass gParcelClazz;
+    static jfieldID gParcelDataFieldID;
+    static jmethodID gParcelObtainMethodID;
+
+    static inline jclass FindClassOrDie(JNIEnv* env, const char* class_name) {
+        jclass clazz = env->FindClass(class_name);
+        LOG_ALWAYS_FATAL_IF(clazz == NULL, "Unable to find class %s", class_name);
+        return clazz;
+    }
+
+    static inline jfieldID GetFieldIDOrDie(JNIEnv* env, jclass clazz, const char* field_name,
+                                           const char* field_signature) {
+        jfieldID res = env->GetFieldID(clazz, field_name, field_signature);
+        LOG_ALWAYS_FATAL_IF(res == NULL, "Unable to find field %s with signature %s", field_name,
+                            field_signature);
+        return res;
+    }
+
+    static inline jmethodID GetStaticMethodIDOrDie(JNIEnv* env, jclass clazz,
+                                                   const char* method_name,
+                                                   const char* method_signature) {
+        jmethodID res = env->GetStaticMethodID(clazz, method_name, method_signature);
+        LOG_ALWAYS_FATAL_IF(res == NULL, "Unable to find method %s with signature %s", method_name,
+                            method_signature);
+        return res;
+    }
+
+    static jobject nativeObtainParcel(JNIEnv* env) {
+        jobject parcel = env->CallStaticObjectMethod(gParcelClazz, gParcelObtainMethodID);
+        if (parcel == nullptr) {
+            jniThrowException(env, "java/lang/IllegalArgumentException", "Obtain parcel failed.");
+        }
+        return parcel;
+    }
+
+    static Parcel* nativeGetParcelData(JNIEnv* env, jobject obj) {
+        Parcel* parcel = reinterpret_cast<Parcel*>(env->GetLongField(obj, gParcelDataFieldID));
+        if (parcel && parcel->objectsCount() != 0) {
+            jniThrowException(env, "java/lang/IllegalArgumentException", "Invalid parcel object.");
+        }
+        parcel->setDataPosition(0);
+        return parcel;
+    }
+
+    static void loadParcelClass(JNIEnv* env) {
+        gParcelClazz = FindClassOrDie(env, "android/os/Parcel");
+        gParcelDataFieldID = GetFieldIDOrDie(env, gParcelClazz, "mNativePtr", "J");
+        gParcelObtainMethodID =
+                GetStaticMethodIDOrDie(env, gParcelClazz, "obtain", "()Landroid/os/Parcel;");
+    }
+
+}
\ No newline at end of file
diff --git a/core/tests/coretests/src/android/os/PowerManagerTest.java b/core/tests/coretests/src/android/os/PowerManagerTest.java
index 0dfc37131f7f..9f85d6f02ec3 100644
--- a/core/tests/coretests/src/android/os/PowerManagerTest.java
+++ b/core/tests/coretests/src/android/os/PowerManagerTest.java
@@ -49,9 +49,6 @@ public class PowerManagerTest extends AndroidTestCase {
     @Mock
     private PowerManager.OnThermalStatusChangedListener mListener2;
     private static final long CALLBACK_TIMEOUT_MILLI_SEC = 5000;
-    private native Parcel nativeObtainWorkSourceParcel(int[] uids, String[] names);
-    private native void nativeUnparcelAndVerifyWorkSource(Parcel parcel, int[] uids,
-            String[] names);
     private native Parcel nativeObtainPowerSaveStateParcel(boolean batterySaverEnabled,
             boolean globalBatterySaverEnabled, int locationMode, int soundTriggerMode,
             float brightnessFactor);
@@ -300,54 +297,6 @@ public class PowerManagerTest extends AndroidTestCase {
     }
 
     /**
-     * Helper function to obtain a WorkSource object as parcel from native, with
-     * specified uids and names and verify the WorkSource object created from the parcel.
-     */
-    private void unparcelWorkSourceFromNativeAndVerify(int[] uids, String[] names) {
-        // Obtain WorkSource as parcel from native, with uids and names.
-        Parcel wsParcel = nativeObtainWorkSourceParcel(uids, names);
-        WorkSource ws = WorkSource.CREATOR.createFromParcel(wsParcel);
-        if (uids == null)  {
-            assertEquals(ws.size(), 0);
-        } else {
-            assertEquals(uids.length, ws.size());
-            for (int i = 0; i < ws.size(); i++) {
-                assertEquals(ws.getUid(i), uids[i]);
-            }
-        }
-        if (names != null)  {
-            for (int i = 0; i < names.length; i++) {
-                assertEquals(ws.getName(i), names[i]);
-            }
-        }
-    }
-
-    /**
-     * Helper function to send a WorkSource as parcel from java to native.
-     * Native will verify the WorkSource in native is expected.
-     */
-    private void parcelWorkSourceToNativeAndVerify(int[] uids, String[] names) {
-        WorkSource ws = new WorkSource();
-        if (uids != null) {
-            if (names == null) {
-                for (int i = 0; i < uids.length; i++) {
-                    ws.add(uids[i]);
-                }
-            } else {
-                assertEquals(uids.length, names.length);
-                for (int i = 0; i < uids.length; i++) {
-                    ws.add(uids[i], names[i]);
-                }
-            }
-        }
-        Parcel wsParcel = Parcel.obtain();
-        ws.writeToParcel(wsParcel, 0 /* flags */);
-        wsParcel.setDataPosition(0);
-        //Set the WorkSource as parcel to native and verify.
-        nativeUnparcelAndVerifyWorkSource(wsParcel, uids, names);
-    }
-
-    /**
      * Helper function to obtain a PowerSaveState as parcel from native, with
      * specified parameters, and verify the PowerSaveState object created from the parcel.
      */
@@ -422,43 +371,6 @@ public class PowerManagerTest extends AndroidTestCase {
     }
 
     /**
-     * Confirm that we can pass WorkSource from native to Java.
-     *
-     * @throws Exception
-     */
-    @Test
-    public void testWorkSourceNativeToJava() {
-        final int[] uids1 = {1000};
-        final int[] uids2 = {1000, 2000};
-        final String[] names1 = {"testWorkSource1"};
-        final String[] names2 = {"testWorkSource1", "testWorkSource2"};
-        unparcelWorkSourceFromNativeAndVerify(null /* uids */, null /* names */);
-        unparcelWorkSourceFromNativeAndVerify(uids1, null /* names */);
-        unparcelWorkSourceFromNativeAndVerify(uids2, null /* names */);
-        unparcelWorkSourceFromNativeAndVerify(null /* uids */, names1);
-        unparcelWorkSourceFromNativeAndVerify(uids1, names1);
-        unparcelWorkSourceFromNativeAndVerify(uids2, names2);
-    }
-
-    /**
-     * Confirm that we can pass WorkSource from Java to native.
-     *
-     * @throws Exception
-     */
-    @Test
-    public void testWorkSourceJavaToNative() {
-        final int[] uids1 = {1000};
-        final int[] uids2 = {1000, 2000};
-        final String[] names1 = {"testGetWorkSource1"};
-        final String[] names2 = {"testGetWorkSource1", "testGetWorkSource2"};
-        parcelWorkSourceToNativeAndVerify(null /* uids */, null /* names */);
-        parcelWorkSourceToNativeAndVerify(uids1, null /* names */);
-        parcelWorkSourceToNativeAndVerify(uids2, null /* names */);
-        parcelWorkSourceToNativeAndVerify(uids1, names1);
-        parcelWorkSourceToNativeAndVerify(uids2, names2);
-    }
-
-    /**
      * Confirm that we can pass PowerSaveState from native to Java.
      *
      * @throws Exception
diff --git a/core/tests/coretests/src/android/os/WorkSourceParcelTest.java b/core/tests/coretests/src/android/os/WorkSourceParcelTest.java
new file mode 100644
index 000000000000..6dcc3a6dfb3c
--- /dev/null
+++ b/core/tests/coretests/src/android/os/WorkSourceParcelTest.java
@@ -0,0 +1,139 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.os;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class WorkSourceParcelTest {
+    /**
+     * END_OF_PARCEL_MARKER is added at the end of Parcel on native or java side on write and
+     * then read on java or native side on read. This way we can ensure that no extra data
+     * is read from the parcel.
+     */
+    private static final int END_OF_PARCEL_MARKER = 99;
+
+    private native Parcel nativeObtainWorkSourceParcel(int[] uids, String[] names,
+            int parcelEndMarker);
+
+    private native void nativeUnparcelAndVerifyWorkSource(Parcel parcel, int[] uids,
+            String[] names, int parcelEndMarker);
+
+    static {
+        System.loadLibrary("worksourceparceltest_jni");
+    }
+    /**
+     * Confirm that we can pass WorkSource from native to Java.
+     */
+    @Test
+    public void testWorkSourceNativeToJava() {
+        final int[] uids1 = {1000};
+        final int[] uids2 = {1000, 2000};
+        final String[] names1 = {"testWorkSource1"};
+        final String[] names2 = {"testWorkSource1", "testWorkSource2"};
+        unparcelWorkSourceFromNativeAndVerify(/* uids= */ null , /* names= */ null);
+        unparcelWorkSourceFromNativeAndVerify(uids1, /* names= */ null);
+        unparcelWorkSourceFromNativeAndVerify(uids2, /* names= */ null);
+        //TODO(b/273752423)
+        unparcelWorkSourceFromNativeAndVerify(/* uids= */ null , names1);
+        unparcelWorkSourceFromNativeAndVerify(uids1, names1);
+        unparcelWorkSourceFromNativeAndVerify(uids2, names2);
+    }
+
+    /**
+     * Confirm that we can pass WorkSource from Java to native.
+     */
+    @Test
+    public void testWorkSourceJavaToNative() {
+        final int[] uids1 = {1000};
+        final int[] uids2 = {1000, 2000};
+        final String[] names1 = {"testGetWorkSource1"};
+        final String[] names2 = {"testGetWorkSource1", "testGetWorkSource2"};
+        //TODO(b/273752423)
+        //parcelWorkSourceToNativeAndVerify(/* uids= */ null , /* names= */ null );
+        parcelWorkSourceToNativeAndVerify(uids1, /* names= */ null);
+        parcelWorkSourceToNativeAndVerify(uids2, /* names= */ null);
+        parcelWorkSourceToNativeAndVerify(uids1, names1);
+        parcelWorkSourceToNativeAndVerify(uids2, names2);
+    }
+
+    /**
+     * Helper function to obtain a WorkSource object as parcel from native, with
+     * specified uids and names and verify the WorkSource object created from the parcel.
+     */
+    private void unparcelWorkSourceFromNativeAndVerify(int[] uids, String[] names) {
+        // Obtain WorkSource as parcel from native, with uids and names.
+        // END_OF_PARCEL_MARKER is written at the end of parcel
+        Parcel wsParcel = nativeObtainWorkSourceParcel(uids, names, END_OF_PARCEL_MARKER);
+        // read WorkSource created on native side
+        WorkSource ws = WorkSource.CREATOR.createFromParcel(wsParcel);
+        // read end marker written on native side
+        int endMarker = wsParcel.readInt();
+
+        assertEquals(0, wsParcel.dataAvail()); // we have read everything
+        assertEquals(END_OF_PARCEL_MARKER, endMarker); // endMarkers match
+
+        if (uids == null) {
+            assertEquals(ws.size(), 0);
+        } else {
+            assertEquals(uids.length, ws.size());
+            for (int i = 0; i < ws.size(); i++) {
+                assertEquals(ws.getUid(i), uids[i]);
+            }
+        }
+        if (names != null) {
+            for (int i = 0; i < names.length; i++) {
+                assertEquals(ws.getPackageName(i), names[i]);
+            }
+        }
+    }
+
+    /**
+     * Helper function to send a WorkSource as parcel from java to native.
+     * Native will verify the WorkSource in native is expected.
+     */
+    private void parcelWorkSourceToNativeAndVerify(int[] uids, String[] names) {
+        WorkSource ws = new WorkSource();
+        if (uids != null) {
+            if (names == null) {
+                for (int uid : uids) {
+                    ws.add(uid);
+                }
+            } else {
+                assertEquals(uids.length, names.length);
+                for (int i = 0; i < uids.length; i++) {
+                    ws.add(uids[i], names[i]);
+                }
+            }
+        }
+        Parcel wsParcel = Parcel.obtain();
+        // write WorkSource on java side
+        ws.writeToParcel(wsParcel, 0 /* flags */);
+        // write end marker on java side
+        wsParcel.writeInt(END_OF_PARCEL_MARKER);
+        wsParcel.setDataPosition(0);
+        //Verify parcel and end marker on native side
+        nativeUnparcelAndVerifyWorkSource(wsParcel, uids, names, END_OF_PARCEL_MARKER);
+    }
+}