diff --git a/drivers/staging/android/ion/heaps/ion_system_heap.c b/drivers/staging/android/ion/heaps/ion_system_heap.c
index 4885702f5d16..e0cc4bab86b8 100644
--- a/drivers/staging/android/ion/heaps/ion_system_heap.c
+++ b/drivers/staging/android/ion/heaps/ion_system_heap.c
@@ -288,7 +288,7 @@ static int ion_system_heap_allocate(struct ion_heap *heap,
 		return -ENOMEM;
 
 	if (ion_heap_is_system_heap_type(buffer->heap->type) &&
-	    is_secure_vmid_valid(vmid)) {
+	    is_secure_allocation(buffer->flags)) {
 		pr_info("%s: System heap doesn't support secure allocations\n",
 			__func__);
 		return -EINVAL;
diff --git a/drivers/staging/android/ion/heaps/ion_system_secure_heap.c b/drivers/staging/android/ion/heaps/ion_system_secure_heap.c
index 188b28f6c5d5..7b252d7ad169 100644
--- a/drivers/staging/android/ion/heaps/ion_system_secure_heap.c
+++ b/drivers/staging/android/ion/heaps/ion_system_secure_heap.c
@@ -79,9 +79,10 @@ static int ion_system_secure_heap_allocate(struct ion_heap *heap,
 	struct ion_system_secure_heap *secure_heap =
 		to_system_secure_heap(heap);
 	enum ion_heap_type type = secure_heap->heap.ion_heap.type;
+	unsigned long cp_flags = buffer->flags & ION_FLAGS_CP_MASK;
 
 	if (!ion_heap_is_system_secure_heap_type(type) ||
-	    !is_cp_flag_present(flags)) {
+	    !is_cp_flag_present(flags) || (hweight_long(cp_flags) != 1)) {
 		pr_info("%s: Incorrect heap type or incorrect flags\n",
 			__func__);
 		return -EINVAL;