From 272d05a4fb92225b9b12501f4555e371d9b78de5 Mon Sep 17 00:00:00 2001
From: Deepmalya Acharya <deepmalyaa@gmail.com>
Date: Thu, 29 Jan 2026 01:55:16 +0530
Subject: [PATCH] Fix UAF when iterator outlives DB during shutdown

---
 db/db_impl.cc | 31 ++++++++++++++++++++-----------
 db/db_impl.h  | 11 +++++++++++
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/db/db_impl.cc b/db/db_impl.cc
index f96d245583..fe7af2b8cd 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -147,7 +147,10 @@ DBImpl::DBImpl(const Options& raw_options, const std::string& dbname)
       background_compaction_scheduled_(false),
       manual_compaction_(nullptr),
       versions_(new VersionSet(dbname_, &options_, table_cache_,
-                               &internal_comparator_)) {}
+                               &internal_comparator_)) {
+  // Initialize mutex wrapper used for lifetime extension of iterators
+  mutex_state_ = std::make_shared<SharedMutexState>();
+}
 
 DBImpl::~DBImpl() {
   // Wait for background work to finish.
@@ -1058,22 +1061,28 @@ Status DBImpl::DoCompactionWork(CompactionState* compact) {
 namespace {
 
 struct IterState {
-  port::Mutex* const mu;
-  Version* const version GUARDED_BY(mu);
-  MemTable* const mem GUARDED_BY(mu);
-  MemTable* const imm GUARDED_BY(mu);
-
-  IterState(port::Mutex* mutex, MemTable* mem, MemTable* imm, Version* version)
-      : mu(mutex), version(version), mem(mem), imm(imm) {}
+  std::shared_ptr<SharedMutexState> shared;  // shared mutex state
+  Version* const version GUARDED_BY(shared);
+  MemTable* const mem GUARDED_BY(shared);
+  MemTable* const imm GUARDED_BY(shared);
+
+  IterState(std::shared_ptr<SharedMutexState> s, MemTable* mem, MemTable* imm,
+            Version* version)
+      : shared(std::move(s)), version(version), mem(mem), imm(imm) {}
 };
 
 static void CleanupIteratorState(void* arg1, void* arg2) {
   IterState* state = reinterpret_cast<IterState*>(arg1);
-  state->mu->Lock();
+  // Use the shared mutex if available to guard cleanup; otherwise skip locking.
+  if (state->shared) {
+    state->shared->mutex.Lock();
+  }
   state->mem->Unref();
   if (state->imm != nullptr) state->imm->Unref();
   state->version->Unref();
-  state->mu->Unlock();
+  if (state->shared) {
+    state->shared->mutex.Unlock();
+  }
   delete state;
 }
 
@@ -1098,7 +1107,7 @@ Iterator* DBImpl::NewInternalIterator(const ReadOptions& options,
       NewMergingIterator(&internal_comparator_, &list[0], list.size());
   versions_->current()->Ref();
 
-  IterState* cleanup = new IterState(&mutex_, mem_, imm_, versions_->current());
+  IterState* cleanup = new IterState(mutex_state_, mem_, imm_, versions_->current());
   internal_iter->RegisterCleanup(CleanupIteratorState, cleanup, nullptr);
 
   *seed = ++seed_;
diff --git a/db/db_impl.h b/db/db_impl.h
index c7b01721b8..938b6bb8a7 100644
--- a/db/db_impl.h
+++ b/db/db_impl.h
@@ -9,6 +9,7 @@
 #include <deque>
 #include <set>
 #include <string>
+#include <memory>
 
 #include "db/dbformat.h"
 #include "db/log_writer.h"
@@ -20,6 +21,14 @@
 
 namespace leveldb {
 
+// Lightweight shared state wrapper to extend the lifetime of the mutex used by
+// iterators. This helps prevent Use-After-Free scenarios when a DB instance is
+// destroyed while iterators created from it are still alive.
+struct SharedMutexState {
+  port::Mutex mutex;
+};
+
+
 class MemTable;
 class TableCache;
 class Version;
@@ -172,6 +181,8 @@ class DBImpl : public DB {
 
   // State below is protected by mutex_
   port::Mutex mutex_;
+  // Shared state wrapper to extend mutex lifetime across iterators.
+  std::shared_ptr<SharedMutexState> mutex_state_;
   std::atomic<bool> shutting_down_;
   port::CondVar background_work_finished_signal_ GUARDED_BY(mutex_);
   MemTable* mem_;