diff --git a/VisionSDK/GeneralTextRecognition/C++/main.cpp b/VisionSDK/GeneralTextRecognition/C++/main.cpp index 647c57c1f4ee8eb0237f50160c9f477ae86d0168..d5d080d1fc85466afffdada353de5ea9bda0b4ee 100644 --- a/VisionSDK/GeneralTextRecognition/C++/main.cpp +++ b/VisionSDK/GeneralTextRecognition/C++/main.cpp @@ -29,6 +29,7 @@ namespace { const int TIME_OUT = 20000; const float SEC2MS = 1000.0; const std::string PICTURE_PATH = "../input_data/"; + const long MAX_FILE_SIZE = 1024 * 1024 * 1024; // 1G } std::string ReadFileContent(const std::string& filePath) @@ -42,6 +43,10 @@ std::string ReadFileContent(const std::string& filePath) file.seekg(0, std::ifstream::end); uint32_t fileSize = file.tellg(); file.seekg(0); + if (fileSize == 0 || fileSize > MAX_FILE_SIZE){ + LogError << "File size invalid"; + return ""; + } std::vector buffer(fileSize); file.read(buffer.data(), fileSize); file.close(); @@ -71,6 +76,7 @@ APP_ERROR GetPicture(const std::string& filePath, std::vector& pict pictureName.push_back(filename); } } + closedir(pDir); return APP_ERR_OK; } diff --git a/VisionSDK/GeneralTextRecognition/python/main_ocr.py b/VisionSDK/GeneralTextRecognition/python/main_ocr.py index efdb38ab5a8fc3764d048e851b2f739711635dc3..58e7248ace18bea01dfa43243e9577a2afa29992 100644 --- a/VisionSDK/GeneralTextRecognition/python/main_ocr.py +++ b/VisionSDK/GeneralTextRecognition/python/main_ocr.py @@ -22,6 +22,10 @@ import json from StreamManagerApi import StreamManagerApi, MxDataInput +MAX_FILE_SIZE = 1024 * 1024 * 10 # 10MB +MAX_IMAGE_SIZE = 1024 * 1024 * 1024 # 1G +PATH = "../data/OCR.pipeline" + if __name__ == '__main__': # init stream manager STREAM_NAME = b'OCR' @@ -33,7 +37,10 @@ if __name__ == '__main__': exit() # create streams by pipeline config file - with open("../data/OCR.pipeline", 'rb') as f: + if os.path.getsize(PATH) <= 0 or os.path.getsize(PATH) > MAX_FILE_SIZE: + print("Pipeline file size invalid!") + exit() + with open(PATH, 'rb') as f: pipelineStr = f.read() ret = streamManagerApi.CreateMultipleStreams(pipelineStr) if ret != 0: @@ -54,6 +61,9 @@ if __name__ == '__main__': if lower_file_name.endswith(".jpeg") or lower_file_name.endswith(".jpg") or lower_file_name.endswith(".png"): img_path = os.path.join(DIR_NAME, file_name) print("img_path: ", img_path) + if os.path.getsize(img_path) <= 0 or os.path.getsize(img_path) > MAX_IMAGE_SIZE: + print("Image file size invalid!") + exit() with open(img_path, 'rb') as f: dataInput.data = f.read() diff --git a/tutorials/OsdSample/C++/main.cpp b/tutorials/OsdSample/C++/main.cpp index 917d22e658ec67c4986fbc5446f2ba82c72f7698..e7a0fcaa174acb8b6f1366df4e20dca8daf71cd8 100644 --- a/tutorials/OsdSample/C++/main.cpp +++ b/tutorials/OsdSample/C++/main.cpp @@ -23,6 +23,8 @@ namespace { const int SLEEP_TIME = 2; +const long MAX_IMAGE_SIZE = 1024 * 1024 * 1024; // 1G +const long MAX_FILE_SIZE = 1024 * 1024 * 100; // 100MB APP_ERROR ReadFile(const std::string& filePath, MxStream::MxstDataInput& dataBuffer) { char c[PATH_MAX + 1] = { 0x00 }; @@ -48,7 +50,7 @@ APP_ERROR ReadFile(const std::string& filePath, MxStream::MxstDataInput& dataBuf long fileSize = ftell(fp); fseek(fp, 0, SEEK_SET); // If file not empty, read it into FileInfo and return it - if (fileSize > 0) { + if (fileSize > 0 && fileSize < MAX_IMAGE_SIZE) { dataBuffer.dataSize = fileSize; dataBuffer.dataPtr = new (std::nothrow) uint32_t[fileSize]; if (dataBuffer.dataPtr == nullptr) { @@ -79,6 +81,10 @@ std::string ReadFileContent(const std::string& filePath) file.seekg(0, std::ifstream::end); uint32_t fileSize = file.tellg(); file.seekg(0); + if (fileSize ==0 || fileSize > MAX_FILE_SIZE){ + LogError << "Invalid file size"; + return ""; + } std::vector buffer; buffer.resize(fileSize); file.read(buffer.data(), fileSize); @@ -97,6 +103,10 @@ std::string ReadPipelineConfig(const std::string& pipelineConfigPath) file.seekg(0, std::ifstream::end); uint32_t fileSize = file.tellg(); file.seekg(0); + if (fileSize ==0 || fileSize > MAX_FILE_SIZE){ + LogError << "Invalid file size"; + return ""; + } std::unique_ptr data(new char[fileSize]); file.read(data.get(), fileSize); file.close(); @@ -163,7 +173,7 @@ int main(int argc, char* argv[]) // destroy streams mxStreamManager.DestroyAllStreams(); - delete dataBuffer.dataPtr; + delete[] dataBuffer.dataPtr; dataBuffer.dataPtr = nullptr; return 0; diff --git a/tutorials/protocolSample/main.cpp b/tutorials/protocolSample/main.cpp index a05cc173693e47dbe9197316b29994f45c9073cf..f2c0a9547390a19c8538b45d59170cd31d63bd9c 100644 --- a/tutorials/protocolSample/main.cpp +++ b/tutorials/protocolSample/main.cpp @@ -110,6 +110,8 @@ namespace { constexpr int COS_LEN_VALUE = 416; constexpr int COS_FORMAT_VALUE = 12; constexpr bool USE_SENDDATA = true; // switch different send data method + const long MAX_IMAGE_SIZE = 1024 * 1024 * 1024; // 1G + const long MAX_FILE_SIZE = 1024 * 1024 * 100; // 100MB // read data File to MxstDataInput structure APP_ERROR ReadFile(const std::string& filePath, MxStream::MxstDataInput& dataBuffer) @@ -137,7 +139,7 @@ namespace { long fileSize = ftell(fp); fseek(fp, 0, SEEK_SET); // If file not empty, read it into FileInfo and return it - if (fileSize > 0) { + if (fileSize > 0 && fileSize < MAX_IMAGE_SIZE) { dataBuffer.dataSize = fileSize; dataBuffer.dataPtr = new (std::nothrow) uint32_t[fileSize]; if (dataBuffer.dataPtr == nullptr) { @@ -169,6 +171,10 @@ namespace { file.seekg(0, std::ifstream::end); uint32_t fileSize = file.tellg(); file.seekg(0); + if (fileSize ==0 || fileSize > MAX_FILE_SIZE){ + LogError << "Invalid file size"; + return ""; + } auto data = std::unique_ptr(new char[fileSize]); file.read(data.get(), fileSize); file.close(); @@ -377,7 +383,7 @@ int main(int argc, char* argv[]) // destroy streams mxStreamManager.DestroyAllStreams(); - delete dataBuffer.dataPtr; + delete[] dataBuffer.dataPtr; dataBuffer.dataPtr = nullptr; return 0; diff --git a/tutorials/sampleCustomProto/protoRead/main.py b/tutorials/sampleCustomProto/protoRead/main.py index 2ff67a3b69e90f8178e42b24b95628dab1bcae82..8df8f236fafe82f929b40d2f0d9395f2084316da 100644 --- a/tutorials/sampleCustomProto/protoRead/main.py +++ b/tutorials/sampleCustomProto/protoRead/main.py @@ -17,12 +17,14 @@ See the License for the specific language governing permissions and limitations under the License. """ -from StreamManagerApi import StreamManagerApi, MxDataInput, StringVector - import sys +import os sys.path.append("../../proto") import mxpiSampleProto_pb2 as mxpiSampleProto +from StreamManagerApi import StreamManagerApi, MxDataInput, StringVector +MAX_FILE_SIZE = 1024 * 1204 * 10 # 10MB +PATH = "../pipeline/Sample_proto.pipeline" if __name__ == '__main__': # init stream manager streamManagerApi = StreamManagerApi() @@ -32,7 +34,10 @@ if __name__ == '__main__': exit() # create streams by pipeline config file - with open("../pipeline/Sample_proto.pipeline", 'rb') as f: + if os.path.getsize(PATH) <= 0 or os.path.getsize(PATH) > MAX_FILE_SIZE: + print("Pipeline file size invalid!") + exit() + with open(PATH, 'rb') as f: pipelineStr = f.read() ret = streamManagerApi.CreateMultipleStreams(pipelineStr) if ret != 0: