From b4eedaaacedfc146ef7899694e56e852e0e174af Mon Sep 17 00:00:00 2001
From: Armando Montanez <amontanez@google.com>
Date: Fri, 22 Apr 2022 10:53:47 -0700
Subject: [PATCH] pw_transfer: Cancel unregistered transfers

When a transfer handler is unregistered, all active transfers that rely
on it are canceled to ensure no dangling references to the handler
remain after unregistration.

Change-Id: I3d6da1416b8e342090bc955bf912058616454b6f
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/92101
Pigweed-Auto-Submit: Armando Montanez <amontanez@google.com>
Reviewed-by: Wyatt Hepler <hepler@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
---
 .../public/pw_transfer/internal/server_context.h    |  3 +++
 pw_transfer/public/pw_transfer/transfer_thread.h    |  3 +++
 pw_transfer/transfer_thread.cc                      | 13 +++++++++++++
 3 files changed, 19 insertions(+)

diff --git a/pw_transfer/public/pw_transfer/internal/server_context.h b/pw_transfer/public/pw_transfer/internal/server_context.h
index 038feece66..ce9a78bbbf 100644
--- a/pw_transfer/public/pw_transfer/internal/server_context.h
+++ b/pw_transfer/public/pw_transfer/internal/server_context.h
@@ -32,6 +32,9 @@ class ServerContext final : public Context {
   // ClientContexts don't track it.
   void set_handler(Handler& handler) { handler_ = &handler; }
 
+  // Returns the pointer to the current handler.
+  const Handler* handler() { return handler_; }
+
  private:
   // Ends the transfer with the given status, calling the handler's Finalize
   // method. No chunks are sent.
diff --git a/pw_transfer/public/pw_transfer/transfer_thread.h b/pw_transfer/public/pw_transfer/transfer_thread.h
index 8b1bcc019c..51c71b7bba 100644
--- a/pw_transfer/public/pw_transfer/transfer_thread.h
+++ b/pw_transfer/public/pw_transfer/transfer_thread.h
@@ -123,6 +123,9 @@ class TransferThread : public thread::ThreadCore {
 
   void RemoveTransferHandler(Handler& handler) {
     TransferHandlerEvent(EventType::kRemoveTransferHandler, handler);
+    // Ensure this function blocks until the transfer handler is fully cleaned
+    // up.
+    WaitUntilEventIsProcessed();
   }
 
   size_t max_chunk_size() const { return chunk_buffer_.size(); }
diff --git a/pw_transfer/transfer_thread.cc b/pw_transfer/transfer_thread.cc
index 2bbcf384cb..f6431ea508 100644
--- a/pw_transfer/transfer_thread.cc
+++ b/pw_transfer/transfer_thread.cc
@@ -270,6 +270,19 @@ void TransferThread::HandleEvent(const internal::Event& event) {
       return;
 
     case EventType::kRemoveTransferHandler:
+      for (ServerContext& server_context : server_transfers_) {
+        if (server_context.handler() == event.remove_transfer_handler) {
+          server_context.HandleEvent({
+              .type = EventType::kServerEndTransfer,
+              .end_transfer =
+                  {
+                      .session_id = server_context.session_id(),
+                      .status = Status::Aborted().code(),
+                      .send_status_chunk = false,
+                  },
+          });
+        }
+      }
       handlers_.remove(*event.remove_transfer_handler);
       return;