From c5e8f134f6d023faa41f10bbd28c478f0699b281 Mon Sep 17 00:00:00 2001
From: hai007 <hais@google.com>
Date: Wed, 3 Mar 2021 17:30:57 -0800
Subject: [PATCH] Internal change

PiperOrigin-RevId: 360794882
---
 cpp/core/internal/endpoint_manager.cc      | 12 ++++++++++++
 cpp/core/internal/endpoint_manager_test.cc |  9 +++++++++
 2 files changed, 21 insertions(+)

diff --git a/cpp/core/internal/endpoint_manager.cc b/cpp/core/internal/endpoint_manager.cc
index 4f20786e..f2ff0833 100644
--- a/cpp/core/internal/endpoint_manager.cc
+++ b/cpp/core/internal/endpoint_manager.cc
@@ -81,6 +81,18 @@ void EndpointManager::EndpointChannelLoopRunnable(
 
     if (!keep_using_channel.ok()) {
       Exception exception = keep_using_channel.GetException();
+      // An "invalid proto" may be a final payload on a channel we're about to
+      // close, so we'll loop back around once. We set |last_failed_medium| to
+      // ensure we don't loop indefinitely. See crbug.com/1182031 for more
+      // detail.
+      if (exception.Raised(Exception::kInvalidProtocolBuffer)) {
+        last_failed_medium = channel->GetMedium();
+        NEARBY_LOG(INFO,
+                   "Received invalid protobuf message, re-fetching endpoint "
+                   "channel; last_failed_medium=%d",
+                   last_failed_medium);
+        continue;
+      }
       if (exception.Raised(Exception::kIo)) {
         last_failed_medium = channel->GetMedium();
         NEARBY_LOG(INFO, "Endpoint channel IO exception; last_failed_medium=%d",
diff --git a/cpp/core/internal/endpoint_manager_test.cc b/cpp/core/internal/endpoint_manager_test.cc
index a43e0766..ae865bff 100644
--- a/cpp/core/internal/endpoint_manager_test.cc
+++ b/cpp/core/internal/endpoint_manager_test.cc
@@ -253,6 +253,15 @@ TEST_F(EndpointManagerTest, SendControlMessageWorks) {
   NEARBY_LOG(INFO, "Will call destructors now");
 }
 
+TEST_F(EndpointManagerTest, SingleReadOnInvalidPayload) {
+  auto endpoint_channel = std::make_unique<MockEndpointChannel>();
+  EXPECT_CALL(*endpoint_channel, Read())
+      .WillOnce(
+          Return(ExceptionOr<ByteArray>(Exception::kInvalidProtocolBuffer)));
+  EXPECT_CALL(*endpoint_channel, Close(_)).Times(1);
+  RegisterEndpoint(std::move(endpoint_channel));
+}
+
 }  // namespace
 }  // namespace connections
 }  // namespace nearby