squashme: address review feedback

Author: ivmaykov

core/src/checks/rpc_checks.c

@@ -11,20 +11,50 @@
 #include <pb_encode.h>
 #include <squareup/subzero/internal.pb.h>
 
-static size_t get_serialized_message_size(const InternalCommandRequest* const cmd) {
+// Helper which returns the size of a buffer that would be needed to hold the serialized version of the given
+// protobuf structure, assuming that pb_encode_delimited() serialization will be used.
+static size_t get_serialized_message_size(const pb_field_t fields[], const void* const proto) {
   pb_ostream_t stream = PB_OSTREAM_SIZING;
-  if (!pb_encode_delimited(&stream, InternalCommandRequest_fields, cmd)) {
+  if (!pb_encode_delimited(&stream, fields, proto)) {
     ERROR("%s: pb_encode_delimited() failed: %s", __func__, PB_GET_ERROR(&stream));
     return 0;
   }
   return stream.bytes_written;
 }
 
+// Serializes the given protobuf structure to the given buffer of the given size, using pb_encode_delimited().
+// Returns true on success or false on failure.
+// Caller brings their own buffer memory, this function does not allocate.
+static bool
+serialize_to_buf(void* const buffer, const size_t buffer_size, const pb_field_t fields[], const void* const proto) {
+  pb_ostream_t ostream = pb_ostream_from_buffer(buffer, buffer_size);
+  if (!pb_encode_delimited(&ostream, fields, proto)) {
+    ERROR("%s: pb_encode_delimited() failed: %s", __func__, PB_GET_ERROR(&ostream));
+    return false;
+  }
+  return true;
+}
+
+// Deserializes the given buffer into the given protobuf structure, using pb_decode_delimited().
+// Returns true on success or false on failure.
+// Caller brings their own protobuf structure, this function does not allocate.
+static bool
+deserialize_from_buf(const void* const buffer, const size_t buffer_size, const pb_field_t fields[], void* const proto) {
+  pb_istream_t istream = pb_istream_from_buffer(buffer, buffer_size);
+  if (!pb_decode_delimited(&istream, fields, proto)) {
+    ERROR("%s: pb_decode_delimited() failed: %s", __func__, PB_GET_ERROR(&istream));
+    return false;
+  }
+  return true;
+}
+
 int verify_rpc_oversized_message_rejected(void) {
   int result = 0;
   uint8_t* serialized_request = NULL;
   uint8_t* serialized_response = NULL;
 
+  // Construct an initial InternalCommandRequest which holds an InitWallet command
+  // with a maximum-allowed-length random_bytes field.
   InternalCommandRequest cmd = InternalCommandRequest_init_default;
   cmd.version = VERSION;
   cmd.wallet_id = 1; // dummy value
@@ -35,12 +65,15 @@ int verify_rpc_oversized_message_rejected(void) {
   cmd.command.InitWallet.random_bytes.size = MASTER_SEED_SIZE;
   random_buffer(cmd.command.InitWallet.random_bytes.bytes, MASTER_SEED_SIZE);
 
-  size_t serialized_size = get_serialized_message_size(&cmd);
+  // Compute the size of the serialized struct.
+  size_t serialized_size = get_serialized_message_size(InternalCommandRequest_fields, &cmd);
   if (serialized_size == 0) {
+    ERROR("%s: error computing serialized request size", __func__);
     result = -1;
     goto out;
   }
 
+  // Allocate a buffer to hold the serialized struct.
   // Note that we allocate 1 extra byte because we'll be extending the message.
   serialized_request = (uint8_t*) calloc(1, serialized_size + 1);
   if (NULL == serialized_request) {
@@ -49,30 +82,30 @@ int verify_rpc_oversized_message_rejected(void) {
     goto out;
   }
 
-  pb_ostream_t ostream = pb_ostream_from_buffer(serialized_request, serialized_size);
-  if (!pb_encode_delimited(&ostream, InternalCommandRequest_fields, &cmd)) {
-    ERROR("%s: pb_encode_delimited() failed: %s", __func__, PB_GET_ERROR(&ostream));
+  // Serialize the struct to a byte array.
+  if (!serialize_to_buf(serialized_request, serialized_size, InternalCommandRequest_fields, &cmd)) {
+    ERROR("%s: serialize_to_buf() failed", __func__);
     result = -1;
     goto out;
   }
 
   // Helper macro used to check our assumptions in the gnarly protobuf mangling code below
-#define ASSERT_BYTE_EQUALS(buf, idx, expected_val)                        \
-  do {                                                                    \
-    const uint8_t* buf_ = (buf);                                          \
-    const size_t idx_ = (idx);                                            \
-    const uint8_t expected_val_ = (expected_val);                         \
-    const uint8_t actual_val_ = buf_[idx_];                               \
-    if (actual_val_ != expected_val_) {                                   \
-      ERROR(                                                              \
-          "%s: buf[%zu] contains unexpected value: %hhu, expected: %hhu", \
-          __func__,                                                       \
-          idx_,                                                           \
-          actual_val_,                                                    \
-          expected_val_);                                                 \
-      result = -1;                                                        \
-      goto out;                                                           \
-    }                                                                     \
+#define ASSERT_BYTE_EQUALS(buf, idx, expected_val)                           \
+  do {                                                                       \
+    const uint8_t* buf_ = (buf);                                             \
+    const size_t idx_ = (idx);                                               \
+    const uint8_t expected_val_ = (expected_val);                            \
+    const uint8_t actual_val_ = buf_[idx_];                                  \
+    if (actual_val_ != expected_val_) {                                      \
+      ERROR(                                                                 \
+          "%s: buf[%zu] contains an unexpected value: %hhu, expected: %hhu", \
+          __func__,                                                          \
+          idx_,                                                              \
+          actual_val_,                                                       \
+          expected_val_);                                                    \
+      result = -1;                                                           \
+      goto out;                                                              \
+    }                                                                        \
   } while (0)
 
   // Corrupt the message by making the random_bytes field 1 byte longer than the max allowed size.
@@ -87,23 +120,23 @@ int verify_rpc_oversized_message_rejected(void) {
   //                           length will actually take more than 1 byte, shifting everything after
   //                           it by a byte.
   //                           *** NOTE: WE NEED TO INCREMENT THIS BY 1. ***
-  //   serialized_request[1] - field id (1 << 3) + tag (0) for field 1 (version). Should equal 0x08.
+  //   serialized_request[1] - field id (1 << 3) + tag (0) for field 1 (version). Should equal 8.
   //   serialized_request[2..3] - varint-encoded value for field 1. Leave this alone, it's the
   //                              contents of the 'version' field (210 at the time of writing). If
   //                              version ever exceeds 16383, this will start taking up an extra byte
   //                              and shift everything after it by a byte.
-  //   serialized_request[4] - field id (2 << 3) + tag (0) for field 2 (wallet_id). Should equal 0x10.
-  //   serialized_request[5] - varint-encoded value for field 2. Leave this allone, it's the dummy
-  //                           'wallet' field which we set to 1 above. Should equal 0x01.
+  //   serialized_request[4] - field id (2 << 3) + tag (0) for field 2 (wallet_id). Should equal 16.
+  //   serialized_request[5] - varint-encoded value for field 2. Leave this alone, it's the dummy
+  //                           'wallet' field which we set to 1 above. Should equal 1.
   //   serialized_request[6] - field id (5 << 3) + tag (2, for 'LEN') for field 5 (command.InitWallet).
-  //                           Should equal 0x2a.
+  //                           Should equal 42.
   //   serialized_request[7] - varint-encoded LEN of the InitWalletRequest submessage.
-  //                           Should equal 0x42 (decimal 66).
+  //                           Should equal 66.
   //                           *** NOTE: WE NEED TO INCREMENT THIS BY 1. ***
   //   serialized_request[8] - field id (1 << 3) + tag (2, for 'LEN') for field 1 of sub-message.
-  //                           Should equal 0x0a.
+  //                           Should equal 10.
   //   serialized_request[9] - varint-encoded LEN of field 1 (random_bytes) of sub-message.
-  //                           Should equal 0x40 (decimal 64).
+  //                           Should equal 64.
   //                           *** NOTE: WE NEED TO INCREMENT THIS BY 1. ***
   //   serialized_request[10..73] - the contents of the random_bytes field. Should be 64 bytes in length.
   //   serialized_request[74] - doesn't exist in the original message. We add an extra data byte here.
@@ -130,17 +163,21 @@ int verify_rpc_oversized_message_rejected(void) {
   serialized_request[7]++;                    // increment LEN byte for top-level field 5
   serialized_request[9]++;                    // increment LEN byte for nested field 1
   serialized_request[serialized_size] = 0xaa; // set the last byte to an arbitrary value
+  serialized_size++;                          // increment serialized_size since we added a byte
 
-  pb_istream_t istream = pb_istream_from_buffer(serialized_request, serialized_size + 1);
+  // Create a stream which will read from the corrupted serialized buffer.
+  pb_istream_t istream = pb_istream_from_buffer(serialized_request, serialized_size);
+
+  // Allocate a buffer for the serialized response.
   const size_t response_buffer_size = 2048; // 2048 bytes should be more than enough
   serialized_response = (uint8_t*) calloc(1, response_buffer_size);
   if (NULL == serialized_response) {
     ERROR("%s: calloc(1, %zu) failed", __func__, response_buffer_size);
     result = -1;
     goto out;
   }
-  pb_ostream_t ostream2 = pb_ostream_from_buffer(serialized_response, response_buffer_size);
-  ERROR("(next line is expected to show red text...)");
+  // Create a stream which will write to the response buffer.
+  pb_ostream_t ostream = pb_ostream_from_buffer(serialized_response, response_buffer_size);
 
   // Now that we have a serialized buffer, try to pass it to handle_incoming_message().
   // This should fail because the InitWallet.random_bytes field has a length of 65 bytes,
@@ -149,23 +186,24 @@ int verify_rpc_oversized_message_rejected(void) {
   // NOTE: when building for nCipher, there are command hooks that would reject the command
   // because it's missing the tickets for key use authorization. But this doesn't matter for
   // this test case, because the protobuf parsing happens before that and fails first.
-  handle_incoming_message(&istream, &ostream2);
-  const size_t actual_response_size = ostream2.bytes_written;
+  ERROR("(next line is expected to show red text...)");
+  handle_incoming_message(&istream, &ostream);
+
+  // Extract the response structure from the serialized_response buffer. It should be an error.
+  const size_t actual_response_size = ostream.bytes_written;
   if (actual_response_size == 0) {
-    ERROR("%s: no response received from handle_incoming_message(): %s", __func__, PB_GET_ERROR(&ostream2));
+    ERROR("%s: no response received from handle_incoming_message(): %s", __func__, PB_GET_ERROR(&ostream));
     result = -1;
     goto out;
   }
-  pb_istream_t istream2 = pb_istream_from_buffer(serialized_response, actual_response_size);
-  InternalCommandResponse response; // note: no need to init, pb_decode_delimited() does it
-  if (!pb_decode_delimited(&istream2, InternalCommandResponse_fields, &response)) {
-    ERROR(
-        "%s: pb_decode_delimited(..., InternalCommandResponse_fields, ...) failed: %s",
-        __func__,
-        PB_GET_ERROR(&istream2));
+  InternalCommandResponse response; // note: no need to init, deserialize_from_buf() does it via pb_decode_delimited().
+  if (!deserialize_from_buf(serialized_response, actual_response_size, InternalCommandResponse_fields, &response)) {
+    ERROR("%s: deserialize_from_buf() failed", __func__);
     result = -1;
     goto out;
   }
+
+  // Check that the response contains an error.
   if (response.which_response != InternalCommandResponse_Error_tag) {
     ERROR(
         "%s: wrong response tag: %d, expected: %d",
@@ -175,6 +213,7 @@ int verify_rpc_oversized_message_rejected(void) {
     result = -1;
     goto out;
   }
+  // Check that the error response contains the expected error code.
   if (response.response.Error.code != Result_COMMAND_DECODE_FAILED) {
     ERROR(
         "%s: wrong response error code: %d, expected: %d",
@@ -184,11 +223,13 @@ int verify_rpc_oversized_message_rejected(void) {
     result = -1;
     goto out;
   }
+  // Check that the error response contains some message.
   if (!response.response.Error.has_message) {
     ERROR("%s: error response does not contain a 'message' field", __func__);
     result = -1;
     goto out;
   }
+  // Check that the error response contains the expected message.
   if (0 != strcmp("Decode Input failed: bytes overflow", response.response.Error.message)) {
     ERROR("%s: error response contains unexpected message: %s", __func__, response.response.Error.message);
     result = -1;
@@ -198,5 +239,8 @@ int verify_rpc_oversized_message_rejected(void) {
 out:
   free(serialized_request);
   free(serialized_response);
+  if (result == 0) {
+    INFO("%s: ok", __func__);
+  }
   return result;
 }