From 38acc8b39421655426c218609986023ad483c241 Mon Sep 17 00:00:00 2001
From: Saifuddin Hitawala <[email protected]>
Date: Thu, 23 Jul 2026 12:29:16 -0700
Subject: [PATCH] [media] Keep ref on MemoryTypeTracker in SharedImageReadLock

Add MemoryTypeTracker ref on SharedImageReadLock to avoid UAF based
on unexpected destruction orders, as the lock can be passed to
IMFSample. We will eventually get rid of MemoryTypeTracker in
SharedImage to instead use MemoryTracker directly which is
RefCountedThreadSafe and should avoid future issues.

Note that this does not pass CommandBufferHelper as we want to get
rid of it as well eventually.

Bug: 537832446
Change-Id: I048a2e20527748c0e169d6b285d0cab5d6c22fa9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8139742
Reviewed-by: Dale Curtis <[email protected]>
Reviewed-by: Vasiliy Telezhnikov <[email protected]>
Commit-Queue: Saifuddin Hitawala <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1667317}
---

diff --git a/media/base/win/mf_helpers.cc b/media/base/win/mf_helpers.cc
index 4e3cfa2..00e115a 100644
--- a/media/base/win/mf_helpers.cc
+++ b/media/base/win/mf_helpers.cc
@@ -51,17 +51,16 @@
     std::unique_ptr<gpu::VideoImageRepresentation> representation,
     std::unique_ptr<gpu::VideoImageRepresentation::ScopedReadAccess>
         scoped_read_access,
-    scoped_refptr<gpu::SharedContextState> context_state) {
-  if (context_state && context_state->MakeCurrent(nullptr, /*needs_gl=*/true)) {
-    scoped_read_access.reset();
-    representation.reset();
-  } else {
-    if (representation) {
-      representation->OnContextLost();
-    }
-    scoped_read_access.reset();
-    representation.reset();
+    scoped_refptr<gpu::SharedContextState> context_state,
+    std::unique_ptr<gpu::MemoryTypeTracker> tracker) {
+  const bool context_current =
+      context_state && context_state->MakeCurrent(nullptr, /*needs_gl=*/true);
+  if (representation && !context_current) {
+    representation->OnContextLost();
   }
+  scoped_read_access.reset();
+  representation.reset();
+  tracker.reset();
 }
 
 }  // namespace
@@ -71,11 +70,13 @@
     std::unique_ptr<gpu::VideoImageRepresentation::ScopedReadAccess>
         scoped_read_access,
     scoped_refptr<VideoFrame> frame,
-    scoped_refptr<gpu::SharedContextState> context_state)
+    scoped_refptr<gpu::SharedContextState> context_state,
+    std::unique_ptr<gpu::MemoryTypeTracker> tracker)
     : representation_(std::move(representation)),
       scoped_read_access_(std::move(scoped_read_access)),
       frame_(std::move(frame)),
       context_state_(std::move(context_state)),
+      tracker_(std::move(tracker)),
       task_runner_(base::SequencedTaskRunner::GetCurrentDefault()) {}
 
 SharedImageReadLock::~SharedImageReadLock() {
@@ -83,7 +84,7 @@
       FROM_HERE,
       base::BindOnce(&DestroySharedImageResourcesOnGpuThread,
                      std::move(representation_), std::move(scoped_read_access_),
-                     std::move(context_state_)));
+                     std::move(context_state_), std::move(tracker_)));
 }
 
 using Microsoft::WRL::ComPtr;
@@ -1134,10 +1135,11 @@
   bool use_same_device = (encoder_device.Get() == shared_d3d11_device.Get());
   gpu::SharedImageManager* shared_image_manager =
       command_buffer_helper->GetSharedImageManager();
