Follow up to #5889 - avoid misaligned peeks (#5905)

CMakeLists.txt

@@ -947,6 +947,7 @@ if(BUILD_TESTS)
       ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/messaging.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/oversized.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/typed_messages.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/serialized.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/serializer.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/hash.cpp
       ${CMAKE_CURRENT_SOURCE_DIR}/src/ds/test/thread_messaging.cpp

src/ds/serialized.h

@@ -30,13 +30,29 @@ namespace serialized
   template <class T>
   T peek(const uint8_t*& data, size_t& size)
   {
+    // This should only be used for numeric types and small structs
+    static constexpr auto max_size = 32u;
+    static_assert(sizeof(T) <= max_size);
+
     if (size < sizeof(T))
     {
       throw InsufficientSpaceException(
         fmt::format("Insufficient space (peek<T>: {} < {})", size, sizeof(T)));
     }
 
-    return *(T*)data;
+    static constexpr auto alignment = alignof(T);
+    if (reinterpret_cast<std::uintptr_t>(data) % alignment != 0)
+    {
+      // Data is not aligned - copy to scratch memory
+      alignas(T) uint8_t scratch[max_size];
+      std::memcpy(scratch, data, sizeof(T));
+      return *(T*)scratch;
+    }
+    else
+    {
+      // Cast directly from source memory
+      return *(T*)data;
+    }
   }
 
   template <class T>

src/ds/test/serialized.cpp

@@ -0,0 +1,86 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the Apache 2.0 License.
+#include "../serialized.h"
+
+#include <doctest/doctest.h>
+
+template <class T>
+T peek_in_vec(const std::vector<uint8_t>& v, size_t offset)
+{
+  auto data = v.data();
+  auto size = v.size();
+  REQUIRE(offset < size);
+  data += offset;
+  size -= offset;
+  return serialized::peek<T>(data, size);
+}
+
+TEST_CASE("peek unaligned" * doctest::test_suite("serialized"))
+{
+  std::vector<uint8_t> src{
+    0x01,
+    0x23,
+    0x45,
+    0x67,
+    0x89,
+    0xab,
+    0xcd,
+    0xef,
+    0xfe,
+    0xdc,
+    0xba,
+    0x98,
+    0x76,
+    0x54,
+    0x32,
+    0x10};
+
+  // Confirm that we can read at any alignment
+  {
+    INFO("uint8_t");
+    REQUIRE(peek_in_vec<uint8_t>(src, 0) == 0x01);
+    REQUIRE(peek_in_vec<uint8_t>(src, 1) == 0x23);
+    REQUIRE(peek_in_vec<uint8_t>(src, 2) == 0x45);
+    REQUIRE(peek_in_vec<uint8_t>(src, 3) == 0x67);
+    REQUIRE(peek_in_vec<uint8_t>(src, 4) == 0x89);
+    REQUIRE(peek_in_vec<uint8_t>(src, 5) == 0xab);
+    REQUIRE(peek_in_vec<uint8_t>(src, 6) == 0xcd);
+    REQUIRE(peek_in_vec<uint8_t>(src, 7) == 0xef);
+  }
+
+  {
+    INFO("uint16_t");
+    REQUIRE(peek_in_vec<uint16_t>(src, 0) == 0x23'01);
+    REQUIRE(peek_in_vec<uint16_t>(src, 1) == 0x45'23);
+    REQUIRE(peek_in_vec<uint16_t>(src, 2) == 0x67'45);
+    REQUIRE(peek_in_vec<uint16_t>(src, 3) == 0x89'67);
+    REQUIRE(peek_in_vec<uint16_t>(src, 4) == 0xab'89);
+    REQUIRE(peek_in_vec<uint16_t>(src, 5) == 0xcd'ab);
+    REQUIRE(peek_in_vec<uint16_t>(src, 6) == 0xef'cd);
+    REQUIRE(peek_in_vec<uint16_t>(src, 7) == 0xfe'ef);
+  }
+
+  {
+    INFO("uint32_t");
+    REQUIRE(peek_in_vec<uint32_t>(src, 0) == 0x67'45'23'01);
+    REQUIRE(peek_in_vec<uint32_t>(src, 1) == 0x89'67'45'23);
+    REQUIRE(peek_in_vec<uint32_t>(src, 2) == 0xab'89'67'45);
+    REQUIRE(peek_in_vec<uint32_t>(src, 3) == 0xcd'ab'89'67);
+    REQUIRE(peek_in_vec<uint32_t>(src, 4) == 0xef'cd'ab'89);
+    REQUIRE(peek_in_vec<uint32_t>(src, 5) == 0xfe'ef'cd'ab);
+    REQUIRE(peek_in_vec<uint32_t>(src, 6) == 0xdc'fe'ef'cd);
+    REQUIRE(peek_in_vec<uint32_t>(src, 7) == 0xba'dc'fe'ef);
+  }
+
+  {
+    INFO("uint64_t");
+    REQUIRE(peek_in_vec<uint64_t>(src, 0) == 0xef'cd'ab'89'67'45'23'01);
+    REQUIRE(peek_in_vec<uint64_t>(src, 1) == 0xfe'ef'cd'ab'89'67'45'23);
+    REQUIRE(peek_in_vec<uint64_t>(src, 2) == 0xdc'fe'ef'cd'ab'89'67'45);
+    REQUIRE(peek_in_vec<uint64_t>(src, 3) == 0xba'dc'fe'ef'cd'ab'89'67);
+    REQUIRE(peek_in_vec<uint64_t>(src, 4) == 0x98'ba'dc'fe'ef'cd'ab'89);
+    REQUIRE(peek_in_vec<uint64_t>(src, 5) == 0x76'98'ba'dc'fe'ef'cd'ab);
+    REQUIRE(peek_in_vec<uint64_t>(src, 6) == 0x54'76'98'ba'dc'fe'ef'cd);
+    REQUIRE(peek_in_vec<uint64_t>(src, 7) == 0x32'54'76'98'ba'dc'fe'ef);
+  }
+}

tests/connections.py

@@ -333,7 +333,13 @@ def run_node_socket_robustness_tests(args):
         try_write(encode_msg())
         try_write(encode_msg(msg_type=1))
         try_write(encode_msg(msg_type=100))
+        try_write(encode_msg(sender="a"))
+        try_write(encode_msg(sender="ab"))
+        try_write(encode_msg(sender="abc"))
         try_write(encode_msg(sender="abcd"))
+        try_write(encode_msg(sender="abcde"))
+        try_write(encode_msg(sender="abcdef"))
+        try_write(encode_msg(sender="abcdefg"))
         try_write(encode_msg(body=struct.pack("<QQQQ", 100, 200, 300, 400)))
         try_write(
             encode_msg(