1 files changed, 25 insertions, 4 deletions
diff --git a/neuralnetworks/aidl/vts/functional/Utils.cpp b/neuralnetworks/aidl/vts/functional/Utils.cpp
index 3c7f5f797d..9af362ea8a 100644
--- a/neuralnetworks/aidl/vts/functional/Utils.cpp
+++ b/neuralnetworks/aidl/vts/functional/Utils.cpp
@@ -23,6 +23,7 @@
 #include <android/binder_status.h>
 #include <android/hardware_buffer.h>
 
+#include <sys/mman.h>
 #include <iostream>
 #include <limits>
 #include <numeric>
@@ -98,19 +99,39 @@ uint32_t sizeOfData(const Operand& operand) {
                            std::multiplies<>{});
 }
 
-std::unique_ptr<TestAshmem> TestAshmem::create(uint32_t size) {
-    auto ashmem = std::make_unique<TestAshmem>(size);
+std::unique_ptr<TestAshmem> TestAshmem::create(uint32_t size, bool aidlReadonly) {
+    auto ashmem = std::make_unique<TestAshmem>(size, aidlReadonly);
     return ashmem->mIsValid ? std::move(ashmem) : nullptr;
 }
 
-void TestAshmem::initialize(uint32_t size) {
+// This function will create a readonly shared memory with PROT_READ only.
+// The input shared memory must be either Ashmem or mapped-FD.
+static nn::SharedMemory convertSharedMemoryToReadonly(const nn::SharedMemory& sharedMemory) {
+    if (std::holds_alternative<nn::Memory::Ashmem>(sharedMemory->handle)) {
+        const auto& memory = std::get<nn::Memory::Ashmem>(sharedMemory->handle);
+        return nn::createSharedMemoryFromFd(memory.size, PROT_READ, memory.fd.get(), /*offset=*/0)
+                .value();
+    } else if (std::holds_alternative<nn::Memory::Fd>(sharedMemory->handle)) {
+        const auto& memory = std::get<nn::Memory::Fd>(sharedMemory->handle);
+        return nn::createSharedMemoryFromFd(memory.size, PROT_READ, memory.fd.get(), memory.offset)
+                .value();
+    }
+    CHECK(false) << "Unexpected shared memory type";
+    return sharedMemory;
+}
+
+void TestAshmem::initialize(uint32_t size, bool aidlReadonly) {
     mIsValid = false;
     ASSERT_GT(size, 0);
     const auto sharedMemory = nn::createSharedMemory(size).value();
     mMappedMemory = nn::map(sharedMemory).value();
     mPtr = static_cast<uint8_t*>(std::get<void*>(mMappedMemory.pointer));
     CHECK_NE(mPtr, nullptr);
-    mAidlMemory = utils::convert(sharedMemory).value();
+    if (aidlReadonly) {
+        mAidlMemory = utils::convert(convertSharedMemoryToReadonly(sharedMemory)).value();
+    } else {
+        mAidlMemory = utils::convert(sharedMemory).value();
+    }
     mIsValid = true;
 }