+  auto tracker = std::make_unique<gpu::MemoryTypeTracker>(
+      base::WrapRefCounted(shared_image_stub->memory_tracker()));
   std::unique_ptr<gpu::VideoImageRepresentation> image_representation =
       shared_image_manager->ProduceVideo(
-          shared_d3d11_device, frame->shared_image()->mailbox(),
-          command_buffer_helper->GetMemoryTypeTracker());
+          shared_d3d11_device, frame->shared_image()->mailbox(), tracker.get());
   RETURN_ON_FAILURE_WITH_CALLBACK(image_representation ? S_OK : E_FAIL,
                                   "Failed to produce video");
 
@@ -1152,7 +1154,7 @@
   ComPtr<SharedImageReadLock> si_lock =
       Microsoft::WRL::Make<SharedImageReadLock>(
           std::move(image_representation), std::move(scoped_read_access), frame,
-          std::move(shared_context_state));
+          std::move(shared_context_state), std::move(tracker));
   if (!si_lock) {
     RETURN_ON_FAILURE_WITH_CALLBACK(E_OUTOFMEMORY,
                                     "Failed to create SharedImageReadLock");
diff --git a/media/base/win/mf_helpers.h b/media/base/win/mf_helpers.h
index a5a2a42..73db391f 100644
--- a/media/base/win/mf_helpers.h
+++ b/media/base/win/mf_helpers.h
@@ -33,6 +33,7 @@
 class IMFMediaType;
 
 namespace gpu {
+class MemoryTypeTracker;
 class SharedContextState;
 }  // namespace gpu
 
@@ -261,7 +262,8 @@
       std::unique_ptr<gpu::VideoImageRepresentation::ScopedReadAccess>
           scoped_read_access,
       scoped_refptr<VideoFrame> frame,
-      scoped_refptr<gpu::SharedContextState> context_state);
+      scoped_refptr<gpu::SharedContextState> context_state,
+      std::unique_ptr<gpu::MemoryTypeTracker> tracker);
 
   gpu::VideoImageRepresentation::ScopedReadAccess* access() const {
     return scoped_read_access_.get();
@@ -275,6 +277,7 @@
       scoped_read_access_;
   scoped_refptr<VideoFrame> frame_;
   scoped_refptr<gpu::SharedContextState> context_state_;
+  std::unique_ptr<gpu::MemoryTypeTracker> tracker_;
   scoped_refptr<base::SequencedTaskRunner> task_runner_;
 };
 
diff --git a/media/gpu/windows/d3d12_video_encode_accelerator.cc b/media/gpu/windows/d3d12_video_encode_accelerator.cc
index e7b6bae..1762d35 100644
--- a/media/gpu/windows/d3d12_video_encode_accelerator.cc
+++ b/media/gpu/windows/d3d12_video_encode_accelerator.cc
@@ -149,10 +149,11 @@
 
   gpu::SharedImageManager* shared_image_manager =
       command_buffer_helper->GetSharedImageManager();
+  auto tracker = std::make_unique<gpu::MemoryTypeTracker>(
+      base::WrapRefCounted(shared_image_stub->memory_tracker()));
   std::unique_ptr<gpu::VideoImageRepresentation> representation =
       shared_image_manager->ProduceVideo(
-          d3d11_device, frame->shared_image()->mailbox(),
-          command_buffer_helper->GetMemoryTypeTracker());
+          d3d11_device, frame->shared_image()->mailbox(), tracker.get());
   RETURN_ON_FAILURE_WITH_CALLBACK(representation ? S_OK : E_FAIL,
                                   "Failed to produce video");
 
@@ -167,7 +168,7 @@
   Microsoft::WRL::ComPtr<SharedImageReadLock> si_lock =
       Microsoft::WRL::Make<SharedImageReadLock>(
           std::move(representation), std::move(scoped_read_access), frame,
-          std::move(shared_context_state));
+          std::move(shared_context_state), std::move(tracker));
   if (!si_lock) {
     RETURN_ON_FAILURE_WITH_CALLBACK(E_OUTOFMEMORY,
                                     "Failed to create SharedImageReadLock");