Add AHB_allocate2 + get/setDataSpace

Also fix that some ADataSpace enum values were not prefixed
with ADATASPACE_*

Test: atest AHardwareBufferTest
Bug: 315475131
Change-Id: Ie92f416a209501fa8fdc9b3593ea6d3bb68eaa48
diff --git a/libs/ui/Gralloc2.cpp b/libs/ui/Gralloc2.cpp
index e9b5dec..a5aca99 100644
--- a/libs/ui/Gralloc2.cpp
+++ b/libs/ui/Gralloc2.cpp
@@ -384,8 +384,8 @@
 
 status_t Gralloc2Allocator::allocate(std::string /*requestorName*/, uint32_t width, uint32_t height,
                                      PixelFormat format, uint32_t layerCount, uint64_t usage,
-                                     uint32_t bufferCount, uint32_t* outStride,
-                                     buffer_handle_t* outBufferHandles, bool importBuffers) const {
+                                     uint32_t* outStride, buffer_handle_t* outBufferHandles,
+                                     bool importBuffers) const {
     IMapper::BufferDescriptorInfo descriptorInfo = {};
     descriptorInfo.width = width;
     descriptorInfo.height = height;
@@ -400,6 +400,8 @@
         return error;
     }
 
+    constexpr auto bufferCount = 1;
+
     auto ret = mAllocator->allocate(descriptor, bufferCount,
                                     [&](const auto& tmpError, const auto& tmpStride,
                                         const auto& tmpBuffers) {
diff --git a/libs/ui/Gralloc3.cpp b/libs/ui/Gralloc3.cpp
index 474d381..152b35a 100644
--- a/libs/ui/Gralloc3.cpp
+++ b/libs/ui/Gralloc3.cpp
@@ -371,7 +371,7 @@
 
 status_t Gralloc3Allocator::allocate(std::string /*requestorName*/, uint32_t width, uint32_t height,
                                      android::PixelFormat format, uint32_t layerCount,
-                                     uint64_t usage, uint32_t bufferCount, uint32_t* outStride,
+                                     uint64_t usage, uint32_t* outStride,
                                      buffer_handle_t* outBufferHandles, bool importBuffers) const {
     IMapper::BufferDescriptorInfo descriptorInfo;
     sBufferDescriptorInfo(width, height, format, layerCount, usage, &descriptorInfo);
@@ -383,6 +383,8 @@
         return error;
     }
 
+    constexpr auto bufferCount = 1;
+
     auto ret = mAllocator->allocate(descriptor, bufferCount,
                                     [&](const auto& tmpError, const auto& tmpStride,
                                         const auto& tmpBuffers) {
diff --git a/libs/ui/Gralloc4.cpp b/libs/ui/Gralloc4.cpp
index 03ff58a..d6970e0 100644
--- a/libs/ui/Gralloc4.cpp
+++ b/libs/ui/Gralloc4.cpp
@@ -1069,7 +1069,7 @@
 
 status_t Gralloc4Allocator::allocate(std::string requestorName, uint32_t width, uint32_t height,
                                      android::PixelFormat format, uint32_t layerCount,
-                                     uint64_t usage, uint32_t bufferCount, uint32_t* outStride,
+                                     uint64_t usage, uint32_t* outStride,
                                      buffer_handle_t* outBufferHandles, bool importBuffers) const {
     IMapper::BufferDescriptorInfo descriptorInfo;
     if (auto error = sBufferDescriptorInfo(requestorName, width, height, format, layerCount, usage,
@@ -1084,6 +1084,8 @@
         return error;
     }
 
+    constexpr auto bufferCount = 1;
+
     if (mAidlAllocator) {
         AllocationResult result;
 #pragma clang diagnostic push
diff --git a/libs/ui/Gralloc5.cpp b/libs/ui/Gralloc5.cpp
index 2ec6d18..b07e155 100644
--- a/libs/ui/Gralloc5.cpp
+++ b/libs/ui/Gralloc5.cpp
@@ -19,6 +19,7 @@
 
 #include <ui/Gralloc5.h>
 
+#include <aidl/android/hardware/graphics/allocator/AllocationError.h>
 #include <aidlcommonsupport/NativeHandle.h>
 #include <android/binder_manager.h>
 #include <android/hardware/graphics/mapper/utils/IMapperMetadataTypes.h>
@@ -215,55 +216,75 @@
 
 status_t Gralloc5Allocator::allocate(std::string requestorName, uint32_t width, uint32_t height,
                                      android::PixelFormat format, uint32_t layerCount,
-                                     uint64_t usage, uint32_t bufferCount, uint32_t *outStride,
-                                     buffer_handle_t *outBufferHandles, bool importBuffers) const {
-    auto descriptorInfo = makeDescriptor(requestorName, width, height, format, layerCount, usage);
+                                     uint64_t usage, uint32_t* outStride,
+                                     buffer_handle_t* outBufferHandles, bool importBuffers) const {
+    auto result = allocate(GraphicBufferAllocator::AllocationRequest{
+            .importBuffer = importBuffers,
+            .width = width,
+            .height = height,
+            .format = format,
+            .layerCount = layerCount,
+            .usage = usage,
+            .requestorName = requestorName,
+    });
+
+    *outStride = result.stride;
+    outBufferHandles[0] = result.handle;
+    return result.status;
+}
+
+GraphicBufferAllocator::AllocationResult Gralloc5Allocator::allocate(
+        const GraphicBufferAllocator::AllocationRequest& request) const {
+    auto descriptorInfo = makeDescriptor(request.requestorName, request.width, request.height,
+                                         request.format, request.layerCount, request.usage);
     if (!descriptorInfo) {
-        return BAD_VALUE;
+        return GraphicBufferAllocator::AllocationResult{BAD_VALUE};
+    }
+
+    descriptorInfo->additionalOptions.reserve(request.extras.size());
+    for (const auto& option : request.extras) {
+        ExtendableType type;
+        type.name = option.name;
+        type.value = option.value;
+        descriptorInfo->additionalOptions.push_back(std::move(type));
     }
 
     AllocationResult result;
-    auto status = mAllocator->allocate2(*descriptorInfo, bufferCount, &result);
+    auto status = mAllocator->allocate2(*descriptorInfo, 1, &result);
     if (!status.isOk()) {
         auto error = status.getExceptionCode();
         if (error == EX_SERVICE_SPECIFIC) {
-            error = status.getServiceSpecificError();
+            switch (static_cast<AllocationError>(status.getServiceSpecificError())) {
+                case AllocationError::BAD_DESCRIPTOR:
+                    error = BAD_VALUE;
+                    break;
+                case AllocationError::NO_RESOURCES:
+                    error = NO_MEMORY;
+                    break;
+                default:
+                    error = UNKNOWN_ERROR;
+                    break;
+            }
         }
-        if (error == OK) {
-            error = UNKNOWN_ERROR;
-        }
-        return error;
+        return GraphicBufferAllocator::AllocationResult{error};
     }
 
-    if (importBuffers) {
-        for (uint32_t i = 0; i < bufferCount; i++) {
-            auto handle = makeFromAidl(result.buffers[i]);
-            auto error = mMapper.importBuffer(handle, &outBufferHandles[i]);
-            native_handle_delete(handle);
-            if (error != NO_ERROR) {
-                for (uint32_t j = 0; j < i; j++) {
-                    mMapper.freeBuffer(outBufferHandles[j]);
-                    outBufferHandles[j] = nullptr;
-                }
-                return error;
-            }
+    GraphicBufferAllocator::AllocationResult ret{OK};
+    if (request.importBuffer) {
+        auto handle = makeFromAidl(result.buffers[0]);
+        auto error = mMapper.importBuffer(handle, &ret.handle);
+        native_handle_delete(handle);
+        if (error != NO_ERROR) {
+            return GraphicBufferAllocator::AllocationResult{error};
         }
     } else {
-        for (uint32_t i = 0; i < bufferCount; i++) {
-            outBufferHandles[i] = dupFromAidl(result.buffers[i]);
-            if (!outBufferHandles[i]) {
-                for (uint32_t j = 0; j < i; j++) {
-                    auto buffer = const_cast<native_handle_t *>(outBufferHandles[j]);
-                    native_handle_close(buffer);
-                    native_handle_delete(buffer);
-                    outBufferHandles[j] = nullptr;
-                }
-                return NO_MEMORY;
-            }
+        ret.handle = dupFromAidl(result.buffers[0]);
+        if (!ret.handle) {
+            return GraphicBufferAllocator::AllocationResult{NO_MEMORY};
         }
     }
 
-    *outStride = result.stride;
+    ret.stride = result.stride;
 
     // Release all the resources held by AllocationResult (specifically any remaining FDs)
     result = {};
@@ -272,7 +293,7 @@
     // is marked apex_available (b/214400477) and libbinder isn't (which of course is correct)
     // IPCThreadState::self()->flushCommands();
 
-    return OK;
+    return ret;
 }
 
 void Gralloc5Mapper::preload() {
diff --git a/libs/ui/GraphicBuffer.cpp b/libs/ui/GraphicBuffer.cpp
index 429760f..c007fdb 100644
--- a/libs/ui/GraphicBuffer.cpp
+++ b/libs/ui/GraphicBuffer.cpp
@@ -106,6 +106,26 @@
                                 inUsage, inStride);
 }
 
+GraphicBuffer::GraphicBuffer(const GraphicBufferAllocator::AllocationRequest& request)
+      : GraphicBuffer() {
+    GraphicBufferAllocator& allocator = GraphicBufferAllocator::get();
+    auto result = allocator.allocate(request);
+    mInitCheck = result.status;
+    if (result.status == NO_ERROR) {
+        handle = result.handle;
+        stride = result.stride;
+
+        mBufferMapper.getTransportSize(handle, &mTransportNumFds, &mTransportNumInts);
+
+        width = static_cast<int>(request.width);
+        height = static_cast<int>(request.height);
+        format = request.format;
+        layerCount = request.layerCount;
+        usage = request.usage;
+        usage_deprecated = int(usage);
+    }
+}
+
 GraphicBuffer::~GraphicBuffer()
 {
     ATRACE_CALL();
@@ -143,6 +163,10 @@
             const_cast<GraphicBuffer*>(this));
 }
 
+status_t GraphicBuffer::getDataspace(ui::Dataspace* outDataspace) const {
+    return mBufferMapper.getDataspace(handle, outDataspace);
+}
+
 status_t GraphicBuffer::reallocate(uint32_t inWidth, uint32_t inHeight,
         PixelFormat inFormat, uint32_t inLayerCount, uint64_t inUsage)
 {
diff --git a/libs/ui/GraphicBufferAllocator.cpp b/libs/ui/GraphicBufferAllocator.cpp
index eb0bd4e..98082fb 100644
--- a/libs/ui/GraphicBufferAllocator.cpp
+++ b/libs/ui/GraphicBufferAllocator.cpp
@@ -113,6 +113,79 @@
     ALOGD("%s", s.c_str());
 }
 
+auto GraphicBufferAllocator::allocate(const AllocationRequest& request) -> AllocationResult {
+    ATRACE_CALL();
+    if (!request.width || !request.height) {
+        return AllocationResult(BAD_VALUE);
+    }
+
+    const auto width = request.width;
+    const auto height = request.height;
+
+    const uint32_t bpp = bytesPerPixel(request.format);
+    if (std::numeric_limits<size_t>::max() / width / height < static_cast<size_t>(bpp)) {
+        ALOGE("Failed to allocate (%u x %u) layerCount %u format %d "
+              "usage %" PRIx64 ": Requesting too large a buffer size",
+              request.width, request.height, request.layerCount, request.format, request.usage);
+        return AllocationResult(BAD_VALUE);
+    }
+
+    if (request.layerCount < 1) {
+        return AllocationResult(BAD_VALUE);
+    }
+
+    auto result = mAllocator->allocate(request);
+    if (result.status == UNKNOWN_TRANSACTION) {
+        if (!request.extras.empty()) {
+            ALOGE("Failed to allocate with additional options, allocator version mis-match? "
+                  "gralloc version = %d",
+                  (int)mMapper.getMapperVersion());
+            return result;
+        }
+        // If there's no additional options, fall back to previous allocate version
+        result.status = mAllocator->allocate(request.requestorName, request.width, request.height,
+                                             request.format, request.layerCount, request.usage,
+                                             &result.stride, &result.handle, request.importBuffer);
+    }
+
+    if (result.status != NO_ERROR) {
+        ALOGE("Failed to allocate (%u x %u) layerCount %u format %d "
+              "usage %" PRIx64 ": %d",
+              request.width, request.height, request.layerCount, request.format, request.usage,
+              result.status);
+        return result;
+    }
+
+    if (!request.importBuffer) {
+        return result;
+    }
+    size_t bufSize;
+
+    // if stride has no meaning or is too large,
+    // approximate size with the input width instead
+    if ((result.stride) != 0 &&
+        std::numeric_limits<size_t>::max() / height / (result.stride) < static_cast<size_t>(bpp)) {
+        bufSize = static_cast<size_t>(width) * height * bpp;
+    } else {
+        bufSize = static_cast<size_t>((result.stride)) * height * bpp;
+    }
+
+    Mutex::Autolock _l(sLock);
+    KeyedVector<buffer_handle_t, alloc_rec_t>& list(sAllocList);
+    alloc_rec_t rec;
+    rec.width = width;
+    rec.height = height;
+    rec.stride = result.stride;
+    rec.format = request.format;
+    rec.layerCount = request.layerCount;
+    rec.usage = request.usage;
+    rec.size = bufSize;
+    rec.requestorName = request.requestorName;
+    list.add(result.handle, rec);
+
+    return result;
+}
+
 status_t GraphicBufferAllocator::allocateHelper(uint32_t width, uint32_t height, PixelFormat format,
                                                 uint32_t layerCount, uint64_t usage,
                                                 buffer_handle_t* handle, uint32_t* stride,
@@ -141,7 +214,7 @@
     usage &= ~static_cast<uint64_t>((1 << 10) | (1 << 13));
 
     status_t error = mAllocator->allocate(requestorName, width, height, format, layerCount, usage,
-                                          1, stride, handle, importBuffer);
+                                          stride, handle, importBuffer);
     if (error != NO_ERROR) {
         ALOGE("Failed to allocate (%u x %u) layerCount %u format %d "
               "usage %" PRIx64 ": %d",
diff --git a/libs/ui/include/ui/Gralloc.h b/libs/ui/include/ui/Gralloc.h
index 496ba57..e6015e0 100644
--- a/libs/ui/include/ui/Gralloc.h
+++ b/libs/ui/include/ui/Gralloc.h
@@ -23,6 +23,7 @@
 #include <ui/PixelFormat.h>
 #include <ui/Rect.h>
 #include <utils/StrongPointer.h>
+#include "GraphicBufferAllocator.h"
 
 #include <string>
 
@@ -218,9 +219,13 @@
      */
     virtual status_t allocate(std::string requestorName, uint32_t width, uint32_t height,
                               PixelFormat format, uint32_t layerCount, uint64_t usage,
-                              uint32_t bufferCount, uint32_t* outStride,
-                              buffer_handle_t* outBufferHandles,
+                              uint32_t* outStride, buffer_handle_t* outBufferHandles,
                               bool importBuffers = true) const = 0;
+
+    virtual GraphicBufferAllocator::AllocationResult allocate(
+            const GraphicBufferAllocator::AllocationRequest&) const {
+        return GraphicBufferAllocator::AllocationResult(UNKNOWN_TRANSACTION);
+    }
 };
 
 } // namespace android
diff --git a/libs/ui/include/ui/Gralloc2.h b/libs/ui/include/ui/Gralloc2.h
index a7b6f492..e50bb3a 100644
--- a/libs/ui/include/ui/Gralloc2.h
+++ b/libs/ui/include/ui/Gralloc2.h
@@ -81,9 +81,8 @@
     std::string dumpDebugInfo(bool less = true) const override;
 
     status_t allocate(std::string requestorName, uint32_t width, uint32_t height,
-                      PixelFormat format, uint32_t layerCount, uint64_t usage, uint32_t bufferCount,
-                      uint32_t* outStride, buffer_handle_t* outBufferHandles,
-                      bool importBuffers = true) const override;
+                      PixelFormat format, uint32_t layerCount, uint64_t usage, uint32_t* outStride,
+                      buffer_handle_t* outBufferHandles, bool importBuffers = true) const override;
 
 private:
     const Gralloc2Mapper& mMapper;
diff --git a/libs/ui/include/ui/Gralloc3.h b/libs/ui/include/ui/Gralloc3.h
index 7367549..035684a 100644
--- a/libs/ui/include/ui/Gralloc3.h
+++ b/libs/ui/include/ui/Gralloc3.h
@@ -82,9 +82,8 @@
     std::string dumpDebugInfo(bool less = true) const override;
 
     status_t allocate(std::string requestorName, uint32_t width, uint32_t height,
-                      PixelFormat format, uint32_t layerCount, uint64_t usage, uint32_t bufferCount,
-                      uint32_t* outStride, buffer_handle_t* outBufferHandles,
-                      bool importBuffers = true) const override;
+                      PixelFormat format, uint32_t layerCount, uint64_t usage, uint32_t* outStride,
+                      buffer_handle_t* outBufferHandles, bool importBuffers = true) const override;
 
 private:
     const Gralloc3Mapper& mMapper;
diff --git a/libs/ui/include/ui/Gralloc4.h b/libs/ui/include/ui/Gralloc4.h
index df43be8..0f469c0 100644
--- a/libs/ui/include/ui/Gralloc4.h
+++ b/libs/ui/include/ui/Gralloc4.h
@@ -174,9 +174,8 @@
     std::string dumpDebugInfo(bool less = true) const override;
 
     status_t allocate(std::string requestorName, uint32_t width, uint32_t height,
-                      PixelFormat format, uint32_t layerCount, uint64_t usage, uint32_t bufferCount,
-                      uint32_t* outStride, buffer_handle_t* outBufferHandles,
-                      bool importBuffers = true) const override;
+                      PixelFormat format, uint32_t layerCount, uint64_t usage, uint32_t* outStride,
+                      buffer_handle_t* outBufferHandles, bool importBuffers = true) const override;
 
 private:
     const Gralloc4Mapper& mMapper;
diff --git a/libs/ui/include/ui/Gralloc5.h b/libs/ui/include/ui/Gralloc5.h
index 44b97d1..f9e8f5e 100644
--- a/libs/ui/include/ui/Gralloc5.h
+++ b/libs/ui/include/ui/Gralloc5.h
@@ -172,10 +172,12 @@
 
     [[nodiscard]] status_t allocate(std::string requestorName, uint32_t width, uint32_t height,
                                     PixelFormat format, uint32_t layerCount, uint64_t usage,
-                                    uint32_t bufferCount, uint32_t *outStride,
-                                    buffer_handle_t *outBufferHandles,
+                                    uint32_t* outStride, buffer_handle_t* outBufferHandles,
                                     bool importBuffers) const override;
 
+    [[nodiscard]] GraphicBufferAllocator::AllocationResult allocate(
+            const GraphicBufferAllocator::AllocationRequest&) const override;
+
 private:
     const Gralloc5Mapper &mMapper;
     std::shared_ptr<aidl::android::hardware::graphics::allocator::IAllocator> mAllocator;
diff --git a/libs/ui/include/ui/GraphicBuffer.h b/libs/ui/include/ui/GraphicBuffer.h
index f859848..652d8ba 100644
--- a/libs/ui/include/ui/GraphicBuffer.h
+++ b/libs/ui/include/ui/GraphicBuffer.h
@@ -26,6 +26,7 @@
 
 #include <android/hardware_buffer.h>
 #include <ui/ANativeObjectBase.h>
+#include <ui/GraphicBufferAllocator.h>
 #include <ui/GraphicBufferMapper.h>
 #include <ui/PixelFormat.h>
 #include <ui/Rect.h>
@@ -103,6 +104,8 @@
             uint32_t inLayerCount, uint64_t inUsage,
             std::string requestorName = "<Unknown>");
 
+    GraphicBuffer(const GraphicBufferAllocator::AllocationRequest&);
+
     // Create a GraphicBuffer from an existing handle.
     enum HandleWrapMethod : uint8_t {
         // Wrap and use the handle directly.  It assumes the handle has been
@@ -169,6 +172,8 @@
         mGenerationNumber = generation;
     }
 
+    status_t getDataspace(ui::Dataspace* outDataspace) const;
+
     // This function is privileged.  It requires access to the allocator
     // device or service, which usually involves adding suitable selinux
     // rules.
diff --git a/libs/ui/include/ui/GraphicBufferAllocator.h b/libs/ui/include/ui/GraphicBufferAllocator.h
index 3ed988c..8f461e1 100644
--- a/libs/ui/include/ui/GraphicBufferAllocator.h
+++ b/libs/ui/include/ui/GraphicBufferAllocator.h
@@ -22,6 +22,7 @@
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include <cutils/native_handle.h>
 
@@ -42,6 +43,35 @@
 public:
     static inline GraphicBufferAllocator& get() { return getInstance(); }
 
+    struct AdditionalOptions {
+        const char* name;
+        int64_t value;
+    };
+
+    struct AllocationRequest {
+        bool importBuffer;
+        uint32_t width;
+        uint32_t height;
+        PixelFormat format;
+        uint32_t layerCount;
+        uint64_t usage;
+        std::string requestorName;
+        std::vector<AdditionalOptions> extras;
+    };
+
+    struct AllocationResult {
+        status_t status;
+        buffer_handle_t handle = nullptr;
+        uint32_t stride = 0;
+
+        explicit AllocationResult(status_t status) : status(status) {}
+
+        explicit AllocationResult(buffer_handle_t handle, uint32_t stride)
+              : status(OK), handle(handle), stride(stride) {}
+    };
+
+    AllocationResult allocate(const AllocationRequest&);
+
     /**
      * Allocates and imports a gralloc buffer.
      *
diff --git a/libs/ui/tests/GraphicBufferAllocator_test.cpp b/libs/ui/tests/GraphicBufferAllocator_test.cpp
index f4c0afa..efca083 100644
--- a/libs/ui/tests/GraphicBufferAllocator_test.cpp
+++ b/libs/ui/tests/GraphicBufferAllocator_test.cpp
@@ -51,7 +51,7 @@
         std::cout << "Setting expected stride to " << stride << std::endl;
         EXPECT_CALL(*(reinterpret_cast<const mock::MockGrallocAllocator*>(mAllocator.get())),
                     allocate)
-                .WillOnce(DoAll(SetArgPointee<7>(stride), Return(err)));
+                .WillOnce(DoAll(SetArgPointee<6>(stride), Return(err)));
     }
     std::unique_ptr<const GrallocAllocator>& getAllocator() { return mAllocator; }
 };
diff --git a/libs/ui/tests/mock/MockGrallocAllocator.h b/libs/ui/tests/mock/MockGrallocAllocator.h
index d62e3e2..d02b387 100644
--- a/libs/ui/tests/mock/MockGrallocAllocator.h
+++ b/libs/ui/tests/mock/MockGrallocAllocator.h
@@ -35,7 +35,7 @@
     MOCK_METHOD(std::string, dumpDebugInfo, (bool less), (const, override));
     MOCK_METHOD(status_t, allocate,
                 (std::string requestorName, uint32_t width, uint32_t height, PixelFormat format,
-                 uint32_t layerCount, uint64_t usage, uint32_t bufferCount, uint32_t* outStride,
+                 uint32_t layerCount, uint64_t usage, uint32_t* outStride,
                  buffer_handle_t* outBufferHandles, bool less),
                 (const, override));
 };