Fix handler unregistration in C++ channels

ci/licenses_golden/licenses_flutter

@@ -752,6 +752,7 @@ FILE: ../../../flutter/shell/platform/android/platform_view_android_jni.cc
 FILE: ../../../flutter/shell/platform/android/platform_view_android_jni.h
 FILE: ../../../flutter/shell/platform/android/vsync_waiter_android.cc
 FILE: ../../../flutter/shell/platform/android/vsync_waiter_android.h
+FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/basic_message_channel_unittests.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/byte_stream_wrappers.h
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/encodable_value_unittests.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/engine_method_result.cc
@@ -774,6 +775,7 @@ FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/include/flutter/
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/json_message_codec.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/json_method_codec.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/method_call_unittests.cc
+FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/method_channel_unittests.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/plugin_registrar.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/plugin_registrar_unittests.cc
 FILE: ../../../flutter/shell/platform/common/cpp/client_wrapper/standard_codec.cc

shell/platform/common/cpp/client_wrapper/BUILD.gn

@@ -46,8 +46,10 @@ executable("client_wrapper_unittests") {
 
   # TODO: Add more unit tests.
   sources = [
+    "basic_message_channel_unittests.cc",
     "encodable_value_unittests.cc",
     "method_call_unittests.cc",
+    "method_channel_unittests.cc",
     "plugin_registrar_unittests.cc",
     "standard_message_codec_unittests.cc",
     "standard_method_codec_unittests.cc",

shell/platform/common/cpp/client_wrapper/basic_message_channel_unittests.cc

@@ -0,0 +1,93 @@
+// Copyright 2013 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "flutter/shell/platform/common/cpp/client_wrapper/include/flutter/basic_message_channel.h"
+
+#include <memory>
+#include <string>
+
+#include "flutter/shell/platform/common/cpp/client_wrapper/include/flutter/binary_messenger.h"
+#include "flutter/shell/platform/common/cpp/client_wrapper/include/flutter/standard_message_codec.h"
+#include "gtest/gtest.h"
+
+namespace flutter {
+
+namespace {
+
+class TestBinaryMessenger : public BinaryMessenger {
+ public:
+  void Send(const std::string& channel,
+            const uint8_t* message,
+            const size_t message_size) const override {}
+
+  void Send(const std::string& channel,
+            const uint8_t* message,
+            const size_t message_size,
+            BinaryReply reply) const override {}
+
+  void SetMessageHandler(const std::string& channel,
+                         BinaryMessageHandler handler) override {
+    last_message_handler_channel_ = channel;
+    last_message_handler_ = handler;
+  }
+
+  std::string last_message_handler_channel() {
+    return last_message_handler_channel_;
+  }
+
+  BinaryMessageHandler last_message_handler() { return last_message_handler_; }
+
+ private:
+  std::string last_message_handler_channel_;
+  BinaryMessageHandler last_message_handler_;
+};
+
+}  // namespace
+
+// Tests that SetMessageHandler sets a handler that correctly interacts with
+// the binary messenger.
+TEST(BasicMessageChannelTest, Registration) {
+  TestBinaryMessenger messenger;
+  const std::string channel_name("some_channel");
+  const StandardMessageCodec& codec = StandardMessageCodec::GetInstance();
+  BasicMessageChannel channel(&messenger, channel_name, &codec);
+
+  bool callback_called = false;
+  const std::string message_value("hello");
+  channel.SetMessageHandler(
+      [&callback_called, message_value](const auto& message, auto reply) {
+        callback_called = true;
+        // Ensure that the wrapper recieved a correctly decoded message and a
+        // reply.
+        EXPECT_EQ(message.StringValue(), message_value);
+        EXPECT_NE(reply, nullptr);
+      });
+  EXPECT_EQ(messenger.last_message_handler_channel(), channel_name);
+  EXPECT_NE(messenger.last_message_handler(), nullptr);
+  // Send a test message to trigger the handler test assertions.
+  auto message = codec.EncodeMessage(EncodableValue(message_value));
+
+  messenger.last_message_handler()(
+      message->data(), message->size(),
+      [](const uint8_t* reply, const size_t reply_size) {});
+  EXPECT_EQ(callback_called, true);
+}
+
+// Tests that SetMessageHandler with a null handler unregisters the handler.
+TEST(BasicMessageChannelTest, Unregistration) {
+  TestBinaryMessenger messenger;
+  const std::string channel_name("some_channel");
+  BasicMessageChannel channel(&messenger, channel_name,
+                              &flutter::StandardMessageCodec::GetInstance());
+
+  channel.SetMessageHandler([](const auto& message, auto reply) {});
+  EXPECT_EQ(messenger.last_message_handler_channel(), channel_name);
+  EXPECT_NE(messenger.last_message_handler(), nullptr);
+
+  channel.SetMessageHandler(nullptr);
+  EXPECT_EQ(messenger.last_message_handler_channel(), channel_name);
+  EXPECT_EQ(messenger.last_message_handler(), nullptr);
+}
+
+}  // namespace flutter

shell/platform/common/cpp/client_wrapper/include/flutter/basic_message_channel.h

@@ -60,8 +60,16 @@ class BasicMessageChannel {
   }
 
   // Registers a handler that should be called any time a message is
-  // received on this channel.
+  // received on this channel. A null handler will remove any previous handler.
+  //
+  // Note that the BasicMessageChannel does not own the handler, and will not
+  // unregister it on destruction, so the caller is responsible for
+  // unregistering explicitly if it should no longer be called.
   void SetMessageHandler(const MessageHandler<T>& handler) const {
+    if (!handler) {
+      messenger_->SetMessageHandler(name_, nullptr);
+      return;
+    }
     const auto* codec = codec_;
     std::string channel_name = name_;
     BinaryMessageHandler binary_handler = [handler, codec, channel_name](

shell/platform/common/cpp/client_wrapper/include/flutter/method_channel.h

@@ -62,8 +62,16 @@ class MethodChannel {
   }
 
   // Registers a handler that should be called any time a method call is
-  // received on this channel.
+  // received on this channel. A null handler will remove any previous handler.
+  //
+  // Note that the MethodChannel does not own the handler, and will not
+  // unregister it on destruction, so the caller is responsible for
+  // unregistering explicitly if it should no longer be called.
   void SetMethodCallHandler(MethodCallHandler<T> handler) const {
+    if (!handler) {
+      messenger_->SetMessageHandler(name_, nullptr);
+      return;
+    }
     const auto* codec = codec_;
     std::string channel_name = name_;
     BinaryMessageHandler binary_handler = [handler, codec, channel_name](

shell/platform/common/cpp/client_wrapper/method_channel_unittests.cc

@@ -0,0 +1,94 @@
+// Copyright 2013 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "flutter/shell/platform/common/cpp/client_wrapper/include/flutter/method_channel.h"
+
+#include <memory>
+#include <string>
+
+#include "flutter/shell/platform/common/cpp/client_wrapper/include/flutter/binary_messenger.h"
+#include "flutter/shell/platform/common/cpp/client_wrapper/include/flutter/standard_method_codec.h"
+#include "gtest/gtest.h"
+
+namespace flutter {
+
+namespace {
+
+class TestBinaryMessenger : public BinaryMessenger {
+ public:
+  void Send(const std::string& channel,
+            const uint8_t* message,
+            const size_t message_size) const override {}
+
+  void Send(const std::string& channel,
+            const uint8_t* message,
+            const size_t message_size,
+            BinaryReply reply) const override {}
+
+  void SetMessageHandler(const std::string& channel,
+                         BinaryMessageHandler handler) override {
+    last_message_handler_channel_ = channel;
+    last_message_handler_ = handler;
+  }
+
+  std::string last_message_handler_channel() {
+    return last_message_handler_channel_;
+  }
+
+  BinaryMessageHandler last_message_handler() { return last_message_handler_; }
+
+ private:
+  std::string last_message_handler_channel_;
+  BinaryMessageHandler last_message_handler_;
+};
+
+}  // namespace
+
+// Tests that SetMethodCallHandler sets a handler that correctly interacts with
+// the binary messenger.
+TEST(MethodChannelTest, Registration) {
+  TestBinaryMessenger messenger;
+  const std::string channel_name("some_channel");
+  const StandardMethodCodec& codec = StandardMethodCodec::GetInstance();
+  MethodChannel channel(&messenger, channel_name, &codec);
+
+  bool callback_called = false;
+  const std::string method_name("hello");
+  channel.SetMethodCallHandler(
+      [&callback_called, method_name](const auto& call, auto result) {
+        callback_called = true;
+        // Ensure that the wrapper recieved a correctly decoded call and a
+        // result.
+        EXPECT_EQ(call.method_name(), method_name);
+        EXPECT_NE(result, nullptr);
+      });
+  EXPECT_EQ(messenger.last_message_handler_channel(), channel_name);
+  EXPECT_NE(messenger.last_message_handler(), nullptr);
+  // Send a test message to trigger the handler test assertions.
+  MethodCall<EncodableValue> call(method_name, nullptr);
+  auto message = codec.EncodeMethodCall(call);
+
+  messenger.last_message_handler()(
+      message->data(), message->size(),
+      [](const uint8_t* reply, const size_t reply_size) {});
+  EXPECT_EQ(callback_called, true);
+}
+
+// Tests that SetMethodCallHandler with a null handler unregisters the handler.
+TEST(MethodChannelTest, Unregistration) {
+  TestBinaryMessenger messenger;
+  const std::string channel_name("some_channel");
+  MethodChannel channel(&messenger, channel_name,
+                        &flutter::StandardMethodCodec::GetInstance());
+
+  channel.SetMethodCallHandler([](const auto& call, auto result) {});
+  EXPECT_EQ(messenger.last_message_handler_channel(), channel_name);
+  EXPECT_NE(messenger.last_message_handler(), nullptr);
+
+  channel.SetMethodCallHandler(nullptr);
+  EXPECT_EQ(messenger.last_message_handler_channel(), channel_name);
+  EXPECT_EQ(messenger.last_message_handler(), nullptr);
+}
+
+}  // namespace flutter

shell/platform/common/cpp/client_wrapper/plugin_registrar_unittests.cc

@@ -23,6 +23,7 @@ class TestApi : public testing::StubFlutterApi {
     last_data_sent_ = message;
     return message_engine_result;
   }
+
   bool MessengerSendWithReply(const char* channel,
                               const uint8_t* message,
                               const size_t message_size,
@@ -32,15 +33,27 @@ class TestApi : public testing::StubFlutterApi {
     return message_engine_result;
   }
 
+  // Called for FlutterDesktopMessengerSetCallback.
+  void MessengerSetCallback(const char* channel,
+                            FlutterDesktopMessageCallback callback,
+                            void* user_data) override {
+    last_callback_set_ = callback;
+  }
+
   const uint8_t* last_data_sent() { return last_data_sent_; }
+  FlutterDesktopMessageCallback last_callback_set() {
+    return last_callback_set_;
+  }
 
  private:
   const uint8_t* last_data_sent_ = nullptr;
+  FlutterDesktopMessageCallback last_callback_set_ = nullptr;
 };
 
 }  // namespace
 
-// Tests that the registrar returns a messenger that calls through to the C API.
+// Tests that the registrar returns a messenger that passes Send through to the
+// C API.
 TEST(MethodCallTest, MessengerSend) {
   testing::ScopedStubFlutterApi scoped_api_stub(std::make_unique<TestApi>());
   auto test_api = static_cast<TestApi*>(scoped_api_stub.stub());
@@ -55,4 +68,28 @@ TEST(MethodCallTest, MessengerSend) {
   EXPECT_EQ(test_api->last_data_sent(), &message[0]);
 }
 
+// Tests that the registrar returns a messenger that passes callback
+// registration and unregistration through to the C API.
+TEST(MethodCallTest, MessengerSetMessageHandler) {
+  testing::ScopedStubFlutterApi scoped_api_stub(std::make_unique<TestApi>());
+  auto test_api = static_cast<TestApi*>(scoped_api_stub.stub());
+
+  auto dummy_registrar_handle =
+      reinterpret_cast<FlutterDesktopPluginRegistrarRef>(1);
+  PluginRegistrar registrar(dummy_registrar_handle);
+  BinaryMessenger* messenger = registrar.messenger();
+  const std::string channel_name("foo");
+
+  // Register.
+  BinaryMessageHandler binary_handler = [](const uint8_t* message,
+                                           const size_t message_size,
+                                           BinaryReply reply) {};
+  messenger->SetMessageHandler(channel_name, std::move(binary_handler));
+  EXPECT_NE(test_api->last_callback_set(), nullptr);
+
+  // Unregister.
+  messenger->SetMessageHandler(channel_name, nullptr);
+  EXPECT_EQ(test_api->last_callback_set(), nullptr);
+}
+
 }  // namespace flutter