From 88b3c1acf9cfc8dd6957671d5c3ab4bcbf53c429 Mon Sep 17 00:00:00 2001
From: Dmitry Muhomor <muhomor.dmitry@gmail.com>
Date: Wed, 1 Nov 2023 21:33:54 +0200
Subject: [PATCH] memtag_test: fix sporadic failures of overflow/underflow
 tests

---
 androidtest/memtag/memtag_test.cc | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/androidtest/memtag/memtag_test.cc b/androidtest/memtag/memtag_test.cc
index e6ec68d..ca491d8 100644
--- a/androidtest/memtag/memtag_test.cc
+++ b/androidtest/memtag/memtag_test.cc
@@ -10,6 +10,7 @@
 #include <unistd.h>
 
 #include <map>
+#include <set>
 #include <string>
 #include <unordered_map>
 
@@ -181,9 +182,26 @@ void tag_distinctness() {
 }
 
 u8* alloc_default() {
-    u8 *p = (u8 *) malloc(DEFAULT_ALLOC_SIZE);
-    assert(p);
-    return p;
+    const size_t full_alloc_size = DEFAULT_ALLOC_SIZE + CANARY_SIZE;
+    set<uptr> addrs;
+
+    // make sure allocation has both left and right neighbors, otherwise overflow/underflow tests
+    // will fail when allocation is at the end/beginning of slab
+    for (;;) {
+        u8 *p = (u8 *) malloc(DEFAULT_ALLOC_SIZE);
+        assert(p);
+
+        uptr addr = (uptr) untag_pointer(p);
+        uptr left = addr - full_alloc_size;
+        if (addrs.find(left) != addrs.end()) {
+            uptr right = addr + full_alloc_size;
+            if (addrs.find(right) != addrs.end()) {
+                return p;
+            }
+        }
+
+        addrs.emplace(addr);
+    }
 }
 
 volatile u8 u8_var;