From 70f35b06c6330d6b46c0689a93545ef33041da79 Mon Sep 17 00:00:00 2001
From: Eugene Zemtsov <[email protected]>
Date: Wed, 01 Apr 2026 01:41:59 -0700
Subject: [PATCH] media: Fix UAF bug in MediaFoundationVideoEncodeAccelerator

The MF VEA used fake COM refcounting while registering its `this`
pointer with the  Media Foundation API. This led to a use after free bug
if the encoder was deleted while an OS callback was pending.

This change decouples the encoder's lifetime from the OS callback's
lifetime by removing `IMFAsyncCallback` from the main class and
introducing `MFAsyncCallbackProxy`, a true COM object. The proxy
receives the events and safely posts them to the parent class via
`base::WeakPtr`, ensuring that the main object is not accessed after
it is freed.

Bug: 497151750
Change-Id: I29d6969cc5db8d30e9e2ee2793d34ab326d560f3
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7712934
Reviewed-by: Dale Curtis <[email protected]>
Commit-Queue: Eugene Zemtsov <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1608366}
---

diff --git a/media/gpu/windows/media_foundation_video_encode_accelerator_win.cc b/media/gpu/windows/media_foundation_video_encode_accelerator_win.cc
index 36c3960c..58983a5 100644
--- a/media/gpu/windows/media_foundation_video_encode_accelerator_win.cc
+++ b/media/gpu/windows/media_foundation_video_encode_accelerator_win.cc
@@ -11,6 +11,7 @@
 #include <mfapi.h>
 #include <mferror.h>
 #include <mftransform.h>
+#include <wrl/implements.h>
 
 #include <algorithm>
 #include <iterator>
@@ -212,6 +213,67 @@
 
 }  // namespace
 
+// A proxy class that implements IMFAsyncCallback and routes the events back to
+// the MediaFoundationVideoEncodeAccelerator safely via a WeakPtr. This
+// decouples the encoder's lifetime from the OS callback's lifetime. If the
+// encoder is destroyed while a callback is pending, the WeakPtr will be
+// invalidated, and the posted task will be safely dropped, preventing a
+// use-after-free.
+class MFAsyncCallbackProxy
+    : public Microsoft::WRL::RuntimeClass<
+          Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>,
+          IMFAsyncCallback> {
+ public:
+  MFAsyncCallbackProxy(
+      scoped_refptr<base::SequencedTaskRunner> task_runner,
+      base::WeakPtr<MediaFoundationVideoEncodeAccelerator> parent)
+      : task_runner_(std::move(task_runner)),
+        parent_weak_ptr_(std::move(parent)) {}
+
+  ~MFAsyncCallbackProxy() override = default;
+
+  IFACEMETHODIMP GetParameters(DWORD* pdwFlags, DWORD* pdwQueue) override {
+    *pdwFlags = MFASYNC_FAST_IO_PROCESSING_CALLBACK;
+    *pdwQueue = MFASYNC_CALLBACK_QUEUE_TIMER;
+    return S_OK;
+  }
+
+  IFACEMETHODIMP Invoke(IMFAsyncResult* pAsyncResult) override {
+    MediaEventType event_type = MEUnknown;
+    HRESULT status = GetEvent(pAsyncResult, &event_type);
+
+    // Invoke() is called on some random OS thread, so we must post to our event
+    // handler since MediaFoundationVideoEncodeAccelerator is single threaded.
+    task_runner_->PostTask(
+        FROM_HERE,
+        base::BindOnce(
+            &MediaFoundationVideoEncodeAccelerator::MediaEventHandler,
+            parent_weak_ptr_, event_type, status));
+    return S_OK;
+  }
+
+ private:
+  HRESULT GetEvent(IMFAsyncResult* pAsyncResult, MediaEventType* event_type) {
+    Microsoft::WRL::ComPtr<IUnknown> state;
+    RETURN_IF_FAILED(pAsyncResult->GetState(&state));
+
+    Microsoft::WRL::ComPtr<IMFMediaEventGenerator> event_generator;
+    RETURN_IF_FAILED(state.As(&event_generator));
+
+    Microsoft::WRL::ComPtr<IMFMediaEvent> media_event;
+    RETURN_IF_FAILED(event_generator->EndGetEvent(pAsyncResult, &media_event));
+
+    RETURN_IF_FAILED(media_event->GetType(event_type));
+
+    HRESULT status = S_OK;
+    RETURN_IF_FAILED(media_event->GetStatus(&status));
+    return status;
+  }
+
+  scoped_refptr<base::SequencedTaskRunner> task_runner_;
+  base::WeakPtr<MediaFoundationVideoEncodeAccelerator> parent_weak_ptr_;
+};
+
 struct MediaFoundationVideoEncodeAccelerator::PendingInput {
   PendingInput() = default;
   ~PendingInput() = default;
@@ -280,7 +342,6 @@
     ~MediaFoundationVideoEncodeAccelerator() {
   DVLOG(3) << __func__;
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  DCHECK(async_callback_ref_.IsOne());
 }
 
 VideoEncodeAccelerator::SupportedProfiles
@@ -617,7 +678,16 @@
     return false;
   }
 
