From de83b242c527c5ac8957a8dd8c75aa2bad3891e7 Mon Sep 17 00:00:00 2001
From: hai007 <hais@google.com>
Date: Fri, 19 Feb 2021 11:28:26 -0800
Subject: [PATCH] 358287309 Cleanup PeerConnectionObserver so it can't touch a
 dead ConnectionFlow

---
 .../mediums/webrtc/connection_flow.cc         |  4 ++++
 .../webrtc/peer_connection_observer_impl.cc   | 23 +++++++++++++++----
 .../webrtc/peer_connection_observer_impl.h    |  6 +++--
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/cpp/core/internal/mediums/webrtc/connection_flow.cc b/cpp/core/internal/mediums/webrtc/connection_flow.cc
index 58c70649..995aac9b 100644
--- a/cpp/core/internal/mediums/webrtc/connection_flow.cc
+++ b/cpp/core/internal/mediums/webrtc/connection_flow.cc
@@ -329,11 +329,15 @@ bool ConnectionFlow::TransitionState(State current_state, State new_state) {
 }
 
 bool ConnectionFlow::CloseLocked() {
+  NEARBY_LOG(INFO, "Closing WebRTC connection.");
   if (state_ == State::kEnded) {
     return false;
   }
   state_ = State::kEnded;
 
+  single_threaded_signaling_offloader_.Shutdown();
+  peer_connection_observer_.Shutdown();
+
   if (peer_connection_) peer_connection_->Close();
 
   data_channel_observer_.reset();
diff --git a/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.cc b/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.cc
index 789f1c51..8d715284 100644
--- a/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.cc
+++ b/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.cc
@@ -14,6 +14,10 @@ PeerConnectionObserverImpl::PeerConnectionObserverImpl(
     : connection_flow_(connection_flow),
       local_ice_candidate_listener_(std::move(local_ice_candidate_listener)) {}
 
+PeerConnectionObserverImpl::~PeerConnectionObserverImpl() {
+    Shutdown();
+}
+
 void PeerConnectionObserverImpl::OnIceCandidate(
     const webrtc::IceCandidateInterface* candidate) {
   local_ice_candidate_listener_.local_ice_candidate_found_cb(candidate);
@@ -24,8 +28,10 @@ void PeerConnectionObserverImpl::OnSignalingChange(
   NEARBY_LOG(INFO, "OnSignalingChange: %d", new_state);
 
   OffloadFromSignalingThread([this, new_state]() {
-    if (new_state == webrtc::PeerConnectionInterface::SignalingState::kStable)
+    if (new_state == webrtc::PeerConnectionInterface::SignalingState::kStable &&
+        connection_flow_) {
       connection_flow_->OnSignalingStable();
+    }
   });
 }
 
@@ -33,8 +39,10 @@ void PeerConnectionObserverImpl::OnDataChannel(
     rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
   NEARBY_LOG(INFO, "OnDataChannel");
 
-  data_channel->RegisterObserver(
-      connection_flow_->CreateDataChannelObserver(data_channel));
+  if (connection_flow_) {
+    data_channel->RegisterObserver(
+        connection_flow_->CreateDataChannelObserver(data_channel));
+  }
 }
 
 void PeerConnectionObserverImpl::OnIceGatheringChange(
@@ -47,7 +55,9 @@ void PeerConnectionObserverImpl::OnConnectionChange(
   NEARBY_LOG(INFO, "OnConnectionChange: %d", new_state);
 
   OffloadFromSignalingThread([this, new_state]() {
-    connection_flow_->ProcessOnPeerConnectionChange(new_state);
+    if (connection_flow_) {
+      connection_flow_->ProcessOnPeerConnectionChange(new_state);
+    }
   });
 }
 
@@ -55,6 +65,11 @@ void PeerConnectionObserverImpl ::OnRenegotiationNeeded() {
   NEARBY_LOG(INFO, "OnRenegotiationNeeded");
 }
 
+void PeerConnectionObserverImpl::Shutdown() {
+  single_threaded_signaling_offloader_.Shutdown();
+  connection_flow_ = nullptr;
+}
+
 void PeerConnectionObserverImpl::OffloadFromSignalingThread(Runnable runnable) {
   single_threaded_signaling_offloader_.Execute(std::move(runnable));
 }
diff --git a/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.h b/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.h
index a8aee854..a88ed15e 100644
--- a/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.h
+++ b/cpp/core/internal/mediums/webrtc/peer_connection_observer_impl.h
@@ -14,10 +14,10 @@ class ConnectionFlow;
 
 class PeerConnectionObserverImpl : public webrtc::PeerConnectionObserver {
  public:
-  ~PeerConnectionObserverImpl() override = default;
   PeerConnectionObserverImpl(
       ConnectionFlow* connection_flow,
       LocalIceCandidateListener local_ice_candidate_listener);
+  ~PeerConnectionObserverImpl() override;
 
   // webrtc::PeerConnectionObserver:
   void OnIceCandidate(const webrtc::IceCandidateInterface* candidate) override;
@@ -31,10 +31,12 @@ class PeerConnectionObserverImpl : public webrtc::PeerConnectionObserver {
       webrtc::PeerConnectionInterface::PeerConnectionState new_state) override;
   void OnRenegotiationNeeded() override;
 
+  void Shutdown();
+
  private:
   void OffloadFromSignalingThread(Runnable runnable);
 
-  ConnectionFlow* connection_flow_;
+  ConnectionFlow* volatile connection_flow_;
   LocalIceCandidateListener local_ice_candidate_listener_;
   SingleThreadExecutor single_threaded_signaling_offloader_;
 };