Miscellaneous fixes (#4584)

* Fix use after move in LevelDbPersistenceForTesting

* Fix size-related errors when compiling in 32-bit mode

* Fix test-only memory leaks

* Fix test-only memory leaks

* Prevent spec tests from appearing to leak

The spec tests aren't actually leaking, but they accumulate enough
garbage without clearing the autorelease pool that Instruments
complains.

* Avoid temporaries in methods that @throw

ARC does not release temporaries when an Objective-C exception is
thrown. Use C++ equivalents directly to avoid leaks as detected by
Instruments.

* Fix memory leak in acknowledge mutation

Any previous value must be freed before the next one is assigned.

* Manually implement assignment operators in ByteString

The single by-value implementation wasn't preventing the compiler
generated standard versions from being generated which was causing
leaks.

* Prevent leaking GrpcCompletions that complete after Shutdown

After shutdown, the AsyncQueue silently drops Enqueue requests which
prevents these blocks from being deleted.

* Avoid racing with enqueued operations

It's possible for the Executor to immediately execute the given block,
at which point the block might observe the async queue in a
not-shutting-down state.

Author: wilhuff

Firestore/Example/Tests/Integration/FSTTransactionTests.mm

@@ -233,8 +233,8 @@ - (void)runSuccessfulTransaction {
       completion:^(id _Nullable result, NSError *_Nullable error) {
         [expectation fulfill];
         NSString *message =
-            [NSString stringWithFormat:@"Expected the sequence %@, to succeed, but got %ld.",
-                                       [self stageNames], [error code]];
+            [NSString stringWithFormat:@"Expected the sequence %@, to succeed, but got %d.",
+                                       [self stageNames], (int)[error code]];
         [self->_testCase assertNilError:error message:message];
       }];
 

Firestore/Example/Tests/SpecTests/FSTSpecTests.mm

@@ -189,6 +189,9 @@ - (void)setUpForSpecWithConfig:(NSDictionary *)config {
 
 - (void)tearDownForSpec {
   [self.driver shutdown];
+
+  // Help ARC realize that everything here can be collected earlier.
+  _driver = nil;
 }
 
 /**
@@ -749,24 +752,26 @@ - (void)validateActiveTargets {
 }
 
 - (void)runSpecTestSteps:(NSArray *)steps config:(NSDictionary *)config {
-  @try {
-    [self setUpForSpecWithConfig:config];
-    for (NSDictionary *step in steps) {
-      LOG_DEBUG("Doing step %s", step);
-      [self doStep:step];
-      [self validateExpectedSnapshotEvents:step[@"expectedSnapshotEvents"]];
-      [self validateExpectedState:step[@"expectedState"]];
-      int expectedSnapshotsInSyncEvents = [step[@"expectedSnapshotsInSyncEvents"] intValue];
-      [self validateSnapshotsInSyncEvents:expectedSnapshotsInSyncEvents];
+  @autoreleasepool {
+    @try {
+      [self setUpForSpecWithConfig:config];
+      for (NSDictionary *step in steps) {
+        LOG_DEBUG("Doing step %s", step);
+        [self doStep:step];
+        [self validateExpectedSnapshotEvents:step[@"expectedSnapshotEvents"]];
+        [self validateExpectedState:step[@"expectedState"]];
+        int expectedSnapshotsInSyncEvents = [step[@"expectedSnapshotsInSyncEvents"] intValue];
+        [self validateSnapshotsInSyncEvents:expectedSnapshotsInSyncEvents];
+      }
+      [self.driver validateUsage];
+    } @finally {
+      // Ensure that the driver is torn down even if the test is failing due to a thrown exception
+      // so that any resources held by the driver are released. This is important when the driver is
+      // backed by LevelDB because LevelDB locks its database. If -tearDownForSpec were not called
+      // after an exception then subsequent attempts to open the LevelDB will fail, making it harder
+      // to zero in on the spec tests as a culprit.
+      [self tearDownForSpec];
     }
-    [self.driver validateUsage];
-  } @finally {
-    // Ensure that the driver is torn down even if the test is failing due to a thrown exception so
-    // that any resources held by the driver are released. This is important when the driver is
-    // backed by LevelDB because LevelDB locks its database. If -tearDownForSpec were not called
-    // after an exception then subsequent attempts to open the LevelDB will fail, making it harder
-    // to zero in on the spec tests as a culprit.
-    [self tearDownForSpec];
   }
 }
 

Firestore/Source/API/FIRDocumentSnapshot.mm

@@ -53,11 +53,13 @@
 using firebase::firestore::model::DatabaseId;
 using firebase::firestore::model::Document;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::FieldPath;
 using firebase::firestore::model::FieldValue;
 using firebase::firestore::model::FieldValueOptions;
 using firebase::firestore::model::ObjectValue;
 using firebase::firestore::model::ServerTimestampBehavior;
 using firebase::firestore::nanopb::MakeNSData;
+using firebase::firestore::util::MakeString;
 using firebase::firestore::util::ThrowInvalidArgument;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -180,16 +182,16 @@ - (nullable id)valueForField:(id)field {
 
 - (nullable id)valueForField:(id)field
      serverTimestampBehavior:(FIRServerTimestampBehavior)serverTimestampBehavior {
-  FIRFieldPath *fieldPath;
+  FieldPath fieldPath;
   if ([field isKindOfClass:[NSString class]]) {
-    fieldPath = [FIRFieldPath pathWithDotSeparatedString:field];
+    fieldPath = FieldPath::FromDotSeparatedString(MakeString(field));
   } else if ([field isKindOfClass:[FIRFieldPath class]]) {
-    fieldPath = field;
+    fieldPath = ((FIRFieldPath *)field).internalValue;
   } else {
     ThrowInvalidArgument("Subscript key must be an NSString or FIRFieldPath.");
   }
 
-  absl::optional<FieldValue> fieldValue = _snapshot.GetValue(fieldPath.internalValue);
+  absl::optional<FieldValue> fieldValue = _snapshot.GetValue(fieldPath);
   FieldValueOptions options = [self optionsForServerTimestampBehavior:serverTimestampBehavior];
   return !fieldValue ? nil : [self convertedValue:*fieldValue options:options];
 }

Firestore/Source/API/FIRQuery.mm

@@ -78,6 +78,7 @@
 using firebase::firestore::model::FieldValue;
 using firebase::firestore::model::ResourcePath;
 using firebase::firestore::util::MakeNSError;
+using firebase::firestore::util::MakeString;
 using firebase::firestore::util::StatusOr;
 using firebase::firestore::util::ThrowInvalidArgument;
 
@@ -86,7 +87,7 @@
 namespace {
 
 FieldPath MakeFieldPath(NSString *field) {
-  return FieldPath::FromDotSeparatedString(util::MakeString(field));
+  return FieldPath::FromDotSeparatedString(MakeString(field));
 }
 
 FIRQuery *Wrap(Query &&query) {
@@ -356,21 +357,25 @@ - (FIRQuery *)queryFilteredUsingPredicate:(NSPredicate *)predicate {
 }
 
 - (FIRQuery *)queryOrderedByField:(NSString *)field {
-  return [self queryOrderedByFieldPath:[FIRFieldPath pathWithDotSeparatedString:field]
-                            descending:NO];
+  return [self queryOrderedByField:field descending:NO];
 }
 
 - (FIRQuery *)queryOrderedByFieldPath:(FIRFieldPath *)fieldPath {
   return [self queryOrderedByFieldPath:fieldPath descending:NO];
 }
 
 - (FIRQuery *)queryOrderedByField:(NSString *)field descending:(BOOL)descending {
-  return [self queryOrderedByFieldPath:[FIRFieldPath pathWithDotSeparatedString:field]
-                            descending:descending];
+  return [self queryOrderedByFieldPath:MakeFieldPath(field)
+                             direction:Direction::FromDescending(descending)];
 }
 
 - (FIRQuery *)queryOrderedByFieldPath:(FIRFieldPath *)fieldPath descending:(BOOL)descending {
-  return Wrap(_query.OrderBy(fieldPath.internalValue, Direction::FromDescending(descending)));
+  return [self queryOrderedByFieldPath:fieldPath.internalValue
+                             direction:Direction::FromDescending(descending)];
+}
+
+- (FIRQuery *)queryOrderedByFieldPath:(model::FieldPath)fieldPath direction:(Direction)direction {
+  return Wrap(_query.OrderBy(std::move(fieldPath), direction));
 }
 
 - (FIRQuery *)queryLimitedTo:(NSInteger)limit {
@@ -468,7 +473,7 @@ - (FIRQuery *)queryWithFilterOperator:(Filter::Operator)filterOperator
                                 value:(id)value {
   FieldValue fieldValue = [self parsedQueryValue:value
                                      allowArrays:filterOperator == Filter::Operator::In];
-  auto describer = [value] { return util::MakeString(NSStringFromClass([value class])); };
+  auto describer = [value] { return MakeString(NSStringFromClass([value class])); };
   return Wrap(_query.Filter(fieldPath, filterOperator, std::move(fieldValue), describer));
 }
 

Firestore/Source/API/FSTUserDataConverter.mm

@@ -200,7 +200,7 @@ - (ParsedUpdateData)parsedUpdateData:(id)input {
     FieldPath path;
 
     if ([key isKindOfClass:[NSString class]]) {
-      path = [FIRFieldPath pathWithDotSeparatedString:key].internalValue;
+      path = FieldPath::FromDotSeparatedString(util::MakeString(key));
     } else if ([key isKindOfClass:[FIRFieldPath class]]) {
       path = ((FIRFieldPath *)key).internalValue;
     } else {

Firestore/Swift/Tests/Codable/FirestoreEncoderTests.swift

@@ -229,7 +229,7 @@ class FirestoreEncoderTests: XCTestCase {
       "s": "abc",
       "d": 123,
       "f": -4,
-      "l": 1_234_567_890_123,
+      "l": Int64(1_234_567_890_123),
       "i": -4444,
       "b": false,
       "sh": 123,

Firestore/core/src/firebase/firestore/local/leveldb_mutation_queue.cc

@@ -421,6 +421,8 @@ ByteString LevelDbMutationQueue::GetLastStreamToken() {
 }
 
 void LevelDbMutationQueue::SetLastStreamToken(ByteString stream_token) {
+  std::free(metadata_->last_stream_token);
+
   metadata_->last_stream_token = stream_token.release();
   db_->current_transaction()->Put(mutation_queue_key(), metadata_);
 }

Firestore/core/src/firebase/firestore/nanopb/byte_string.cc

@@ -55,13 +55,31 @@ ByteString::ByteString(const ByteString& other)
 }
 
 ByteString::ByteString(ByteString&& other) noexcept {
-  swap(*this, other);
+  bytes_ = other.bytes_;
+  other.bytes_ = nullptr;
 }
 
 ByteString::~ByteString() {
   std::free(bytes_);
 }
 
+ByteString& ByteString::operator=(const ByteString& other) {
+  if (bytes_ != other.bytes_) {
+    std::free(bytes_);
+    bytes_ = MakeBytesArray(other.data(), other.size());
+  }
+  return *this;
+}
+
+ByteString& ByteString::operator=(ByteString&& other) noexcept {
+  if (bytes_ != other.bytes_) {
+    std::free(bytes_);
+    bytes_ = other.bytes_;
+    other.bytes_ = nullptr;
+  }
+  return *this;
+}
+
 /* static */ ByteString ByteString::Take(pb_bytes_array_t* bytes) {
   return ByteString{bytes, 0};
 }

Firestore/core/src/firebase/firestore/nanopb/byte_string.h

@@ -80,10 +80,8 @@ class ByteString : public util::Comparable<ByteString> {
 
   ~ByteString();
 
-  ByteString& operator=(ByteString other) {
-    swap(*this, other);
-    return *this;
-  }
+  ByteString& operator=(const ByteString& other);
+  ByteString& operator=(ByteString&& other) noexcept;
 
   friend void swap(ByteString& lhs, ByteString& rhs) noexcept;
 

Firestore/core/src/firebase/firestore/remote/grpc_completion.cc

@@ -25,6 +25,23 @@ namespace remote {
 
 using util::AsyncQueue;
 
+std::shared_ptr<GrpcCompletion> GrpcCompletion::Create(
+    Type type,
+    const std::shared_ptr<util::AsyncQueue>& worker_queue,
+    Callback&& callback) {
+  // Construct in two steps to use the private constructor.
+  GrpcCompletion partial(type, worker_queue, std::move(callback));
+  auto completion = std::make_shared<GrpcCompletion>(std::move(partial));
+
+  // Prepare the `GrpcCompletion` for submission to gRPC.
+  //
+  // Note: this is done in a separate step because `shared_from_this` cannot be
+  // called in a constructor.
+  completion->grpc_ownership_ = completion;
+
+  return completion;
+}
+
 GrpcCompletion::GrpcCompletion(
     Type type,
     const std::shared_ptr<util::AsyncQueue>& worker_queue,
@@ -35,13 +52,17 @@ GrpcCompletion::GrpcCompletion(
 void GrpcCompletion::Cancel() {
   worker_queue_->VerifyIsCurrentQueue();
   callback_ = {};
+
+  // Does not release grpc_ownership_. If gRPC still holds this completion it
+  // must remain valid to avoid a use-after-free once Complete is actually
+  // called.
 }
 
 void GrpcCompletion::WaitUntilOffQueue() {
   worker_queue_->VerifyIsCurrentQueue();
 
   EnsureValidFuture();
-  return off_queue_future_.wait();
+  off_queue_future_.wait();
 }
 
 std::future_status GrpcCompletion::WaitUntilOffQueue(
@@ -64,12 +85,21 @@ void GrpcCompletion::Complete(bool ok) {
   // objects to be valid).
   off_queue_.set_value();
 
-  worker_queue_->Enqueue([this, ok] {
-    if (callback_) {
-      callback_(ok, this);
+  // The queued operation needs to also retain this completion. It's possible
+  // for Complete to fire, shutdown to start, and then have this queued
+  // operation run. If this weren't a retain that ordering would have the
+  // callback use after free.
+  auto shared_this = grpc_ownership_;
+  worker_queue_->Enqueue([shared_this, ok] {
+    if (shared_this->callback_) {
+      shared_this->callback_(ok, shared_this);
     }
-    delete this;
   });
+
+  // Having called Complete, gRPC has released its ownership interest in this
+  // object. Once the queued operation completes the `GrpcCompletion` will be
+  // deleted.
+  grpc_ownership_.reset();
 }
 
 }  // namespace remote