-  event_generator_->BeginGetEvent(this, nullptr);
+  proxy_callback_ = Microsoft::WRL::Make<MFAsyncCallbackProxy>(
+      task_runner_, weak_factory_.GetWeakPtr());
+
+  hr = event_generator_->BeginGetEvent(proxy_callback_.Get(),
+                                       event_generator_.Get());
+  if (FAILED(hr)) {
+    NotifyErrorStatus({EncoderStatus::Codes::kEncoderInitializationError,
+                       "Couldn't begin get event: " + PrintHr(hr)});
+    return false;
+  }
 
   // Start the asynchronous processing model
   hr = encoder_->ProcessMessage(MFT_MESSAGE_COMMAND_FLUSH, 0);
@@ -2615,7 +2685,12 @@
     default:
       break;
   }
-  event_generator_->BeginGetEvent(this, nullptr);
+  HRESULT hr = event_generator_->BeginGetEvent(proxy_callback_.Get(),
+                                               event_generator_.Get());
+  if (FAILED(hr)) {
+    NotifyErrorStatus({EncoderStatus::Codes::kSystemAPICallError,
+                       "Failed to begin get event: " + PrintHr(hr)});
+  }
 }
 
 void MediaFoundationVideoEncodeAccelerator::SetState(State state) {
@@ -3004,47 +3079,4 @@
   return initialized;
 }
 
-HRESULT MediaFoundationVideoEncodeAccelerator::GetParameters(DWORD* pdwFlags,
-                                                             DWORD* pdwQueue) {
-  *pdwFlags = MFASYNC_FAST_IO_PROCESSING_CALLBACK;
-  *pdwQueue = MFASYNC_CALLBACK_QUEUE_TIMER;
-  return S_OK;
-}
-
-HRESULT MediaFoundationVideoEncodeAccelerator::Invoke(
-    IMFAsyncResult* pAsyncResult) {
-  ComMFMediaEvent media_event;
-  RETURN_IF_FAILED(event_generator_->EndGetEvent(pAsyncResult, &media_event));
-
-  MediaEventType event_type = MEUnknown;
-  RETURN_IF_FAILED(media_event->GetType(&event_type));
-
-  HRESULT status = S_OK;
-  media_event->GetStatus(&status);
-
-  // Invoke() is called on some random OS thread, so we must post to our event
-  // handler since MediaFoundationVideoEncodeAccelerator is single threaded.
-  task_runner_->PostTask(
-      FROM_HERE,
-      base::BindOnce(&MediaFoundationVideoEncodeAccelerator::MediaEventHandler,
-                     weak_ptr_, event_type, status));
-  return status;
-}
-
-ULONG MediaFoundationVideoEncodeAccelerator::AddRef() {
-  return async_callback_ref_.Increment();
-}
-
-ULONG MediaFoundationVideoEncodeAccelerator::Release() {
-  DCHECK(!async_callback_ref_.IsOne());
-  return async_callback_ref_.Decrement() ? 1 : 0;
-}
-
-HRESULT MediaFoundationVideoEncodeAccelerator::QueryInterface(REFIID riid,
-                                                              void** ppv) {
-  static const QITAB kQI[] = {
-      QITABENT(MediaFoundationVideoEncodeAccelerator, IMFAsyncCallback), {0}};
-  return QISearch(this, kQI, riid, ppv);
-}
-
 }  // namespace media
diff --git a/media/gpu/windows/media_foundation_video_encode_accelerator_win.h b/media/gpu/windows/media_foundation_video_encode_accelerator_win.h
index e59a26c2..746b4cd 100644
--- a/media/gpu/windows/media_foundation_video_encode_accelerator_win.h
+++ b/media/gpu/windows/media_foundation_video_encode_accelerator_win.h
@@ -45,6 +45,7 @@
 class VideoRateControlWrapper;
 class TemporalScalabilityIdExtractor;
 class VEAEncodingLatencyMetricsHelper;