pw_rpc: Guard the Nanopb and pwpb on_completed callbacks

Move the on_completed_ callbacks for Nanopb and pwpb RPCs to a local variable while the lock is held. Bug: b/234876851 Change-Id: I5a2b16cdb658925ba1985f5700821483d2deab23 Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/125230 Pigweed-Auto-Submit: Wyatt Hepler <[email protected]> Commit-Queue: Auto-Submit <[email protected]> Reviewed-by: Alexei Frolov <[email protected]>
google · Dec 23, 2022 · 580e279 · 580e279
1 parent c4a481b
commit 580e279
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/pw_rpc/nanopb/method_union_test.cc b/pw_rpc/nanopb/method_union_test.cc
@@ -25,8 +25,6 @@
 namespace pw::rpc::internal {
 namespace {
 
-using std::byte;
-
 template <typename Implementation>
 class FakeGeneratedService : public Service {
  public:

diff --git a/pw_rpc/nanopb/public/pw_rpc/nanopb/client_reader_writer.h b/pw_rpc/nanopb/public/pw_rpc/nanopb/client_reader_writer.h
@@ -103,10 +103,14 @@ class NanopbUnaryResponseClientCall : public UnaryResponseClientCall {
 
     UnaryResponseClientCall::set_on_completed_locked(
         [this](ConstByteSpan payload, Status status) {
-          if (nanopb_on_completed_) {
+          rpc_lock().lock();
+          auto nanopb_on_completed_local = std::move(nanopb_on_completed_);
+          rpc_lock().unlock();
+
+          if (nanopb_on_completed_local) {
             Response response_struct{};
             if (serde_->DecodeResponse(payload, &response_struct)) {
-              nanopb_on_completed_(response_struct, status);
+              nanopb_on_completed_local(response_struct, status);
             } else {
               // TODO(hepler): This should send a DATA_LOSS error and call the
               //     error callback.
@@ -118,7 +122,8 @@ class NanopbUnaryResponseClientCall : public UnaryResponseClientCall {
   }
 
   const NanopbMethodSerde* serde_;
-  Function<void(const Response&, Status)> nanopb_on_completed_;
+  Function<void(const Response&, Status)> nanopb_on_completed_
+      PW_GUARDED_BY(rpc_lock());
 };
 
 // Base class for server and bidirectional streaming calls.

diff --git a/pw_rpc/pwpb/public/pw_rpc/pwpb/client_reader_writer.h b/pw_rpc/pwpb/public/pw_rpc/pwpb/client_reader_writer.h
@@ -136,12 +136,16 @@ class PwpbUnaryResponseClientCall : public UnaryResponseClientCall {
 
     UnaryResponseClientCall::set_on_completed_locked(
         [this](ConstByteSpan payload, Status status) {
-          if (pwpb_on_completed_) {
+          rpc_lock().lock();
+          auto pwpb_on_completed_local = std::move(pwpb_on_completed_);
+          rpc_lock().unlock();
+
+          if (pwpb_on_completed_local) {
             Response response{};
             const Status decode_status =
                 serde().DecodeResponse(payload, response);
             if (decode_status.ok()) {
-              pwpb_on_completed_(response, status);
+              pwpb_on_completed_local(response, status);
             } else {
               rpc_lock().lock();
               CallOnError(Status::DataLoss());
@@ -151,7 +155,8 @@ class PwpbUnaryResponseClientCall : public UnaryResponseClientCall {
   }
 
   const PwpbMethodSerde* serde_;
-  Function<void(const Response&, Status)> pwpb_on_completed_;
+  Function<void(const Response&, Status)> pwpb_on_completed_
+      PW_GUARDED_BY(rpc_lock());
 };
 
 // internal::PwpbStreamResponseClientCall extends