Remove buffer reference counting from handle

Bug: 212803946
Test: Boots to home
Change-Id: I97a95c0c5738a1d7a40bc371b33582ec81d8f5b7

6 files changed, 219 insertions, 158 deletions

diff --git a/gralloc4/Android.bp b/gralloc4/Android.bp
index bbac925..8edc25a 100644
--- a/gralloc4/Android.bp
+++ b/gralloc4/Android.bp
@@ -70,6 +70,7 @@ arm_gralloc_cc_defaults {
 	cflags: [
 		"-Wundef",
 		"-Werror",
+		"-Wthread-safety",
 		"-DGRALLOC_LIBRARY_BUILD=1",
 		"-DDISABLE_FRAMEBUFFER_HAL=1",
 		"-DGRALLOC_USE_LEGACY_CALCS=0",
diff --git a/gralloc4/src/Android.bp b/gralloc4/src/Android.bp
index 12a3769..867ff2e 100644
--- a/gralloc4/src/Android.bp
+++ b/gralloc4/src/Android.bp
@@ -51,6 +51,7 @@ cc_library_shared {
 	],
 	cflags: [
 		"-DGRALLOC_LIBRARY_BUILD=1",
+		"-Wthread-safety",
 	],
 	static_libs: [
 		"libgralloc_capabilities",
diff --git a/gralloc4/src/core/mali_gralloc_bufferaccess.cpp b/gralloc4/src/core/mali_gralloc_bufferaccess.cpp
index 0c7dd7c..adda46b 100644
--- a/gralloc4/src/core/mali_gralloc_bufferaccess.cpp
+++ b/gralloc4/src/core/mali_gralloc_bufferaccess.cpp
@@ -113,8 +113,6 @@ int validate_lock_input_parameters(const buffer_handle_t buffer, const int l,
                                    const int t, const int w, const int h,
                                    uint64_t usage)
 {
-	bool is_registered_process = false;
-	const int lock_pid = getpid();
 	const private_handle_t * const hnd = (private_handle_t *)buffer;
 
 	/* TODO: do not check access region for blob formats.
@@ -126,8 +124,7 @@ int validate_lock_input_parameters(const buffer_handle_t buffer, const int l,
 		if ((l < 0) || (t < 0) || (w < 0) || (h < 0))
 		{
 			MALI_GRALLOC_LOGW("Negative values for access region (l = %d t = %d w = %d and "
-			     "h = %d) in buffer lock request are invalid. Locking PID:%d",
-			      l, t, w, h, lock_pid);
+			     "h = %d) in buffer lock request are invalid.", l, t, w, h);
 			return -EINVAL;
 		}
 
@@ -135,8 +132,7 @@ int validate_lock_input_parameters(const buffer_handle_t buffer, const int l,
 		if (((l + w) < 0) || ((t + h) < 0))
 		{
 			MALI_GRALLOC_LOGW("Encountered overflow with access region (l = %d t = %d w = %d and"
-			     " h = %d) in buffer lock request. Locking PID:%d",
-			       l, t, w, h, lock_pid);
+			     " h = %d) in buffer lock request.", l, t, w, h);
 			return -EINVAL;
 		}
 
@@ -144,28 +140,12 @@ int validate_lock_input_parameters(const buffer_handle_t buffer, const int l,
 		if (((t + h) > hnd->height)  || ((l + w) > hnd->width))
 		{
 			MALI_GRALLOC_LOGW("Buffer lock access region (l = %d t = %d w = %d "
-			     "and h = %d) is outside allocated buffer (width = %d and height = %d)"
-			     " Locking PID:%d", l, t, w, h, hnd->width, hnd->height, lock_pid);
+			     "and h = %d) is outside allocated buffer (width = %d and height = %d)",
+			     l, t, w, h, hnd->width, hnd->height);
 			return -EINVAL;
 		}
 	}
 
-	/* Locking process should have a valid buffer virtual address. A process
-	 * will have a valid buffer virtual address if it is the allocating
-	 * process or it retained / registered a cloned buffer handle
-	 */
-	if (hnd->remote_pid == lock_pid)
-	{
-		is_registered_process = true;
-	}
-
-	/* TODO: differentiate between mappable and unmappable(secure, hfr etc) buffers */
-	if (is_registered_process == false)
-	{
-		MALI_GRALLOC_LOGE("The buffer must be retained before lock request");
-		return -EINVAL;
-	}
-
 	/* Reject lock requests for AFBC (compressed format) enabled buffers */
 	if ((hnd->alloc_format & MALI_GRALLOC_INTFMT_EXT_MASK) != 0)
 	{
@@ -210,7 +190,7 @@ int mali_gralloc_lock(buffer_handle_t buffer,
 {
 	int status;
 
-	if (private_handle_t::validate(buffer) < 0)
+	if (mali_gralloc_reference_validate(buffer))
 	{
 		MALI_GRALLOC_LOGE("Locking invalid buffer %p, returning error", buffer);
 		return -EINVAL;
diff --git a/gralloc4/src/core/mali_gralloc_reference.cpp b/gralloc4/src/core/mali_gralloc_reference.cpp
index 074e189..6dc7ddd 100644
--- a/gralloc4/src/core/mali_gralloc_reference.cpp
+++ b/gralloc4/src/core/mali_gralloc_reference.cpp
@@ -16,141 +16,223 @@
  * limitations under the License.
  */
 
+#include "mali_gralloc_reference.h"
+
+#include <android-base/thread_annotations.h>
 #include <hardware/gralloc1.h>
 
-#include "mali_gralloc_buffer.h"
+#include <map>
+#include <mutex>
+
 #include "allocator/mali_gralloc_ion.h"
 #include "allocator/mali_gralloc_shared_memory.h"
+#include "mali_gralloc_buffer.h"
 #include "mali_gralloc_bufferallocation.h"
-#include "mali_gralloc_reference.h"
 #include "mali_gralloc_usages.h"
 
-static pthread_mutex_t s_map_lock = PTHREAD_MUTEX_INITIALIZER;
-
-int mali_gralloc_reference_retain(buffer_handle_t handle)
-{
-	if (private_handle_t::validate(handle) < 0)
-	{
-		MALI_GRALLOC_LOGE("Registering/Retaining invalid buffer %p, returning error", handle);
-		return -EINVAL;
-	}
-
-	private_handle_t *hnd = (private_handle_t *)handle;
-	pthread_mutex_lock(&s_map_lock);
-	int retval = 0;
-
-	if (hnd->remote_pid == getpid())
-	{
-		hnd->ref_count++;
-	}
-	else
-	{
-		hnd->remote_pid = getpid();
-		hnd->ref_count = 1;
-
-		// Reset the handle bases, this is used to check if a buffer is mapped
-		for (int fidx = 0; fidx < hnd->fd_count; fidx++) {
-			hnd->bases[fidx] = 0;
-		}
-	}
-
-	pthread_mutex_unlock(&s_map_lock);
-
-	// TODO(b/187145254): CPU_READ/WRITE buffer is not being properly locked from
-	// MFC. This is a WA for the time being.
-	constexpr auto cpu_access_usage = (
-			GRALLOC_USAGE_SW_WRITE_OFTEN |
-			GRALLOC_USAGE_SW_READ_OFTEN |
-			GRALLOC_USAGE_SW_WRITE_RARELY |
-			GRALLOC_USAGE_SW_READ_RARELY
-	);
-
-	if (hnd->get_usage() & cpu_access_usage)
-		retval = mali_gralloc_reference_map(handle);
-
-	return retval;
+class BufferManager {
+private:
+    // This struct for now is just for validation and reference counting. When we
+    // are sure that private_handle_t::bases is not being used outside gralloc, this
+    // should become the only place where address mapping is maintained and can be
+    // queried from.
+    struct MappedData {
+        void *bases[MAX_FDS] = {};
+        size_t alloc_sizes[MAX_FDS] = {};
+        uint64_t ref_count = 0;
+    };
+
+    BufferManager() = default;
+
+    std::mutex lock;
+    std::map<const private_handle_t *, std::unique_ptr<MappedData>> buffer_map GUARDED_BY(lock);
+
+    static bool should_map_dmabuf(buffer_handle_t handle) {
+        private_handle_t *hnd = (private_handle_t *)handle;
+
+        // TODO(b/187145254): CPU_READ/WRITE buffer is not being properly locked from
+        // MFC. This is a WA for the time being.
+        constexpr auto cpu_access_usage =
+                (GRALLOC_USAGE_SW_WRITE_OFTEN | GRALLOC_USAGE_SW_READ_OFTEN |
+                 GRALLOC_USAGE_SW_WRITE_RARELY | GRALLOC_USAGE_SW_READ_RARELY);
+        return hnd->get_usage() & cpu_access_usage;
+    }
+
+    int map_locked(buffer_handle_t handle) REQUIRES(lock) {
+        private_handle_t *hnd = (private_handle_t *)handle;
+        auto it = buffer_map.find(hnd);
+
+        if (it == buffer_map.end()) {
+            MALI_GRALLOC_LOGE("BUG: Map called without importing buffer");
+            return -EINVAL;
+        }
+
+        auto &data = *(it->second.get());
+        if (data.ref_count == 0) {
+            MALI_GRALLOC_LOGE("BUG: Found an imported buffer with ref count 0, expect errors");
+        }
+
+        if (data.bases[0] != nullptr) {
+            return 0;
+        }
+
+        int error = mali_gralloc_ion_map(hnd);
+        if (error != 0) {
+            return error;
+        }
+
+        for (auto i = 0; i < hnd->fd_count; i++) {
+            data.bases[i] = reinterpret_cast<void *>(hnd->bases[i]);
+            data.alloc_sizes[i] = hnd->alloc_sizes[i];
+        }
+
+        return 0;
+    }
+
+    int validate_locked(buffer_handle_t handle) REQUIRES(lock) {
+        if (private_handle_t::validate(handle) < 0) {
+            MALI_GRALLOC_LOGE("Reference invalid buffer %p, returning error", handle);
+            return -EINVAL;
+        }
+
+        const auto *hnd = (private_handle_t *)handle;
+        auto it = buffer_map.find(hnd);
+        if (it == buffer_map.end()) {
+            MALI_GRALLOC_LOGE("Reference unimported buffer %p, returning error", handle);
+            return -EINVAL;
+        }
+
+        auto &data = *(it->second.get());
+        if (data.bases[0] != nullptr) {
+            for (auto i = 0; i < hnd->fd_count; i++) {
+                if (data.bases[i] != reinterpret_cast<void *>(hnd->bases[i]) ||
+                    data.alloc_sizes[i] != hnd->alloc_sizes[i]) {
+                    MALI_GRALLOC_LOGE(
+                            "Validation failed: Buffer attributes inconsistent with mapper");
+                    return -EINVAL;
+                }
+            }
+        } else {
+            for (auto i = 0; i < hnd->fd_count; i++) {
+                if (hnd->bases[i] != 0 || data.bases[i] != nullptr) {
+                    MALI_GRALLOC_LOGE("Validation failed: Expected nullptr for unmaped buffer");
+                    return -EINVAL;
+                }
+            }
+        }
+
+        return 0;
+    }
+
+public:
+    static BufferManager &getInstance() {
+        static BufferManager instance;
+        return instance;
+    }
+
+    int retain(buffer_handle_t handle) EXCLUDES(lock) {
+        if (private_handle_t::validate(handle) < 0) {
+            MALI_GRALLOC_LOGE("Registering/Retaining invalid buffer %p, returning error", handle);
+            return -EINVAL;
+        }
+        std::lock_guard<std::mutex> _l(lock);
+
+        private_handle_t *hnd = (private_handle_t *)handle;
+
+        auto it = buffer_map.find(hnd);
+        if (it == buffer_map.end()) {
+            bool success = false;
+            auto _data = std::make_unique<MappedData>();
+
+            std::tie(it, success) = buffer_map.insert({hnd, std::move(_data)});
+            if (!success) {
+                MALI_GRALLOC_LOGE("Failed to create buffer data mapping");
+                return -EINVAL;
+            }
+
+            for (int i = 0; i < hnd->fd_count; i++) {
+                hnd->bases[i] = 0;
+            }
+        } else if (it->second->ref_count == 0) {
+            MALI_GRALLOC_LOGE("BUG: Import counter of an imported buffer is 0, expect errors");
+        }
+        auto &data = *(it->second.get());
+
+        data.ref_count++;
+        if (!should_map_dmabuf(handle)) return 0;
+        return map_locked(handle);
+    }
+
+    int map(buffer_handle_t handle) EXCLUDES(lock) {
+        std::lock_guard<std::mutex> _l(lock);
+        auto error = validate_locked(handle);
+        if (error != 0) {
+            return error;
+        }
+
+        return map_locked(handle);
+    }
+
+    int release(buffer_handle_t handle) EXCLUDES(lock) {
+        std::lock_guard<std::mutex> _l(lock);
+
+        // Always call locked variant of validate from this function. On calling
+        // the other validate variant, an attacker might launch a timing attack
+        // where they would try to time their attack between the return of
+        // validate and before taking the lock in this function again.
+        auto error = validate_locked(handle);
+        if (error != 0) {
+            return error;
+        }
+
+        private_handle_t *hnd = (private_handle_t *)handle;
+        auto it = buffer_map.find(hnd);
+        if (it == buffer_map.end()) {
+            MALI_GRALLOC_LOGE("Trying to release a non-imported buffer");
+            return -EINVAL;
+        }
+
+        auto &data = *(it->second.get());
+        if (data.ref_count == 0) {
+            MALI_GRALLOC_LOGE("BUG: Reference held for buffer whose counter is 0");
+            return -EINVAL;
+        }
+
+        data.ref_count--;
+        if (data.ref_count == 0) {
+            if (data.bases[0] != nullptr) {
+                mali_gralloc_ion_unmap(hnd);
+            }
+
+            /* TODO: Make this unmapping of shared meta fd into a function? */
+            if (hnd->attr_base) {
+                munmap(hnd->attr_base, hnd->attr_size);
+                hnd->attr_base = nullptr;
+            }
+            buffer_map.erase(it);
+        }
+
+        return 0;
+    }
+
+    int validate(buffer_handle_t handle) EXCLUDES(lock) {
+        std::lock_guard<std::mutex> _l(lock);
+        return validate_locked(handle);
+    }
+};
+
+int mali_gralloc_reference_retain(buffer_handle_t handle) {
+    return BufferManager::getInstance().retain(handle);
 }
 
 int mali_gralloc_reference_map(buffer_handle_t handle) {
-	private_handle_t *hnd = (private_handle_t *)handle;
-
-	pthread_mutex_lock(&s_map_lock);
-
-	if (hnd->bases[0]) {
-		MALI_GRALLOC_LOGV("Buffer is already mapped");
-		pthread_mutex_unlock(&s_map_lock);
-		return 0;
-	}
-
-	int retval = mali_gralloc_ion_map(hnd);
-
-	pthread_mutex_unlock(&s_map_lock);
-
-	return retval;
+    return BufferManager::getInstance().map(handle);
 }
 
-int mali_gralloc_reference_release(buffer_handle_t handle)
-{
-	if (private_handle_t::validate(handle) < 0)
-	{
-		MALI_GRALLOC_LOGE("unregistering/releasing invalid buffer %p, returning error", handle);
-		return -EINVAL;
-	}
-
-	private_handle_t *hnd = (private_handle_t *)handle;
-	pthread_mutex_lock(&s_map_lock);
-
-	if (hnd->ref_count == 0)
-	{
-		MALI_GRALLOC_LOGE("Buffer %p should have already been released", handle);
-		pthread_mutex_unlock(&s_map_lock);
-		return -EINVAL;
-	}
-
-	if (hnd->remote_pid == getpid()) // never unmap buffers that were not imported into this process
-	{
-		hnd->ref_count--;
-
-		if (hnd->ref_count == 0)
-		{
-			mali_gralloc_ion_unmap(hnd);
-
-			/* TODO: Make this unmapping of shared meta fd into a function? */
-			if (hnd->attr_base)
-			{
-				munmap(hnd->attr_base, hnd->attr_size);
-				hnd->attr_base = nullptr;
-			}
-		}
-	}
-	else
-	{
-		MALI_GRALLOC_LOGE("Trying to unregister buffer %p from process %d that was not imported into current process: %d", hnd,
-		     hnd->remote_pid, getpid());
-	}
-
-	pthread_mutex_unlock(&s_map_lock);
-	return 0;
+int mali_gralloc_reference_release(buffer_handle_t handle) {
+    return BufferManager::getInstance().release(handle);
 }
 
-int mali_gralloc_reference_validate(buffer_handle_t handle)
-{
-	if (private_handle_t::validate(handle) < 0)
-	{
-		MALI_GRALLOC_LOGE("Reference invalid buffer %p, returning error", handle);
-		return -EINVAL;
-	}
-
-	const auto *hnd = (private_handle_t *)handle;
-	pthread_mutex_lock(&s_map_lock);
-
-	if (hnd->remote_pid == getpid()) {
-		pthread_mutex_unlock(&s_map_lock);
-		return 0;
-	} else {
-		pthread_mutex_unlock(&s_map_lock);
-		MALI_GRALLOC_LOGE("Reference unimported buffer %p, returning error", handle);
-		return -EINVAL;
-	}
+int mali_gralloc_reference_validate(buffer_handle_t handle) {
+    return BufferManager::getInstance().validate(handle);
 }
-
diff --git a/gralloc4/src/mali_gralloc_buffer.h b/gralloc4/src/mali_gralloc_buffer.h
index d186f02..908c36c 100644
--- a/gralloc4/src/mali_gralloc_buffer.h
+++ b/gralloc4/src/mali_gralloc_buffer.h
@@ -50,6 +50,11 @@
  */
 #define MAX_PLANES 3
 
+/*
+ * Maximum number of fds in a private_handle_t.
+ */
+#define MAX_FDS 5
+
 #ifdef __cplusplus
 #define DEFAULT_INITIALIZER(x) = x
 #else
@@ -155,7 +160,7 @@ struct private_handle_t
 	 * DO NOT MOVE THIS ELEMENT!
 	 */
 	union {
-		int fds[5];
+		int fds[MAX_FDS];
 	};
 
 	// ints
@@ -215,8 +220,6 @@ struct private_handle_t
 	uint64_t backing_store_id DEFAULT_INITIALIZER(0x0);
 	int cpu_read DEFAULT_INITIALIZER(0);               /**< Buffer is locked for CPU read when non-zero. */
 	int cpu_write DEFAULT_INITIALIZER(0);              /**< Buffer is locked for CPU write when non-zero. */
-	int remote_pid DEFAULT_INITIALIZER(-1);
-	int ref_count DEFAULT_INITIALIZER(0);
 	// locally mapped shared attribute area
 
 	int ion_handles[3];
@@ -264,7 +267,6 @@ struct private_handle_t
 		stride = _stride;
 		alloc_format = _alloc_format;
 		layer_count = _layer_count;
-		ref_count = 1;
 		version = sizeof(native_handle);
 		set_numfds(fd_count);
 		memcpy(plane_info, _plane_info, sizeof(plane_info_t) * MAX_PLANES);
diff --git a/libvendorgraphicbuffer/gralloc4/vendor_graphicbuffer_meta.cpp b/libvendorgraphicbuffer/gralloc4/vendor_graphicbuffer_meta.cpp
index 030d509..a155519 100644
--- a/libvendorgraphicbuffer/gralloc4/vendor_graphicbuffer_meta.cpp
+++ b/libvendorgraphicbuffer/gralloc4/vendor_graphicbuffer_meta.cpp
@@ -46,12 +46,7 @@ using android::hardware::graphics::mapper::V4_0::Error;
 // Gralloc.
 int mali_gralloc_reference_validate(buffer_handle_t handle) {
 	auto hnd = static_cast<const private_handle_t *>(handle);
-
-	if (hnd->remote_pid != getpid()) {
-		return -EINVAL;
-	}
-
-	return 0;
+  return private_handle_t::validate(hnd);
 }
 
 const private_handle_t * convertNativeHandleToPrivateHandle(buffer_handle_t handle) {