1 /**
2 * Copyright 2020-2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "tools/common/protobuf_utils.h"
18 #include <fstream>
19 #include <string>
20 #include "google/protobuf/io/zero_copy_stream_impl.h"
21 #include "google/protobuf/text_format.h"
22 #include "google/protobuf/io/coded_stream.h"
23 #include "src/common/file_utils.h"
24
25 namespace mindspore {
26 namespace lite {
27 static const int PROTO_READ_BYTES_LIMIT = INT_MAX; // Max size of 2 GB minus 1 byte.
28 static const int WARNING_THRESHOLD = 536870912 * 2;
29
ReadProtoFromCodedInputStream(google::protobuf::io::CodedInputStream * coded_stream,google::protobuf::Message * proto)30 bool ReadProtoFromCodedInputStream(google::protobuf::io::CodedInputStream *coded_stream,
31 google::protobuf::Message *proto) {
32 if (proto == nullptr) {
33 MS_LOG(ERROR) << "incorrect parameter. nullptr == proto";
34 return false;
35 }
36 coded_stream->SetTotalBytesLimit(PROTO_READ_BYTES_LIMIT, WARNING_THRESHOLD);
37 return proto->ParseFromCodedStream(coded_stream);
38 }
39
ReadProtoFromText(const std::string & file,google::protobuf::Message * message)40 STATUS ReadProtoFromText(const std::string &file, google::protobuf::Message *message) {
41 if (file.empty() || message == nullptr) {
42 return RET_ERROR;
43 }
44
45 std::string realPath = RealPath(file.c_str());
46 if (realPath.empty()) {
47 MS_LOG(ERROR) << "Proto file path " << file << " is not valid";
48 return RET_ERROR;
49 }
50
51 std::ifstream fs(realPath.c_str(), std::ifstream::in);
52
53 if (!fs.is_open()) {
54 MS_LOG(ERROR) << "Open proto file " << file << " failed.";
55 return RET_ERROR;
56 }
57
58 google::protobuf::io::IstreamInputStream input(&fs);
59 bool status = google::protobuf::TextFormat::Parse(&input, message);
60 if (!status) {
61 MS_LOG(ERROR) << "call [google::protobuf::TextFormat::Parse] func status fail, please check your text file.";
62 fs.close();
63 return RET_ERROR;
64 }
65
66 fs.close();
67 return RET_OK;
68 }
69
ReadProtoFromBinaryFile(const std::string & file,google::protobuf::Message * message)70 STATUS ReadProtoFromBinaryFile(const std::string &file, google::protobuf::Message *message) {
71 if (file.empty() || message == nullptr) {
72 return RET_ERROR;
73 }
74
75 std::string realPath = RealPath(file.c_str());
76 if (realPath.empty()) {
77 MS_LOG(ERROR) << "Binary proto file path " << file << " is not valid";
78 return RET_ERROR;
79 }
80
81 std::ifstream fs(realPath, std::ifstream::in | std::ifstream::binary);
82 if (!fs.is_open()) {
83 MS_LOG(ERROR) << "Open binary proto file " << file << " failed.";
84 return RET_ERROR;
85 }
86
87 google::protobuf::io::IstreamInputStream istream(&fs);
88 google::protobuf::io::CodedInputStream coded_stream(&istream);
89
90 bool success = ReadProtoFromCodedInputStream(&coded_stream, message);
91 fs.close();
92
93 if (!success) {
94 MS_LOG(WARNING) << "Parse " << file << " failed.";
95 return RET_ERROR;
96 }
97
98 return RET_OK;
99 }
100 } // namespace lite
101 } // namespace mindspore
102