[Asan] Return nullptr for invalid chunks

[vitalybuka]

Summary

CHUNK_ALLOCATED. CHUNK_QUARANTINE are only states
which make AsanChunk useful for GetAsanChunk callers.
In either case member of AsanChunk are not useful.

Fix few cases which didn't expect nullptr. Most of the callers are already
expects nullptr.

[morehouse]

compiler-rt/lib/asan/asan_allocator.cpp

757

Are we guaranteed that chunk_state == CHUNK_INVALID when p = alloc_beg?

If so, this seems tricky. Please add a comment explaining why this is the case.

[vitalybuka]

compiler-rt/lib/asan/asan_allocator.cpp

757

Actually I cannot to guaranty that.
Still the patch should be useful.

In followup I'll try to guaranty.

[vitalybuka]

fixes

[morehouse]

LGTM

compiler-rt/lib/asan/asan_allocator.cpp

struct Allocator {
   void SetRssLimitExceeded(bool limit_exceeded) {
     atomic_store(&rss_limit_exceeded, limit_exceeded, memory_order_relaxed);
   }
 
   void RePoisonChunk(uptr chunk) {
     // This could be a user-facing chunk (with redzones), or some internal
     // housekeeping chunk, like TransferBatch. Start by assuming the former.
     AsanChunk *ac = GetAsanChunk((void *)chunk);
-    uptr allocated_size = allocator.GetActuallyAllocatedSize((void *)ac);
-    if (atomic_load(&ac->chunk_state, memory_order_acquire) ==
+    uptr allocated_size = allocator.GetActuallyAllocatedSize((void *)chunk);
+    if (ac && atomic_load(&ac->chunk_state, memory_order_acquire) ==
                   CHUNK_ALLOCATED) {
       uptr beg = ac->Beg();
       uptr end = ac->Beg() + ac->UsedSize(true);
       uptr chunk_end = chunk + allocated_size;
       if (chunk < beg && beg < end && end <= chunk_end) {
         // Looks like a valid AsanChunk in use, poison redzones only.
         PoisonShadow(chunk, beg - chunk, kAsanHeapLeftRedzoneMagic);
         uptr end_aligned_down = RoundDownTo(end, SHADOW_GRANULARITY);
         FastPoisonShadowPartialRightRedzone(
static uptr ComputeUserAlignment(uptr user_requested_alignment_log) {
     if (user_requested_alignment_log == 0)
       return 0;
     return 1LL << (user_requested_alignment_log + 2);
   }
 
   // We have an address between two chunks, and we want to report just one.
   AsanChunk *ChooseChunk(uptr addr, AsanChunk *left_chunk,
                          AsanChunk *right_chunk) {
+    if (!left_chunk)
+      return right_chunk;
+    if (!right_chunk)
+      return left_chunk;
     // Prefer an allocated chunk over freed chunk and freed chunk
     // over available chunk.
     u8 left_state = atomic_load(&left_chunk->chunk_state, memory_order_relaxed);
     u8 right_state =
         atomic_load(&right_chunk->chunk_state, memory_order_relaxed);
     if (left_state != right_state) {
       if (left_state == CHUNK_ALLOCATED)
         return left_chunk;
#endif
 
   // Assumes alloc_beg == allocator.GetBlockBegin(alloc_beg).
   // Returns nullptr if AsanChunk is not yet initialized just after
   // get_allocator().Allocate(), or is being destroyed just before
   // get_allocator().Deallocate().
   AsanChunk *GetAsanChunk(void *alloc_beg) {
     if (!alloc_beg)
       return nullptr;
+    AsanChunk *p = nullptr;
     if (!allocator.FromPrimary(alloc_beg)) {
       uptr *meta = reinterpret_cast<uptr *>(allocator.GetMetaData(alloc_beg));
-      AsanChunk *m = reinterpret_cast<AsanChunk *>(meta[1]);
-      return m;
-    }
+      p = reinterpret_cast<AsanChunk *>(meta[1]);
+    } else {
       uptr *alloc_magic = reinterpret_cast<uptr *>(alloc_beg);
       if (alloc_magic[0] == kAllocBegMagic)
-      return reinterpret_cast<AsanChunk *>(alloc_magic[1]);
-    // FIXME: This is either valid small chunk with tiny redzone or invalid
-    // chunk which is beeing allocated/deallocated. The latter case should
-    // return nullptr like secondary allocator does.
-    return reinterpret_cast<AsanChunk *>(alloc_beg);
+        p = reinterpret_cast<AsanChunk *>(alloc_magic[1]);
+      else
+        p = reinterpret_cast<AsanChunk *>(alloc_beg);
+    }
+    if (!p)
+      return nullptr;
+    u8 state = atomic_load(&p->chunk_state, memory_order_relaxed);
+    // It does not guaranty that Chunk is initialized, but it's
+    // definitely not for any other value.
+    if (state == CHUNK_ALLOCATED || state == CHUNK_QUARANTINE)
+      return p;
+    return nullptr;
   }
 
   AsanChunk *GetAsanChunkByAddr(uptr p) {
     void *alloc_beg = allocator.GetBlockBegin(reinterpret_cast<void *>(p));
     return GetAsanChunk(alloc_beg);
   }
 
   // Allocator must be locked when this function is called.
uptr AllocationSize(uptr p) {
     if (atomic_load(&m->chunk_state, memory_order_acquire) != CHUNK_ALLOCATED)
       return 0;
     if (m->Beg() != p) return 0;
     return m->UsedSize();
   }
 
   AsanChunkView FindHeapChunkByAddress(uptr addr) {
     AsanChunk *m1 = GetAsanChunkByAddr(addr);
-    if (!m1) return AsanChunkView(m1);
     sptr offset = 0;
-    if (AsanChunkView(m1).AddrIsAtLeft(addr, 1, &offset)) {
+    if (!m1 || AsanChunkView(m1).AddrIsAtLeft(addr, 1, &offset)) {
       // The address is in the chunk's left redzone, so maybe it is actually
       // a right buffer overflow from the other chunk to the left.
       // Search a bit to the left to see if there is another chunk.
       AsanChunk *m2 = nullptr;
       for (uptr l = 1; l < GetPageSizeCached(); l++) {
         m2 = GetAsanChunkByAddr(addr - l);
         if (m2 == m1) continue;  // Still the same chunk.
         break;