Add [[nodiscard]] to ErrorCode

AnnService/inc/Core/Common.h

@@ -130,6 +130,8 @@ extern std::shared_ptr<Helper::DiskIO>(*f_createIO)();
 
 #define IOBINARY(ptr, func, bytes, ...) if (ptr->func(bytes, __VA_ARGS__) != bytes) return ErrorCode::DiskIOFail
 #define IOSTRING(ptr, func, ...) if (ptr->func(__VA_ARGS__) == 0) return ErrorCode::DiskIOFail
+#define RETURN_IF_ERROR(expr) { ErrorCode ret = expr; if (ret != ErrorCode::Success) return ret; }
+#define RETURN_IF_ERROR_WITH_LOG(expr, log_level, ...) { ErrorCode ret = expr; if (ret != ErrorCode::Success) { SPTAGLIB_LOG(log_level, __VA_ARGS__); return ret; } }
 
 extern Helper::LoggerHolder& GetLoggerHolder();
 extern std::shared_ptr<Helper::Logger> GetLogger();
@@ -150,7 +152,7 @@ class MyException : public std::exception
 #endif
 };
 
-enum class ErrorCode : std::uint16_t
+enum class [[nodiscard]] ErrorCode : std::uint16_t
 {
 #define DefineErrorCode(Name, Value) Name = Value,
 #include "DefinitionList.h"

AnnService/inc/Core/Common/IQuantizer.h

@@ -39,7 +39,7 @@ namespace SPTAG
 
             virtual ErrorCode LoadQuantizer(std::shared_ptr<Helper::DiskIO> p_in) = 0;
 
-            virtual ErrorCode LoadQuantizer(uint8_t* raw_bytes) = 0;
+            virtual ErrorCode LoadQuantizer(uint8_t* raw_bytes, SizeType num_bytes) = 0;
 
             static std::shared_ptr<IQuantizer> LoadIQuantizer(std::shared_ptr<Helper::DiskIO> p_in);
 

AnnService/inc/Core/Common/NeighborhoodGraph.h

@@ -506,6 +506,7 @@ break;
                 newGraph->m_iGraphSize = R;
                 newGraph->m_iNeighborhoodSize = m_iNeighborhoodSize;
 
+                ErrorCode ret = ErrorCode::Success;
 #pragma omp parallel for schedule(dynamic)
                 for (SizeType i = 0; i < R; i++)
                 {
@@ -514,7 +515,11 @@ break;
                     SizeType* outnodes = newGraph->m_pNeighborhoodGraph[i];
 
                     COMMON::QueryResultSet<T> query((const T*)index->GetSample(indices[i]), m_iCEF + 1);
-                    index->RefineSearchIndex(query, false);
+                    ErrorCode internal_ret =index->RefineSearchIndex(query, false);
+                    if (internal_ret != ErrorCode::Success) {
+                        SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "RefineSearchIndex failed \n");
+                        ret = internal_ret;
+                    }
                     RebuildNeighbors(index, indices[i], outnodes, query.GetResults(), m_iCEF + 1);
 
                     std::unordered_map<SizeType, SizeType>::const_iterator iter;
@@ -527,8 +532,8 @@ break;
                         outnodes[m_iNeighborhoodSize - 1] = -2 - iter->second;
                 }
 
-                if (output != nullptr) newGraph->SaveGraph(output);
-                return ErrorCode::Success;
+                if (output != nullptr) RETURN_IF_ERROR(newGraph->SaveGraph(output));
+                return ret;
             }
 
             template <typename T>

AnnService/inc/Core/Common/OPQQuantizer.h

@@ -34,7 +34,7 @@ namespace SPTAG
 
             virtual ErrorCode LoadQuantizer(std::shared_ptr<Helper::DiskIO> p_in);
 
-            virtual ErrorCode LoadQuantizer(std::uint8_t* raw_bytes);
+            virtual ErrorCode LoadQuantizer(std::uint8_t* raw_bytes, SizeType num_bytes);
 
             virtual SizeType ReconstructSize() const;
 
@@ -164,13 +164,26 @@ namespace SPTAG
         }
 
         template <typename T>
