diff --git a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp index 4aad1566d6995e497bcdcd27a5df63d9446b21de..d9bd049c985939c4d804f0b61b12ba0ecca0b88e 100644 --- a/frameworks/native/neural_network_runtime/nncompiled_cache.cpp +++ b/frameworks/native/neural_network_runtime/nncompiled_cache.cpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include "common/utils.h" #include "backend_manager.h" @@ -197,12 +199,19 @@ OH_NN_ReturnCode NNCompiledCache::GenerateCacheModel(const std::vector(version); *cacheInfoPtr++ = static_cast(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 49dc108e56319a5b311dd13cb833ccdc0699ebe3..1140fab89bb37ad3c28f62cc5e65dbe7348114a7 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 56ab88632363f1d65b9d2a3514243a5a4cd9f8a4..cbc89e6d75d09e49549418d742426e3d6a80c7f0 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; }