pw_rpc: Function for reading channel ID from a packet

ExtractChannelId(ConstByteSpan) parses just the channel ID from an
encoded pw_rpc packet.

Fixed: b/216849468
Change-Id: I9fed3c9b11c139bb7f4da5b984923ac8b489f3fb
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/89640
Reviewed-by: Alexei Frolov <frolv@google.com>
Commit-Queue: Wyatt Hepler <hepler@google.com>
Pigweed-Auto-Submit: Wyatt Hepler <hepler@google.com>

3 files changed, 54 insertions, 5 deletions

diff --git a/pw_rpc/channel.cc b/pw_rpc/channel.cc
index 2167cbc96..ea74168cc 100644
--- a/pw_rpc/channel.cc
+++ b/pw_rpc/channel.cc
@@ -18,20 +18,42 @@
 #include "pw_rpc/internal/channel.h"
 // clang-format on
 
+#include "pw_bytes/span.h"
 #include "pw_log/log.h"
+#include "pw_protobuf/decoder.h"
 #include "pw_rpc/internal/config.h"
 
-namespace pw::rpc::internal {
-
+namespace pw::rpc {
 namespace {
 
 // TODO(pwbug/615): Dynamically allocate this buffer if
 //     PW_RPC_DYNAMIC_ALLOCATION is enabled.
 std::array<std::byte, cfg::kEncodingBufferSizeBytes> encoding_buffer
-    PW_GUARDED_BY(rpc_lock());
+    PW_GUARDED_BY(internal::rpc_lock());
 
 }  // namespace
 
+Result<uint32_t> ExtractChannelId(ConstByteSpan packet) {
+  protobuf::Decoder decoder(packet);
+
+  while (decoder.Next().ok()) {
+    switch (static_cast<internal::RpcPacket::Fields>(decoder.FieldNumber())) {
+      case internal::RpcPacket::Fields::CHANNEL_ID: {
+        uint32_t channel_id;
+        PW_TRY(decoder.ReadUint32(&channel_id));
+        return channel_id;
+      }
+
+      default:
+        continue;
+    }
+  }
+
+  return Status::DataLoss();
+}
+
+namespace internal {
+
 ByteSpan GetPayloadBuffer() PW_EXCLUSIVE_LOCKS_REQUIRED(rpc_lock()) {
   return ByteSpan(encoding_buffer)
       .subspan(Packet::kMinEncodedSizeWithoutPayload);
@@ -61,4 +83,5 @@ Status Channel::Send(const Packet& packet) {
   return OkStatus();
 }
 
-}  // namespace pw::rpc::internal
+}  // namespace internal
+}  // namespace pw::rpc
diff --git a/pw_rpc/channel_test.cc b/pw_rpc/channel_test.cc
index 185d257a6..b61fd8d75 100644
--- a/pw_rpc/channel_test.cc
+++ b/pw_rpc/channel_test.cc
@@ -14,6 +14,8 @@
 
 #include "pw_rpc/channel.h"
 
+#include <cstddef>
+
 #include "gtest/gtest.h"
 #include "pw_rpc/internal/packet.h"
 #include "pw_rpc/internal/test_utils.h"
@@ -33,7 +35,7 @@ TEST(ChannelOutput, Name) {
 }
 
 constexpr Packet kTestPacket(
-    PacketType::RESPONSE, 1, 42, 100, 0, {}, Status::NotFound());
+    PacketType::RESPONSE, 23, 42, 100, 0, {}, Status::NotFound());
 const size_t kReservedSize = 2 /* type */ + 2 /* channel */ + 5 /* service */ +
                              5 /* method */ + 2 /* payload key */ +
                              2 /* status (if not OK) */;
@@ -54,5 +56,23 @@ TEST(Channel, TestPacket_ReservedSizeMatchesMinEncodedSizeBytes) {
   EXPECT_EQ(kReservedSize, kTestPacket.MinEncodedSizeBytes());
 }
 
+TEST(ExtractChannelId, ValidPacket) {
+  std::byte buffer[64] = {};
+  Result<ConstByteSpan> result = kTestPacket.Encode(buffer);
+  ASSERT_EQ(result.status(), OkStatus());
+
+  Result<uint32_t> channel_id = ExtractChannelId(*result);
+  ASSERT_EQ(channel_id.status(), OkStatus());
+  EXPECT_EQ(*channel_id, 23u);
+}
+
+TEST(ExtractChannelId, InvalidPacket) {
+  constexpr std::byte buffer[64] = {std::byte{1}, std::byte{2}};
+
+  Result<uint32_t> channel_id = ExtractChannelId(buffer);
+
+  EXPECT_EQ(channel_id.status(), Status::DataLoss());
+}
+
 }  // namespace
 }  // namespace pw::rpc::internal
diff --git a/pw_rpc/public/pw_rpc/channel.h b/pw_rpc/public/pw_rpc/channel.h
index 099a994b2..849595602 100644
--- a/pw_rpc/public/pw_rpc/channel.h
+++ b/pw_rpc/public/pw_rpc/channel.h
@@ -19,11 +19,17 @@
 #include <type_traits>
 
 #include "pw_assert/assert.h"
+#include "pw_bytes/span.h"
+#include "pw_result/result.h"
 #include "pw_rpc/internal/lock.h"
 #include "pw_status/status.h"
 
 namespace pw::rpc {
 
+// Extracts the channel ID from a pw_rpc packet. Returns DATA_LOSS if the
+// packet is corrupt and the channel ID could not be found.
+Result<uint32_t> ExtractChannelId(ConstByteSpan packet);
+
 class ChannelOutput {
  public:
   // Returned from MaximumTransmissionUnit() to indicate that this ChannelOutput