diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp index 498fc5ac7c0cd5751567c7a60484d0812eb20898..23ed62c8ca539e903b39f3912497f2bc07ef6aa3 100644 --- a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp +++ b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp @@ -190,6 +190,10 @@ OH_NN_ReturnCode NNCompiledCache::GenerateCacheModel(const std::vector NN_CACHE_FILE_NUMBER_MAX) { + LOGE("[NNCompiledCache] Caches size is equal 0 or greater than 100."); + return OH_NN_FAILED; + } auto cacheInfoPtr = cacheInfo.get(); *cacheInfoPtr++ = static_cast(cacheNumber); diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.h b/frameworks/native/neural_network_runtime/nncompiled_cache.h index df84886a564a9bf78200ca8afa4699dbb2ae21e5..fff2c1f50ef1315156401ba6c949065b45e8e337 100644 --- a/frameworks/native/neural_network_runtime/nncompiled_cache.h +++ b/frameworks/native/neural_network_runtime/nncompiled_cache.h @@ -27,6 +27,7 @@ namespace OHOS { namespace NeuralNetworkRuntime { const uint32_t INVALID_CAHCE_VERSION = UINT32_MAX; // UINT32_MAX is reserved for invalid cache version. +constexpr size_t NN_CACHE_FILE_NUMBER_MAX = 100; // 限制cache文件数量最大为100 struct NNCompiledCacheInfo { int64_t fileNumber{0}; diff --git a/frameworks/native/neural_network_runtime/nncompiler.cpp b/frameworks/native/neural_network_runtime/nncompiler.cpp index d1aa30e03e8e4e2887f2940d952587462874baa3..ed72e0808baf5e378f1c376b1c70d070b7b54871 100644 --- a/frameworks/native/neural_network_runtime/nncompiler.cpp +++ b/frameworks/native/neural_network_runtime/nncompiler.cpp @@ -531,6 +531,12 @@ OH_NN_ReturnCode NNCompiler::SaveToCacheFile() const return ret; } + size_t cacheNumber = caches.size(); + if (cacheNumber == 0 || cacheNumber > NN_CACHE_FILE_NUMBER_MAX) { + LOGE("[NNCompiler] Caches size is equal 0 or greater than 100."); + return OH_NN_FAILED; + } + NNCompiledCache compiledCache; ret = compiledCache.SetBackend(m_backendID); if (ret != OH_NN_SUCCESS) {