-        ErrorCode OPQQuantizer<T>::LoadQuantizer(std::uint8_t* raw_bytes)
+        ErrorCode OPQQuantizer<T>::LoadQuantizer(std::uint8_t* raw_bytes, SizeType num_bytes)
         {
-            PQQuantizer<OPQMatrixType>::LoadQuantizer(raw_bytes);
+            std::uint8_t* original_ptr = raw_bytes;
+            ErrorCode code = PQQuantizer<OPQMatrixType>::LoadQuantizer(raw_bytes, num_bytes);
+            if (code != ErrorCode::Success)
+            {
+                return code;
+            }
+
+            // read OPQ matrix
             raw_bytes += sizeof(DimensionType) + sizeof(SizeType) + sizeof(DimensionType) + (sizeof(OPQMatrixType) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector);
             m_matrixDim = m_NumSubvectors * m_DimPerSubvector;
             m_OPQMatrix = std::make_unique<OPQMatrixType[]>(m_matrixDim * m_matrixDim);
-            std::memcpy(m_OPQMatrix.get(), raw_bytes, sizeof(OPQMatrixType) * m_matrixDim * m_matrixDim);
+
+            SizeType matrix_bytes = sizeof(OPQMatrixType) * m_matrixDim * m_matrixDim;
+            if (raw_bytes - original_ptr + matrix_bytes > num_bytes) {
+                SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Not enough bytes for OPQ matrix\n");
+                return ErrorCode::Fail;
+            }
+            std::memcpy(m_OPQMatrix.get(), raw_bytes, matrix_bytes);
             raw_bytes += sizeof(OPQMatrixType) * m_matrixDim * m_matrixDim;
 
             m_InitMatrixTranspose();

AnnService/inc/Core/Common/PQQuantizer.h

@@ -49,7 +49,7 @@ namespace SPTAG
 
             virtual ErrorCode LoadQuantizer(std::shared_ptr<Helper::DiskIO> p_in);
 
-            virtual ErrorCode LoadQuantizer(std::uint8_t* raw_bytes);
+            virtual ErrorCode LoadQuantizer(std::uint8_t* raw_bytes, SizeType num_bytes);
 
             virtual DimensionType GetNumSubvectors() const;
 
@@ -260,20 +260,38 @@ namespace SPTAG
         }
 
         template <typename T>
-        ErrorCode PQQuantizer<T>::LoadQuantizer(std::uint8_t* raw_bytes)
+        ErrorCode PQQuantizer<T>::LoadQuantizer(std::uint8_t* raw_bytes, SizeType num_bytes)
         {
+            std::uint8_t* original_ptr = raw_bytes;
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Loading Quantizer.\n");
             m_NumSubvectors = *(DimensionType*)raw_bytes;
             raw_bytes += sizeof(DimensionType);
+            if (raw_bytes - original_ptr > num_bytes) {
+                SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Not enough bytes to read NumSubvectors.\n");
+                return ErrorCode::MemoryOverFlow;
+            }
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read subvecs: %s.\n", std::to_string(m_NumSubvectors).c_str());
             m_KsPerSubvector = *(SizeType*)raw_bytes;
             raw_bytes += sizeof(SizeType);
+            if (raw_bytes - original_ptr > num_bytes) {
+                SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Not enough bytes to read KsPerSubvector.\n");
+                return ErrorCode::MemoryOverFlow;
+            }
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read ks: %s.\n", std::to_string(m_KsPerSubvector).c_str());
             m_DimPerSubvector = *(DimensionType*)raw_bytes;
             raw_bytes += sizeof(DimensionType);
+            if (raw_bytes - original_ptr > num_bytes) {
+                SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Not enough bytes to read DimPerSubvector.\n");
+                return ErrorCode::MemoryOverFlow;
+            }
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read dim: %s.\n", std::to_string(m_DimPerSubvector).c_str());
             m_codebooks = std::make_unique<T[]>(m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector);
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "sizeof(T): %s.\n", std::to_string(sizeof(T)).c_str());
+            SizeType codebook_bytes = sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector;
+            if (raw_bytes - original_ptr + codebook_bytes > num_bytes) {
+                SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Not enough bytes to read codebooks.\n");
+                return ErrorCode::MemoryOverFlow;
+            }
             std::memcpy(m_codebooks.get(), raw_bytes, sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector);
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read codebooks.\n");
 

AnnService/inc/Core/VectorIndex.h

@@ -188,7 +188,7 @@ class VectorIndex
 
     inline SizeType GetMetaMapping(std::string& meta) const;
 
-    void UpdateMetaMapping(const std::string& meta, SizeType i);
+    ErrorCode UpdateMetaMapping(const std::string& meta, SizeType i);
 
     void BuildMetaMapping(bool p_checkDeleted = true);
 

AnnService/src/Core/BKT/BKTIndex.cpp

