Fix use-after-free in DocumentReference::AddSnapshotListener (#2686) (#2688)

Author: mikelehen

Firestore/CHANGELOG.md

@@ -1,4 +1,8 @@
 # Unreleased
+- [fixed] Fixed a use-after-free bug that could be observed when using snapshot
+  listeners on temporary document references (#2682).
+
+# 1.2.0
 - [fixed] Fixed the way gRPC certificates are loaded on macOS (#2604).
 
 # 1.1.0

Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.mm

@@ -128,4 +128,30 @@ - (void)testCanBeRemovedIndependently {
   [two remove];
 }
 
+- (void)testCanOutliveDocumentReference {
+  FIRCollectionReference *collectionRef = [self collectionRef];
+
+  XCTestExpectation *seen = [self expectationWithDescription:@"seen document"];
+
+  __block id<FIRListenerRegistration> registration;
+  NSString *documentID;
+  @autoreleasepool {
+    FIRDocumentReference *docRef = [collectionRef documentWithAutoID];
+    documentID = docRef.documentID;
+    registration = [docRef addSnapshotListener:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+      if (snapshot.exists) {
+        [seen fulfill];
+      }
+    }];
+    docRef = nil;
+  }
+
+  XCTAssertNotNil(registration);
+
+  FIRDocumentReference *docRef2 = [collectionRef documentWithPath:documentID];
+  [self writeDocumentRef:docRef2 data:@{@"foo" : @"bar"}];
+
+  [registration remove];
+}
+
 @end

Firestore/core/src/firebase/firestore/api/document_reference.mm

@@ -202,34 +202,35 @@ void Resolve(ListenerRegistration&& registration) {
    public:
     Converter(DocumentReference* parent,
               DocumentSnapshot::Listener&& user_listener)
-        : parent_(parent), user_listener_(std::move(user_listener)) {
+        : firestore_(parent->firestore_),
+          key_(parent->key_),
+          user_listener_(std::move(user_listener)) {
     }
 
     void OnEvent(StatusOr<ViewSnapshot> maybe_snapshot) override {
       if (!maybe_snapshot.ok()) {
         user_listener_->OnEvent(maybe_snapshot.status());
         return;
       }
-      Firestore* firestore = parent_->firestore_;
-      DocumentKey key = parent_->key_;
 
       ViewSnapshot snapshot = std::move(maybe_snapshot).ValueOrDie();
       HARD_ASSERT(snapshot.documents().size() <= 1,
                   "Too many documents returned on a document query");
-      FSTDocument* document = snapshot.documents().GetDocument(key);
+      FSTDocument* document = snapshot.documents().GetDocument(key_);
 
       bool has_pending_writes =
-          document ? snapshot.mutated_keys().contains(key)
+          document ? snapshot.mutated_keys().contains(key_)
                    // We don't raise `has_pending_writes` for deleted documents.
                    : false;
 
-      DocumentSnapshot result{firestore, std::move(key), document,
-                              snapshot.from_cache(), has_pending_writes};
+      DocumentSnapshot result{firestore_, key_, document, snapshot.from_cache(),
+                              has_pending_writes};
       user_listener_->OnEvent(std::move(result));
     }
 
    private:
-    DocumentReference* parent_;
+    Firestore* firestore_;
+    DocumentKey key_;
     DocumentSnapshot::Listener user_listener_;
   };
   auto view_listener =