From 5954e8e49c5a0622115ea03e0ddacab3c870aaf7 Mon Sep 17 00:00:00 2001
From: Wiktor Garbacz <wiktorg@google.com>
Date: Mon, 12 Aug 2024 08:02:02 -0700
Subject: [PATCH] Better error message on mismatching Send/Recv

Report tag mismatch before reporting length mismatch.

PiperOrigin-RevId: 662084983
Change-Id: Ie1c8fdca5488a2ddbbcc2ba2100b42de65959843
---
 sandboxed_api/rpcchannel.cc     |  7 +------
 sandboxed_api/sandbox2/comms.cc | 14 ++++++++------
 sandboxed_api/sandbox2/comms.h  |  4 +++-
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/sandboxed_api/rpcchannel.cc b/sandboxed_api/rpcchannel.cc
index 7dd79a7e..b9f04007 100644
--- a/sandboxed_api/rpcchannel.cc
+++ b/sandboxed_api/rpcchannel.cc
@@ -44,14 +44,9 @@ absl::StatusOr<FuncRet> RPCChannel::Return(v::Type exp_type) {
   uint32_t tag;
   size_t len;
   FuncRet ret;
-  if (!comms_->RecvTLV(&tag, &len, &ret, sizeof(ret))) {
+  if (!comms_->RecvTLV(&tag, &len, &ret, sizeof(ret), comms::kMsgReturn)) {
     return absl::UnavailableError("Receiving TLV value failed");
   }
-  if (tag != comms::kMsgReturn) {
-    LOG(ERROR) << "tag != comms::kMsgReturn (" << absl::StrCat(absl::Hex(tag))
-               << " != " << absl::StrCat(absl::Hex(comms::kMsgReturn)) << ")";
-    return absl::UnavailableError("Received TLV has incorrect tag");
-  }
   if (len != sizeof(FuncRet)) {
     LOG(ERROR) << "len != sizeof(FuncReturn) (" << len
                << " != " << sizeof(FuncRet) << ")";
diff --git a/sandboxed_api/sandbox2/comms.cc b/sandboxed_api/sandbox2/comms.cc
index 08b85431..05703cd8 100644
--- a/sandboxed_api/sandbox2/comms.cc
+++ b/sandboxed_api/sandbox2/comms.cc
@@ -34,6 +34,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -618,11 +619,16 @@ bool Comms::RecvTLVGeneric(uint32_t* tag, T* value) {
 }
 
 bool Comms::RecvTLV(uint32_t* tag, size_t* length, void* buffer,
-                    size_t buffer_size) {
+                    size_t buffer_size, std::optional<uint32_t> expected_tag) {
   if (!RecvTL(tag, length)) {
     return false;
   }
 
+  if (expected_tag.has_value() && *tag != *expected_tag) {
+    SAPI_RAW_LOG(ERROR, "Expected tag: 0x%08x, got: 0x%x", *expected_tag, *tag);
+    return false;
+  }
+
   if (*length == 0) {
     return true;
   }
@@ -639,14 +645,10 @@ bool Comms::RecvTLV(uint32_t* tag, size_t* length, void* buffer,
 bool Comms::RecvInt(void* buffer, size_t len, uint32_t tag) {
   uint32_t received_tag;
   size_t received_length;
-  if (!RecvTLV(&received_tag, &received_length, buffer, len)) {
+  if (!RecvTLV(&received_tag, &received_length, buffer, len, tag)) {
     return false;
   }
 
-  if (received_tag != tag) {
-    SAPI_RAW_LOG(ERROR, "Expected tag: 0x%08x, got: 0x%x", tag, received_tag);
-    return false;
-  }
   if (received_length != len) {
     SAPI_RAW_LOG(ERROR, "Expected length: %zu, got: %zu", len, received_length);
     return false;
diff --git a/sandboxed_api/sandbox2/comms.h b/sandboxed_api/sandbox2/comms.h
index 96329ba5..b9337370 100644
--- a/sandboxed_api/sandbox2/comms.h
+++ b/sandboxed_api/sandbox2/comms.h
@@ -30,6 +30,7 @@
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <variant>
@@ -144,7 +145,8 @@ class Comms {
   // by std::string.
   bool RecvTLV(uint32_t* tag, std::string* value);
   // Receives a TLV value into a specified buffer without allocating memory.
-  bool RecvTLV(uint32_t* tag, size_t* length, void* buffer, size_t buffer_size);
+  bool RecvTLV(uint32_t* tag, size_t* length, void* buffer, size_t buffer_size,
+               std::optional<uint32_t> expected_tag = std::nullopt);
 
   // Sends/receives various types of data.
   bool RecvUint8(uint8_t* v) { return RecvIntGeneric(v, kTagUint8); }