pw_rpc: Guard the on_client_stream_end callback

Bug: b/234876851
Change-Id: Ia6588afe588923d64c8c4a27be4e761bed8a111b
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/126050
Pigweed-Auto-Submit: Wyatt Hepler <hepler@google.com>
Reviewed-by: Alexei Frolov <frolv@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>

1 files changed, 10 insertions, 7 deletions

diff --git a/pw_rpc/public/pw_rpc/internal/server_call.h b/pw_rpc/public/pw_rpc/internal/server_call.h
index c80e06334..2f998dc0e 100644
--- a/pw_rpc/public/pw_rpc/internal/server_call.h
+++ b/pw_rpc/public/pw_rpc/internal/server_call.h
@@ -25,13 +25,15 @@ class ServerCall : public Call {
  public:
   void HandleClientStreamEnd() PW_UNLOCK_FUNCTION(rpc_lock()) {
     MarkClientStreamCompleted();
-    // TODO(b/234876851): Ensure on_client_stream_end_ is properly guarded.
-    rpc_lock().unlock();
 
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
-    if (on_client_stream_end_) {
-      on_client_stream_end_();
+    auto on_client_stream_end_local = std::move(on_client_stream_end_);
+    rpc_lock().unlock();
+    if (on_client_stream_end_local) {
+      on_client_stream_end_local();
     }
+#else
+    rpc_lock().unlock();
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
   }
 
@@ -63,14 +65,15 @@ class ServerCall : public Call {
   // disabled with a helpful static_assert message.
   template <typename UnusedType = void>
   void set_on_client_stream_end(
-      [[maybe_unused]] Function<void()>&& on_client_stream_end) {
-    // TODO(b/234876851): Ensure on_client_stream_end_ is properly guarded.
+      [[maybe_unused]] Function<void()>&& on_client_stream_end)
+      PW_LOCKS_EXCLUDED(rpc_lock()) {
     static_assert(
         cfg::kClientStreamEndCallbackEnabled<UnusedType>,
         "The client stream end callback is disabled, so "
         "set_on_client_stream_end cannot be called. To enable the client end "
         "callback, set PW_RPC_CLIENT_STREAM_END_CALLBACK to 1.");
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
+    LockGuard lock(rpc_lock());
     on_client_stream_end_ = std::move(on_client_stream_end);
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
   }
@@ -78,7 +81,7 @@ class ServerCall : public Call {
  private:
 #if PW_RPC_CLIENT_STREAM_END_CALLBACK
   // Called when a client stream completes.
-  Function<void()> on_client_stream_end_;
+  Function<void()> on_client_stream_end_ PW_GUARDED_BY(rpc_lock());
 #endif  // PW_RPC_CLIENT_STREAM_END_CALLBACK
 };