From 781e3c8ee33e2566fe561948c4128b2105ad690d Mon Sep 17 00:00:00 2001
From: Daniel Cheng <dcheng@chromium.org>
Date: Fri, 30 Oct 2026 08:34:11 -0700
Subject: [PATCH] Handle reentrant destruction from clients in `mojo::DataPipeDrainer`

`DataPipeDrainer::ReadData()` called `source_->EndReadData()`
unconditionally after `client_->OnDataAvailable()`; however, a client
that reentrantly destroyed the `mojo::DataPipeDrainer` would trigger a
use-after-free. Avoiding this use-after-free requires additional
handling (https://crrev.com/c/2672379, https://crrev.com/c/7781209);
forgetting to handle this edge case results in a security bug instead
(crbug.com/550181232, crbug.com/566136673).

Since `mojo::DataPipeDrainer` is intended for general consumption, it
should handle this more gracefully instead of requiring all of its
callers to be careful. Add a WeakPtr guard and document the lifetime and
validity expectations.

Bug: 566347710
Change-Id: I7f43003e044757c91c1fee967eae4e465eee1892
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8683777
Reviewed-by: Tom Sepez <tsepez@chromium.org>
Commit-Queue: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1731506}
---
diff --git a/mojo/public/cpp/system/data_pipe_drainer.cc b/mojo/public/cpp/system/data_pipe_drainer.cc
index 0d2262b..60d53ee 100644
--- a/mojo/public/cpp/system/data_pipe_drainer.cc
+++ b/mojo/public/cpp/system/data_pipe_drainer.cc
@@ -32,7 +32,11 @@
   base::span<const uint8_t> buffer;
   MojoResult rv = source_->BeginReadData(MOJO_READ_DATA_FLAG_NONE, buffer);
   if (rv == MOJO_RESULT_OK) {
+    base::WeakPtr<DataPipeDrainer> weak_self = weak_factory_.GetWeakPtr();
     client_->OnDataAvailable(buffer);
+    if (!weak_self) {
+      return;
+    }
     source_->EndReadData(buffer.size());
   } else if (rv == MOJO_RESULT_FAILED_PRECONDITION) {
     client_->OnDataComplete();
diff --git a/mojo/public/cpp/system/data_pipe_drainer.h b/mojo/public/cpp/system/data_pipe_drainer.h
index 741d786..4f1730ee 100644
--- a/mojo/public/cpp/system/data_pipe_drainer.h
+++ b/mojo/public/cpp/system/data_pipe_drainer.h
@@ -5,7 +5,7 @@
 #ifndef MOJO_PUBLIC_CPP_SYSTEM_DATA_PIPE_DRAINER_H_
 #define MOJO_PUBLIC_CPP_SYSTEM_DATA_PIPE_DRAINER_H_
 
-#include <stddef.h>
+#include <stdint.h>
 
 #include "base/containers/span.h"
 #include "base/memory/raw_ptr.h"
@@ -20,6 +20,9 @@
  public:
   class Client {
    public:
+    // The memory backing `data` is transitively owned by the `DataPipeDrainer`
+    // and is guaranteed to be valid until `OnDataAvailable()` returns, or until
+    // the `DataPipeDrainer` is destroyed.
     virtual void OnDataAvailable(base::span<const uint8_t> data) = 0;
     virtual void OnDataComplete() = 0;
 
@@ -27,7 +30,8 @@
     virtual ~Client() {}
   };
 
-  DataPipeDrainer(Client*, mojo::ScopedDataPipeConsumerHandle source);
+  // `client` must outlive this `DataPipeDrainer`.
+  DataPipeDrainer(Client* client, mojo::ScopedDataPipeConsumerHandle source);
 
   DataPipeDrainer(const DataPipeDrainer&) = delete;
   DataPipeDrainer& operator=(const DataPipeDrainer&) = delete;
diff --git a/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc b/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc
index 8683385..71ebadad 100644
--- a/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc
+++ b/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc
@@ -4,72 +4,191 @@
 
 #include "mojo/public/cpp/system/data_pipe_drainer.h"
 
+#include <stdint.h>
+
+#include <memory>
+#include <optional>
+#include <string>
 #include <string_view>
+#include <utility>
 
 #include "base/containers/span.h"
-#include "base/functional/bind.h"
 #include "base/functional/callback.h"
 #include "base/run_loop.h"
 #include "base/strings/string_view_util.h"
+#include "base/test/bind.h"
 #include "base/test/task_environment.h"
-#include "base/values.h"
+#include "mojo/public/c/system/types.h"
+#include "mojo/public/cpp/system/data_pipe.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace mojo {
-namespace test {
+namespace {
 
-template <typename Functor>
-base::RepeatingClosure BindLambda(Functor callable) {
-  return base::BindRepeating([](Functor callable) { callable(); }, callable);
-}
-
-class DataPipeDrainerTest : public testing::Test,
-                            public DataPipeDrainer::Client {
+class TestClient : public DataPipeDrainer::Client {
  public:
-  DataPipeDrainerTest(const DataPipeDrainerTest&) = delete;
-  DataPipeDrainerTest& operator=(const DataPipeDrainerTest&) = delete;
-
- protected:
-  DataPipeDrainerTest() {
-    ScopedDataPipeProducerHandle producer_handle;
-    ScopedDataPipeConsumerHandle consumer_handle;
-    EXPECT_EQ(CreateDataPipe(nullptr, producer_handle, consumer_handle),
-              MOJO_RESULT_OK);
-    drainer_ =
-        std::make_unique<DataPipeDrainer>(this, std::move(consumer_handle));
-    producer_handle_ = std::move(producer_handle);
+  explicit TestClient(ScopedDataPipeConsumerHandle consumer) {
+    drainer_.emplace(this, std::move(consumer));
   }
 
-  ScopedDataPipeProducerHandle producer_handle_;
-  base::RepeatingClosure completion_callback_;
+  TestClient(const TestClient&) = delete;
+  TestClient& operator=(const TestClient&) = delete;
+
+  bool WaitForData(std::string_view expected_data) {
+    while (data_ != expected_data && !is_complete_) {
+      data_run_loop_.emplace();
+      data_run_loop_->Run();
+    }
+    return data_ == expected_data;
+  }
+
+  void WaitForComplete() { completion_run_loop_.Run(); }
 
   void OnDataAvailable(base::span<const uint8_t> data) override {
     data_.append(base::as_string_view(data));
+    if (data_run_loop_) {
+      data_run_loop_->Quit();
+    }
+    if (data_available_callback_) {
+      std::move(data_available_callback_).Run();
+    }
   }
 
-  void OnDataComplete() override { completion_callback_.Run(); }
+  void OnDataComplete() override {
+    EXPECT_FALSE(is_complete_);
+    is_complete_ = true;
+    if (data_run_loop_) {
+      data_run_loop_->Quit();
+    }
+    completion_run_loop_.Quit();
+    if (completion_callback_) {
+      std::move(completion_callback_).Run();
+    }
+  }
 
-  base::test::SingleThreadTaskEnvironment task_environment_;
+  void ResetDrainer() { drainer_.reset(); }
+  bool has_drainer() const { return drainer_.has_value(); }
+  const std::string& data() const { return data_; }
+
+  void set_data_available_callback(base::RepeatingClosure callback) {
+    data_available_callback_ = std::move(callback);
+  }
+  void set_completion_callback(base::RepeatingClosure callback) {
+    completion_callback_ = std::move(callback);
+  }
+
+ private:
   std::string data_;
-  std::unique_ptr<DataPipeDrainer> drainer_;
+  bool is_complete_ = false;
+  std::optional<base::RunLoop> data_run_loop_;
+  base::RunLoop completion_run_loop_;
+  base::RepeatingClosure data_available_callback_;
+  base::RepeatingClosure completion_callback_;
+  std::optional<DataPipeDrainer> drainer_;
+};
+
+class DataPipeDrainerTest : public testing::Test {
+ public:
+  DataPipeDrainerTest() {
+    ScopedDataPipeConsumerHandle consumer_handle;
+    EXPECT_EQ(CreateDataPipe(nullptr, producer_handle_, consumer_handle),
+              MOJO_RESULT_OK);
+    client_ = std::make_unique<TestClient>(std::move(consumer_handle));
+  }
+
+ protected:
+  base::test::SingleThreadTaskEnvironment task_environment_;
+  ScopedDataPipeProducerHandle producer_handle_;
+  std::unique_ptr<TestClient> client_;
 };
 
 TEST_F(DataPipeDrainerTest, TestCompleteIsCalledOnce) {
-  bool had_data_complete = false;
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  EXPECT_TRUE(client_->WaitForData("hello"));
 
-  completion_callback_ = BindLambda([&had_data_complete]() {
-    EXPECT_FALSE(had_data_complete);
-    had_data_complete = true;
-  });
-  size_t bytes_written = 0;
-  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteData(
-                                base::byte_span_from_cstring("hello"),
-                                MOJO_WRITE_DATA_FLAG_NONE, bytes_written));
-  EXPECT_EQ(bytes_written, 5u);
-  base::RunLoop().RunUntilIdle();
   producer_handle_.reset();
-  base::RunLoop().RunUntilIdle();
+  client_->WaitForComplete();
 }
 
-}  // namespace test
+TEST_F(DataPipeDrainerTest, DestroyDrainerOnDataAvailable) {
+  client_->set_data_available_callback(
+      base::BindLambdaForTesting([&]() { client_->ResetDrainer(); }));
+  client_->set_completion_callback(base::BindLambdaForTesting(
+      []() { ADD_FAILURE() << "OnDataComplete should not be called."; }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+
+  EXPECT_TRUE(client_->WaitForData("hello"));
+  EXPECT_FALSE(client_->has_drainer());
+  EXPECT_EQ(
+      MOJO_RESULT_FAILED_PRECONDITION,
+      producer_handle_->WriteAllData(base::byte_span_from_cstring("world")));
+}
+
+TEST_F(DataPipeDrainerTest, DestroyDrainerOnDataAvailableWithQueuedEOF) {
+  client_->set_data_available_callback(
+      base::BindLambdaForTesting([&]() { client_->ResetDrainer(); }));
+  client_->set_completion_callback(base::BindLambdaForTesting(
+      []() { ADD_FAILURE() << "OnDataComplete should not be called."; }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+
+  EXPECT_TRUE(client_->WaitForData("hello"));
+  EXPECT_FALSE(client_->has_drainer());
+}
+
+TEST_F(DataPipeDrainerTest, DestroyDrainerOnDataComplete) {
+  client_->set_completion_callback(
+      base::BindLambdaForTesting([&]() { client_->ResetDrainer(); }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+  client_->WaitForComplete();
+
+  EXPECT_EQ(client_->data(), "hello");
+  EXPECT_FALSE(client_->has_drainer());
+}
+
+TEST_F(DataPipeDrainerTest, DestroyClientOnDataAvailable) {
+  base::RunLoop run_loop;
+  std::string captured_data;
+  client_->set_data_available_callback(base::BindLambdaForTesting([&]() {
+    captured_data = client_->data();
+    client_.reset();
+    run_loop.Quit();
+  }));
+  client_->set_completion_callback(base::BindLambdaForTesting(
+      []() { ADD_FAILURE() << "OnDataComplete should not be called."; }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+  run_loop.Run();
+
+  EXPECT_EQ(captured_data, "hello");
+  EXPECT_EQ(client_, nullptr);
+}
+
+TEST_F(DataPipeDrainerTest, DestroyClientOnDataComplete) {
+  base::RunLoop run_loop;
+  client_->set_completion_callback(base::BindLambdaForTesting([&]() {
+    EXPECT_EQ(client_->data(), "hello");
+    client_.reset();
+    run_loop.Quit();
+  }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+  run_loop.Run();
+
+  EXPECT_EQ(client_, nullptr);
+}
+
+}  // namespace
 }  // namespace mojo

diff --git a/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc b/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc
index 8683385..71ebadad 100644
--- a/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc
+++ b/mojo/public/cpp/system/tests/data_pipe_drainer_unittest.cc
@@ -4,72 +4,191 @@
 
 #include "mojo/public/cpp/system/data_pipe_drainer.h"
 
+#include <stdint.h>
+
+#include <memory>
+#include <optional>
+#include <string>
 #include <string_view>
+#include <utility>
 
 #include "base/containers/span.h"
-#include "base/functional/bind.h"
 #include "base/functional/callback.h"
 #include "base/run_loop.h"
 #include "base/strings/string_view_util.h"
+#include "base/test/bind.h"
 #include "base/test/task_environment.h"
-#include "base/values.h"
+#include "mojo/public/c/system/types.h"
+#include "mojo/public/cpp/system/data_pipe.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace mojo {
-namespace test {
+namespace {
 
-template <typename Functor>
-base::RepeatingClosure BindLambda(Functor callable) {
-  return base::BindRepeating([](Functor callable) { callable(); }, callable);
-}
-
-class DataPipeDrainerTest : public testing::Test,
-                            public DataPipeDrainer::Client {
+class TestClient : public DataPipeDrainer::Client {
  public:
-  DataPipeDrainerTest(const DataPipeDrainerTest&) = delete;
-  DataPipeDrainerTest& operator=(const DataPipeDrainerTest&) = delete;
-
- protected:
-  DataPipeDrainerTest() {
-    ScopedDataPipeProducerHandle producer_handle;
-    ScopedDataPipeConsumerHandle consumer_handle;
-    EXPECT_EQ(CreateDataPipe(nullptr, producer_handle, consumer_handle),
-              MOJO_RESULT_OK);
-    drainer_ =
-        std::make_unique<DataPipeDrainer>(this, std::move(consumer_handle));
-    producer_handle_ = std::move(producer_handle);
+  explicit TestClient(ScopedDataPipeConsumerHandle consumer) {
+    drainer_.emplace(this, std::move(consumer));
   }
 
-  ScopedDataPipeProducerHandle producer_handle_;
-  base::RepeatingClosure completion_callback_;
+  TestClient(const TestClient&) = delete;
+  TestClient& operator=(const TestClient&) = delete;
+
+  bool WaitForData(std::string_view expected_data) {
+    while (data_ != expected_data && !is_complete_) {
+      data_run_loop_.emplace();
+      data_run_loop_->Run();
+    }
+    return data_ == expected_data;
+  }
+
+  void WaitForComplete() { completion_run_loop_.Run(); }
 
   void OnDataAvailable(base::span<const uint8_t> data) override {
     data_.append(base::as_string_view(data));
+    if (data_run_loop_) {
+      data_run_loop_->Quit();
+    }
+    if (data_available_callback_) {
+      std::move(data_available_callback_).Run();
+    }
   }
 
-  void OnDataComplete() override { completion_callback_.Run(); }
+  void OnDataComplete() override {
+    EXPECT_FALSE(is_complete_);
+    is_complete_ = true;
+    if (data_run_loop_) {
+      data_run_loop_->Quit();
+    }
+    completion_run_loop_.Quit();
+    if (completion_callback_) {
+      std::move(completion_callback_).Run();
+    }
+  }
 
-  base::test::SingleThreadTaskEnvironment task_environment_;
+  void ResetDrainer() { drainer_.reset(); }
+  bool has_drainer() const { return drainer_.has_value(); }
+  const std::string& data() const { return data_; }
+
+  void set_data_available_callback(base::RepeatingClosure callback) {
+    data_available_callback_ = std::move(callback);
+  }
+  void set_completion_callback(base::RepeatingClosure callback) {
+    completion_callback_ = std::move(callback);
+  }
+
+ private:
   std::string data_;
-  std::unique_ptr<DataPipeDrainer> drainer_;
+  bool is_complete_ = false;
+  std::optional<base::RunLoop> data_run_loop_;
+  base::RunLoop completion_run_loop_;
+  base::RepeatingClosure data_available_callback_;
+  base::RepeatingClosure completion_callback_;
+  std::optional<DataPipeDrainer> drainer_;
+};
+
+class DataPipeDrainerTest : public testing::Test {
+ public:
+  DataPipeDrainerTest() {
+    ScopedDataPipeConsumerHandle consumer_handle;
+    EXPECT_EQ(CreateDataPipe(nullptr, producer_handle_, consumer_handle),
+              MOJO_RESULT_OK);
+    client_ = std::make_unique<TestClient>(std::move(consumer_handle));
+  }
+
+ protected:
+  base::test::SingleThreadTaskEnvironment task_environment_;
+  ScopedDataPipeProducerHandle producer_handle_;
+  std::unique_ptr<TestClient> client_;
 };
 
 TEST_F(DataPipeDrainerTest, TestCompleteIsCalledOnce) {
-  bool had_data_complete = false;
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  EXPECT_TRUE(client_->WaitForData("hello"));
 
-  completion_callback_ = BindLambda([&had_data_complete]() {
-    EXPECT_FALSE(had_data_complete);
-    had_data_complete = true;
-  });
-  size_t bytes_written = 0;
-  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteData(
-                                base::byte_span_from_cstring("hello"),
-                                MOJO_WRITE_DATA_FLAG_NONE, bytes_written));
-  EXPECT_EQ(bytes_written, 5u);
-  base::RunLoop().RunUntilIdle();
   producer_handle_.reset();
-  base::RunLoop().RunUntilIdle();
+  client_->WaitForComplete();
 }
 
-}  // namespace test
+TEST_F(DataPipeDrainerTest, DestroyDrainerOnDataAvailable) {
+  client_->set_data_available_callback(
+      base::BindLambdaForTesting([&]() { client_->ResetDrainer(); }));
+  client_->set_completion_callback(base::BindLambdaForTesting(
+      []() { ADD_FAILURE() << "OnDataComplete should not be called."; }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+
+  EXPECT_TRUE(client_->WaitForData("hello"));
+  EXPECT_FALSE(client_->has_drainer());
+  EXPECT_EQ(
+      MOJO_RESULT_FAILED_PRECONDITION,
+      producer_handle_->WriteAllData(base::byte_span_from_cstring("world")));
+}
+
+TEST_F(DataPipeDrainerTest, DestroyDrainerOnDataAvailableWithQueuedEOF) {
+  client_->set_data_available_callback(
+      base::BindLambdaForTesting([&]() { client_->ResetDrainer(); }));
+  client_->set_completion_callback(base::BindLambdaForTesting(
+      []() { ADD_FAILURE() << "OnDataComplete should not be called."; }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+
+  EXPECT_TRUE(client_->WaitForData("hello"));
+  EXPECT_FALSE(client_->has_drainer());
+}
+
+TEST_F(DataPipeDrainerTest, DestroyDrainerOnDataComplete) {
+  client_->set_completion_callback(
+      base::BindLambdaForTesting([&]() { client_->ResetDrainer(); }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+  client_->WaitForComplete();
+
+  EXPECT_EQ(client_->data(), "hello");
+  EXPECT_FALSE(client_->has_drainer());
+}
+
+TEST_F(DataPipeDrainerTest, DestroyClientOnDataAvailable) {
+  base::RunLoop run_loop;
+  std::string captured_data;
+  client_->set_data_available_callback(base::BindLambdaForTesting([&]() {
+    captured_data = client_->data();
+    client_.reset();
+    run_loop.Quit();
+  }));
+  client_->set_completion_callback(base::BindLambdaForTesting(
+      []() { ADD_FAILURE() << "OnDataComplete should not be called."; }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+  run_loop.Run();
+
+  EXPECT_EQ(captured_data, "hello");
+  EXPECT_EQ(client_, nullptr);
+}
+
+TEST_F(DataPipeDrainerTest, DestroyClientOnDataComplete) {
+  base::RunLoop run_loop;
+  client_->set_completion_callback(base::BindLambdaForTesting([&]() {
+    EXPECT_EQ(client_->data(), "hello");
+    client_.reset();
+    run_loop.Quit();
+  }));
+
+  EXPECT_EQ(MOJO_RESULT_OK, producer_handle_->WriteAllData(
+                                base::byte_span_from_cstring("hello")));
+  producer_handle_.reset();
+  run_loop.Run();
+
+  EXPECT_EQ(client_, nullptr);
+}
+
+}  // namespace
 }  // namespace mojo