Add dtype prefix check

Signed-off-by: Mickael Ide <mide@nvidia.com>

Author: lowener

cpp/src/neighbors/brute_force_serialize.cu

@@ -82,8 +82,10 @@ void serialize(raft::resources const& handle,
 template <typename T, typename DistT>
 auto deserialize(raft::resources const& handle, std::istream& is)
 {
-  auto dtype_string = std::array<char, 4>{};
-  is.read(dtype_string.data(), 4);
+  char dtype_string[4];
+  RAFT_EXPECTS(is.read(dtype_string, 4), "brute_force::deserialize: failed to read dtype prefix");
+  RAFT_EXPECTS(cuvs::util::validate_serialized_dtype<T>(dtype_string, sizeof(dtype_string)),
+               "brute_force::deserialize: serialized dtype prefix does not match requested type");
 
   auto ver = raft::deserialize_scalar<int>(handle, is);
   if (ver != serialization_version) {

cpp/src/neighbors/detail/cagra/cagra_serialize.cuh

@@ -269,7 +269,9 @@ void deserialize(raft::resources const& res, std::istream& is, index<T, IdxT>* i
   raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope("cagra::deserialize");
 
   char dtype_string[4];
-  is.read(dtype_string, 4);
+  RAFT_EXPECTS(is.read(dtype_string, 4), "cagra::deserialize: failed to read dtype prefix");
+  RAFT_EXPECTS(cuvs::util::validate_serialized_dtype<T>(dtype_string, sizeof(dtype_string)),
+               "cagra::deserialize: serialized dtype prefix does not match requested type");
 
   auto ver = raft::deserialize_scalar<int>(res, is);
   if (ver != serialization_version) {

cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh

@@ -108,7 +108,9 @@ template <typename T, typename IdxT>
 auto deserialize(raft::resources const& handle, std::istream& is) -> index<T, IdxT>
 {
   char dtype_string[4];
-  is.read(dtype_string, 4);
+  RAFT_EXPECTS(is.read(dtype_string, 4), "ivf_flat::deserialize: failed to read dtype prefix");
+  RAFT_EXPECTS(cuvs::util::validate_serialized_dtype<T>(dtype_string, sizeof(dtype_string)),
+               "ivf_flat::deserialize: serialized dtype prefix does not match requested type");
 
   auto ver = raft::deserialize_scalar<int>(handle, is);
   if (ver != serialization_version) {

cpp/src/util/serialize_validation.hpp

@@ -7,11 +7,14 @@
 #include <cuvs/distance/distance.hpp>
 #include <cuvs/neighbors/ivf_pq.hpp>
 
+#include <raft/core/detail/mdspan_numpy_serializer.hpp>
 #include <raft/core/error.hpp>
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <limits>
+#include <string>
 #include <type_traits>
 
 namespace cuvs::util {
@@ -36,6 +39,17 @@ inline bool is_mul_no_overflow(T a, T b, Rest... rest)
   }
 }
 
+template <typename T>
+inline bool validate_serialized_dtype(const char* dtype_prefix, std::size_t dtype_prefix_size)
+{
+  if (dtype_prefix == nullptr || dtype_prefix_size != 4) { return false; }
+
+  auto expected_dtype = raft::detail::numpy_serializer::get_numpy_dtype<T>().to_string();
+  expected_dtype.resize(dtype_prefix_size, '\0');
+
+  return std::equal(dtype_prefix, dtype_prefix + dtype_prefix_size, expected_dtype.begin());
+}
+
 inline bool is_valid_distance_type(cuvs::distance::DistanceType m)
 {
   using cuvs::distance::DistanceType;