@@ -20,11 +20,13 @@ namespace SPTAG
         template <typename T>
         ErrorCode Index<T>::LoadConfig(Helper::IniReader& p_reader)
         {
+            ErrorCode code = ErrorCode::Success;
 #define DefineBKTParameter(VarName, VarType, DefaultValue, RepresentStr) \
-            SetParameter(RepresentStr, \
+            code = SetParameter(RepresentStr, \
                          p_reader.GetParameter("Index", \
                          RepresentStr, \
                          std::string(#DefaultValue)).c_str()); \
+            if (code != ErrorCode::Success) return code;
 
 #include "inc/Core/BKT/ParameterDefinitionList.h"
 #undef DefineBKTParameter
@@ -816,7 +818,7 @@ namespace SPTAG
             ptr->m_deletedID.Initialize(newR, m_iDataBlockSize, m_iDataCapacity, COMMON::Labelset::InvalidIDBehavior::AlwaysContains);
             COMMON::BKTree* newtree = &(ptr->m_pTrees);
             (*newtree).BuildTrees<T>(ptr->m_pSamples, ptr->m_iDistCalcMethod, omp_get_num_threads());
-            m_pGraph.RefineGraph<T>(this, indices, reverseIndices, nullptr, &(ptr->m_pGraph), &(ptr->m_pTrees.GetSampleMap()));
+            if ((ret = m_pGraph.RefineGraph<T>(this, indices, reverseIndices, nullptr, &(ptr->m_pGraph), &(ptr->m_pTrees.GetSampleMap()))) != ErrorCode::Success) return ret;
             if (HasMetaMapping()) ptr->BuildMetaMapping(false);
             ptr->m_bReady = true;
             return ret;
@@ -875,18 +877,28 @@ namespace SPTAG
         template <typename T>
         ErrorCode Index<T>::DeleteIndex(const void* p_vectors, SizeType p_vectorNum) {
             const T* ptr_v = (const T*)p_vectors;
+            ErrorCode ret = ErrorCode::Success;
 #pragma omp parallel for schedule(dynamic)
             for (SizeType i = 0; i < p_vectorNum; i++) {
                 COMMON::QueryResultSet<T> query(ptr_v + i * GetFeatureDim(), m_pGraph.m_iCEF);
-                SearchIndex(query);
+                ErrorCode search_ret = SearchIndex(query);
+                if (search_ret != ErrorCode::Success || query.GetResultNum() == 0) {
+                    SPTAGLIB_LOG(Helper::LogLevel::LL_Warning, "Cannot find vector to delete!\n");
+                    ret = search_ret;
+                }
 
                 for (int j = 0; j < m_pGraph.m_iCEF; j++) {
                     if (query.GetResult(j)->Dist < 1e-6) {
-                        DeleteIndex(query.GetResult(j)->VID);
+                        SizeType vid = query.GetResult(j)->VID;
+                        ErrorCode delete_ret = DeleteIndex(vid);
+                        if (delete_ret != ErrorCode::Success) {
+                            SPTAGLIB_LOG(Helper::LogLevel::LL_Warning, "Cannot delete vector! ID: %llu\n", vid);
+                            ret = delete_ret;
+                        }
                     }
                 }
             }
-            return ErrorCode::Success;
+            return ret;
         }
 
         template <typename T>
@@ -940,7 +952,7 @@ namespace SPTAG
                             for (SizeType i = begin; i < end; i++) {
                                 ByteArray meta = m_pMetadata->GetMetadata(i);
                                 std::string metastr((char*)meta.Data(), meta.Length());
-                                UpdateMetaMapping(metastr, i);
+                                RETURN_IF_ERROR(UpdateMetaMapping(metastr, i));
                             }
                         }
                     }

AnnService/src/Core/Common/IQuantizer.cpp

@@ -55,11 +55,14 @@ namespace SPTAG
         std::shared_ptr<IQuantizer> IQuantizer::LoadIQuantizer(SPTAG::ByteArray bytes)
         {
             auto raw_bytes = bytes.Data();
+            SizeType num_bytes = bytes.Length();
             QuantizerType quantizerType = *(QuantizerType*) raw_bytes;
             raw_bytes += sizeof(QuantizerType);
+            num_bytes -= sizeof(QuantizerType);
 
             VectorValueType reconstructType = *(VectorValueType*)raw_bytes;
             raw_bytes += sizeof(VectorValueType);
+            num_bytes -= sizeof(VectorValueType);
             SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Loading quantizer of type %s with reconstructtype %s.\n", Helper::Convert::ConvertToString<QuantizerType>(quantizerType).c_str(), Helper::Convert::ConvertToString<VectorValueType>(reconstructType).c_str());
             std::shared_ptr<IQuantizer> ret = nullptr;
 
@@ -80,7 +83,7 @@ namespace SPTAG
                 default: break;
                 }
 
-                if (ret->LoadQuantizer(raw_bytes) != ErrorCode::Success) ret.reset();
+                if (ret->LoadQuantizer(raw_bytes, num_bytes) != ErrorCode::Success) ret.reset();
                 return ret;
             case QuantizerType::OPQQuantizer:
                 switch (reconstructType) {
@@ -94,7 +97,7 @@ namespace SPTAG
                 default: break;
                 }
 
-                if (ret->LoadQuantizer(raw_bytes) != ErrorCode::Success) ret.reset();
+                if (ret->LoadQuantizer(raw_bytes, num_bytes) != ErrorCode::Success) ret.reset();
                 return ret;
             }
             return ret;

AnnService/src/Core/KDT/KDTIndex.cpp

@@ -666,7 +666,7 @@ case VectorValueType::Name: \
                             for (SizeType i = begin; i < end; i++) {
                                 ByteArray meta = m_pMetadata->GetMetadata(i);
                                 std::string metastr((char*)meta.Data(), meta.Length());
-                                UpdateMetaMapping(metastr, i);
+                                RETURN_IF_ERROR(UpdateMetaMapping(metastr, i));
                             }
                         }
                     }