use after free fix for perf-client on shutdown (#218)

Summary:
The perf client could crash on shutdown. SubscriberState handed the session a non-owning
pointer to its own callback member, but the session keeps subgroup receivers alive after
SubscriberState is destroyed, so a late callback dereferenced freed memory.

The callback is now a shared_ptr that outlives SubscriberState, and the destructor detaches
it so any callback arriving afterwards returns early instead of touching the dead object.

Pull Request resolved: #218

Reviewed By: sandarsh

Differential Revision: D116952494

Pulled By: afrind

fbshipit-source-id: d03e01fc365b35271db7bcb4af331ed1b165cc9b

Author: peterchave

moxygen/moqtest/MoQPerfTestClient.cpp

@@ -58,11 +58,12 @@ SubscriberState::SubscriberState(
       receiver_(
           std::make_shared<ObjectReceiver>(
               ObjectReceiver::SUBSCRIBE,
-              std::shared_ptr<ObjectReceiverCallback>(
-                  std::shared_ptr<void>(),
-                  &callback_))) {}
+              callback_)) {}
 
 SubscriberState::~SubscriberState() {
+  // The session may still hold subgroup receivers referencing callback_;
+  // detach first so any late callback can't touch this destroyed object.
+  callback_->detach();
   try {
     // Unsubscribe if we have a valid subHandle_
     if (subHandle_) {
@@ -193,9 +194,12 @@ ObjectReceiverCallback::FlowControlState SubscriberState::Callback::onObject(
     std::optional<TrackAlias> /* trackAlias */,
     const ObjectHeader& objHeader,
     Payload payload) {
-  state_.objectsReceived_++;
+  if (!state_) {
+    return FlowControlState::UNBLOCKED;
+  }
+  state_->objectsReceived_++;
   if (payload) {
-    state_.bytesReceived_ += payload->computeChainDataLength();
+    state_->bytesReceived_ += payload->computeChainDataLength();
   }
 
   if (auto sendTs =
@@ -205,15 +209,15 @@ ObjectReceiverCallback::FlowControlState SubscriberState::Callback::onObject(
                          .count();
     if (nowMs >= *sendTs) {
       uint64_t latencyMs = nowMs - *sendTs;
-      state_.totalLatencyMs_ += latencyMs;
-      state_.latencyObjects_++;
-      state_.testClient_.recordLatency(latencyMs);
+      state_->totalLatencyMs_ += latencyMs;
+      state_->latencyObjects_++;
+      state_->testClient_.recordLatency(latencyMs);
     }
   }
 
   // Update largest object seen for track restart detection
   AbsoluteLocation location(objHeader.group, objHeader.id);
-  state_.testClient_.updateLargestObjectSeen(location);
+  state_->testClient_.updateLargestObjectSeen(location);
 
   return FlowControlState::UNBLOCKED;
 }
@@ -229,37 +233,46 @@ void SubscriberState::Callback::onEndOfStream() {
 }
 
 void SubscriberState::Callback::onError(ResetStreamErrorCode code) {
-  XLOG(ERR) << "Subscriber " << state_.id_
+  if (!state_) {
+    return;
+  }
+  XLOG(ERR) << "Subscriber " << state_->id_
             << " received stream reset: " << static_cast<uint64_t>(code);
   // Don't unsubscribe immediately - let removeSubscriber handle it
-  state_.testClient_.recordReset();
+  state_->testClient_.recordReset();
 }
 
 void SubscriberState::Callback::onPublishDone(PublishDone done) {
-  XLOG(DBG1) << "Subscriber " << state_.id_
+  if (!state_) {
+    return;
+  }
+  XLOG(DBG1) << "Subscriber " << state_->id_
              << " received PublishDone - status: "
              << static_cast<uint32_t>(done.statusCode)
              << ", reason: " << done.reasonPhrase;
 
   // Library has already cleaned up subscription state; just release our handle
-  state_.subHandle_.reset();
+  state_->subHandle_.reset();
 
   // Only signal completion if track ended naturally
   if (done.statusCode == PublishDoneStatusCode::TRACK_ENDED) {
-    XLOG(DBG1) << "Subscriber " << state_.id_ << " - track ended naturally";
-    state_.testClient_.completed();
+    XLOG(DBG1) << "Subscriber " << state_->id_ << " - track ended naturally";
+    state_->testClient_.completed();
   } else {
     // Other status codes (errors, going away, etc.) don't end the test
-    XLOG(DBG1) << "Subscriber " << state_.id_
+    XLOG(DBG1) << "Subscriber " << state_->id_
                << " - PublishDone with non-TRACK_ENDED status, not ending test";
   }
 }
 
 void SubscriberState::Callback::onAllDataReceived() {
-  XLOG(DBG1) << "Subscriber " << state_.id_
+  if (!state_) {
+    return;
+  }
+  XLOG(DBG1) << "Subscriber " << state_->id_
              << " - all data received, removing subscriber";
   // Remove this subscriber from the client's map now that all streams are done
-  state_.testClient_.removeSubscriber(state_.id_);
+  state_->testClient_.removeSubscriber(state_->id_);
 }
 
 // ============================================================================

moxygen/moqtest/MoQPerfTestClient.h

@@ -59,10 +59,17 @@ class SubscriberState {
   bool hasError_{false};
 
  private:
-  // ObjectReceiverCallback implementation
+  // ObjectReceiverCallback implementation.
+  // The MoQ session keeps subgroup receivers (and therefore this callback)
+  // alive past SubscriberState destruction, so detach() severs the back
+  // pointer and later callbacks become no-ops.
   class Callback : public ObjectReceiverCallback {
    public:
-    explicit Callback(SubscriberState& state) : state_(state) {}
+    explicit Callback(SubscriberState& state) : state_(&state) {}
+
+    void detach() {
+      state_ = nullptr;
+    }
 
     FlowControlState onObject(
         std::optional<TrackAlias> trackAlias,
@@ -79,10 +86,10 @@ class SubscriberState {
     void onAllDataReceived() override;
 
    private:
-    SubscriberState& state_;
+    SubscriberState* state_;
   };
 
-  Callback callback_{*this};
+  std::shared_ptr<Callback> callback_{std::make_shared<Callback>(*this)};
   std::shared_ptr<MoQFollyExecutorImpl> moqExecutor_;
   std::unique_ptr<MoQClientBase> moqClient_;
   std::shared_ptr<ObjectReceiver> receiver_;