From 7e6d3709744e286554e5fef1e4d40af083eb1443 Mon Sep 17 00:00:00 2001
From: w30052974 <wangyifan94@huawei.com>
Date: Fri, 20 Dec 2024 15:44:50 +0800
Subject: [PATCH] nnrt cache file number limited

Signed-off-by: w30052974 <wangyifan94@huawei.com>
---
 .../native/neural_network_runtime/nncompiled_cache.cpp      | 4 ++++
 frameworks/native/neural_network_runtime/nncompiled_cache.h | 1 +
 frameworks/native/neural_network_runtime/nncompiler.cpp     | 6 ++++++
 3 files changed, 11 insertions(+)

diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp
index 498fc5a..23ed62c 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<OHOS::Neu
                                                      uint32_t version) const
 {
     size_t cacheNumber = caches.size();
+    if (cacheNumber == 0 || cacheNumber > 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<int64_t>(cacheNumber);
diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.h b/frameworks/native/neural_network_runtime/nncompiled_cache.h
index df84886..fff2c1f 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 d1aa30e..ed72e08 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) {
-- 
Gitee