From f40dcd6733f56dad0d1ded2a366089120a74cef4 Mon Sep 17 00:00:00 2001
From: maoyong <maoyong5@huawei.com>
Date: Mon, 6 May 2024 16:05:16 +0800
Subject: [PATCH] fix cache bug

Signed-off-by: maoyong <maoyong5@huawei.com>
---
 .../nncompiled_cache.cpp                      | 38 ++++++++++++++++++-
 .../neural_network_runtime/nncompiled_cache.h |  1 +
 .../neural_network_runtime/nncompiler.cpp     |  4 +-
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp
index 4aad156..d9bd049 100644
--- a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp
+++ b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp
@@ -18,6 +18,8 @@
 #include <unistd.h>
 #include <functional>
 #include <memory>
+#include <cstdlib>
+#include <limits>
 
 #include "common/utils.h"
 #include "backend_manager.h"
@@ -197,12 +199,19 @@ OH_NN_ReturnCode NNCompiledCache::GenerateCacheModel(const std::vector<OHOS::Neu
     *cacheInfoPtr++ = static_cast<uint64_t>(version);
     *cacheInfoPtr++ = static_cast<uint64_t>(m_backendID); // Should call SetBackend first.
 
+    OH_NN_ReturnCode ret = OH_NN_SUCCESS;
     for (size_t i = 0; i < cacheNumber; ++i) {
         std::string cacheModelFile = cacheDir + "/" + m_modelName + std::to_string(i) + ".nncache";
         char path[PATH_MAX];
         if (realpath(cacheModelFile.c_str(), path) == nullptr) {
             LOGE("[NNCompiledCache] GenerateCacheModel failed, fail to get the real path of cacheModelFile.");
-            return OH_NN_INVALID_PARAMETER;
+            return OH_NN_SAVE_CACHE_EXCEPTION;
+        }
+
+        ret = VerifyCachePath(path);
+        if (ret != OH_NN_SUCCESS) {
+            LOGE("[NNCompiledCache] GenerateCacheModel failed, fail to verify the file path of cacheModelFile.");
+            return ret;
         }
 
         std::ofstream cacheModelStream(cacheModelFile, std::ios::binary | std::ios::out | std::ios::trunc);
@@ -234,7 +243,13 @@ OH_NN_ReturnCode NNCompiledCache::WriteCacheInfo(uint32_t cacheSize,
     char path[PATH_MAX];
     if (realpath(cacheInfoPath.c_str(), path) == nullptr) {
         LOGE("[NNCompiledCache] WriteCacheInfo failed, fail to get the real path of cacheInfoPath.");
-        return OH_NN_INVALID_PARAMETER;
+        return OH_NN_SAVE_CACHE_EXCEPTION;
+    }
+
+    OH_NN_ReturnCode ret = VerifyCachePath(path);
+    if (ret != OH_NN_SUCCESS) {
+        LOGE("[NNCompiledCache] WriteCacheInfo failed, fail to verify the file path of cacheInfoPath.");
+        return ret;
     }
 
     std::ofstream cacheInfoStream(cacheInfoPath, std::ios::binary | std::ios::out | std::ios::trunc);
@@ -392,5 +407,24 @@ OH_NN_ReturnCode NNCompiledCache::GetCacheFileLength(std::ifstream& ifs, int& fi
     fileSize = handleValue;
     return OH_NN_SUCCESS;
 }
+
+OH_NN_ReturnCode NNCompiledCache::VerifyCachePath(const std::string& cachePath) const
+{
+    if (cachePath.find('/') != size_t(0)) {
+        LOGE("[NNCompiledCache] VerifyCachePath failed, input file dir=%{public}s is invalid, "
+             "should start with '/'.",
+             cachePath.c_str());
+        return OH_NN_INVALID_FILE;
+    }
+
+    if (cachePath.find("//") != std::string::npos) {
+        LOGE("[NNCompiledCache] VerifyCachePath failed, input file dir=%{public}s is invalid, "
+             "containing double '/'.",
+             cachePath.c_str());
+        return OH_NN_INVALID_FILE;
+    }
+
+    return OH_NN_SUCCESS;
+}
 } // namespace NeuralNetworkRuntime
 } // namespace OHOS
diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.h b/frameworks/native/neural_network_runtime/nncompiled_cache.h
index 49dc108..1140fab 100644
--- a/frameworks/native/neural_network_runtime/nncompiled_cache.h
+++ b/frameworks/native/neural_network_runtime/nncompiled_cache.h
@@ -65,6 +65,7 @@ private:
     OH_NN_ReturnCode ReadCacheModelFile(const std::string& file, Buffer& cache) const;
     unsigned short GetCrc16(char* buffer, size_t length) const;
     OH_NN_ReturnCode GetCacheFileLength(std::ifstream& ifs, int& fileSize) const;
+    OH_NN_ReturnCode VerifyCachePath(const std::string& cachePath) const;
 
 private:
     size_t m_backendID {0};
diff --git a/frameworks/native/neural_network_runtime/nncompiler.cpp b/frameworks/native/neural_network_runtime/nncompiler.cpp
index 56ab886..cbc89e6 100644
--- a/frameworks/native/neural_network_runtime/nncompiler.cpp
+++ b/frameworks/native/neural_network_runtime/nncompiler.cpp
@@ -391,7 +391,9 @@ OH_NN_ReturnCode NNCompiler::NormalBuild()
     // 保存cache
     if (!m_cachePath.empty()) {
         ret = SaveToCacheFile();
-        if (ret != OH_NN_SUCCESS) {
+        if (ret == OH_NN_SAVE_CACHE_EXCEPTION) {
+            LOGE("[NNCompiler] Build success, but some error happened when write cache, continue running.");
+        } else if (ret != OH_NN_SUCCESS) {
             LOGE("[NNCompiler] Build success, but fail to save cache to file.");
             return ret;
         }
-- 
Gitee