Fix UAF in adapter lifecycle: serialize Adapters ops with a mutex

src/models/adapters.cpp

@@ -10,25 +10,32 @@ Adapter::Adapter(const char* adapter_file_path, Ort::Allocator* allocator)
     : adapter_{OrtLoraAdapter::Create(fs::path(adapter_file_path).c_str(), *allocator)} {}
 
 const OrtLoraAdapter* Adapter::AcquireRef() {
+  // Caller (Adapters::AcquireAdapter) holds Adapters::mutex_, which
+  // serializes all access to ref_count_.
   ref_count_++;
-
   return adapter_.get();
 }
 
 void Adapter::ReleaseRef() {
+  // Caller (Adapters::ReleaseAdapter) holds Adapters::mutex_.
   ref_count_--;
   if (ref_count_ < 0) {
+    // Restore invariant so a caller catching the exception doesn't leave the
+    // counter in a negative state that would trip later releases too.
+    ref_count_++;
     throw std::runtime_error("Adapter ref count went negative.");
   }
 }
 
 int32_t Adapter::RefCount() const {
+  // Caller (Adapters::UnloadAdapter) holds Adapters::mutex_.
   return ref_count_;
 }
 
 Adapters::Adapters(const Model* model) : model_{model} {}
 
 void Adapters::LoadAdapter(const char* adapter_file_path, const std::string& adapter_name) {
+  std::lock_guard<std::mutex> lock(mutex_);
   if (adapters_.find(adapter_name) != adapters_.end()) {
     throw std::runtime_error("Adapter already loaded: " + std::string{adapter_name});
   }
@@ -40,11 +47,16 @@ void Adapters::LoadAdapter(const char* adapter_file_path, const std::string& ada
 }
 
 void Adapters::UnloadAdapter(const std::string& adapter_name) {
+  std::lock_guard<std::mutex> lock(mutex_);
   auto adapter = adapters_.find(adapter_name);
   if (adapter == adapters_.end()) {
     throw std::runtime_error("Adapter not found: " + std::string{adapter_name});
   }
 
+  // Check-and-erase must happen atomically with respect to AcquireAdapter /
+  // ReleaseAdapter, which also acquire mutex_. This closes the TOCTOU window
+  // where another thread could AcquireRef() between the RefCount() check and
+  // the erase(), producing a use-after-free.
   if (adapter->second->RefCount() > 0) {
     throw std::runtime_error("Adapter still in use: " + std::string{adapter_name});
   }
@@ -53,6 +65,7 @@ void Adapters::UnloadAdapter(const std::string& adapter_name) {
 }
 
 const OrtLoraAdapter* Adapters::AcquireAdapter(const std::string& adapter_name) {
+  std::lock_guard<std::mutex> lock(mutex_);
   auto adapter = adapters_.find(adapter_name);
   if (adapter == adapters_.end()) {
     throw std::runtime_error("Adapter not found: " + std::string{adapter_name});
@@ -62,6 +75,7 @@ const OrtLoraAdapter* Adapters::AcquireAdapter(const std::string& adapter_name)
 }
 
 void Adapters::ReleaseAdapter(const std::string& adapter_name) {
+  std::lock_guard<std::mutex> lock(mutex_);
   auto adapter = adapters_.find(adapter_name);
   if (adapter == adapters_.end()) {
     throw std::runtime_error("Adapter not found: " + std::string{adapter_name});

src/models/adapters.h

@@ -20,6 +20,10 @@ struct Adapter {
   int32_t RefCount() const;
 
  private:
+  // All access to ref_count_ is serialized by Adapters::mutex_, which is the
+  // sole entry point for AcquireRef/ReleaseRef/RefCount. This closes both the
+  // data race on the counter and the TOCTOU window between RefCount() and
+  // container erasure in Adapters::UnloadAdapter().
   int32_t ref_count_{};
   std::unique_ptr<OrtLoraAdapter> adapter_;
 };
@@ -41,6 +45,11 @@ struct Adapters : std::enable_shared_from_this<Adapters>, ExternalRefCounted<Ada
 
  private:
   const Model* model_;
+  // Serializes all access to adapters_ and to per-Adapter ref counts so that
+  // load/unload/acquire/release cannot race. Without this, the check-then-erase
+  // pattern in UnloadAdapter (and concurrent std::unordered_map mutation) is a
+  // use-after-free hazard.
+  mutable std::mutex mutex_;
   std::unordered_map<std::string, std::unique_ptr<Adapter>> adapters_;
 };