1 /**
2 * Copyright 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 #ifndef MINDSPORE_LITE_TOOLS_OBFUSCATOR_INCLUDE_DEOBFUSCATOR_H
18 #define MINDSPORE_LITE_TOOLS_OBFUSCATOR_INCLUDE_DEOBFUSCATOR_H
19
20 #include "src/common/common.h"
21 #include "include/model.h"
22
23 #define IV_SIZE 16
24
25 namespace mindspore::lite {
26 struct DeObfuscator {
27 Uint32Vector junk_tensor_Indices_;
28 Uint32Vector junk_node_Indices_;
29 Uint32Vector masking_values_;
30 using PrimTypeVector = Vector<schema::PrimitiveType>;
31 PrimTypeVector all_prims_type_;
32 unsigned char *obf_meta_data_;
33 uint32_t all_tensor_size_;
34 uint32_t all_node_size_;
35 bool with_sub_graph_;
36 void Free();
~DeObfuscatorDeObfuscator37 ~DeObfuscator() { Free(); }
38 };
39
40 int DeObfuscateModel(Model *model, DeObfuscator *model_deobf);
41 bool DeObfuscatePrimitive(const schema::Primitive *src, uint32_t src_node_stat, unsigned char **dst_prim,
42 schema::PrimitiveType dst_type);
43 bool InitModelDeObfuscator(Model *model, DeObfuscator *model_deobf, const flatbuffers::Vector<uint8_t> *meta_data,
44 const flatbuffers::Vector<uint8_t> *decrypt_table, size_t node_num);
45
46 template <typename T = schema::MetaGraph>
IsMetaGraphObfuscated(const T & meta_graph)47 bool IsMetaGraphObfuscated(const T &meta_graph) {
48 if (meta_graph.obfMetaData() == nullptr) {
49 MS_LOG(INFO) << "obfMetaData is null.";
50 return false;
51 }
52 return true;
53 }
54
55 template <typename T = schema::MetaGraph>
GetModelDeObfuscator(const T & meta_graph,Model * model)56 DeObfuscator *GetModelDeObfuscator(const T &meta_graph, Model *model) {
57 auto meta_data = meta_graph.obfMetaData();
58 auto *model_deobfuscator = new (std::nothrow) DeObfuscator();
59 if (model_deobfuscator == nullptr) {
60 MS_LOG(ERROR) << "new model deobfuscator fail!";
61 return nullptr;
62 }
63 if (!InitModelDeObfuscator(model, model_deobfuscator, meta_data, nullptr, meta_graph.nodes()->size())) {
64 MS_LOG(ERROR) << "init model deobfuscator fail!";
65 delete model_deobfuscator;
66 return nullptr;
67 }
68 model_deobfuscator->all_tensor_size_ = meta_graph.allTensors()->size();
69 model_deobfuscator->all_node_size_ = meta_graph.nodes()->size();
70 if (meta_graph.subGraph() == nullptr) {
71 model_deobfuscator->with_sub_graph_ = false;
72 } else {
73 model_deobfuscator->with_sub_graph_ = true;
74 }
75 return model_deobfuscator;
76 }
77 } // namespace mindspore::lite
78
79 #endif // MINDSPORE_LITE_TOOLS_OBFUSCATOR_INCLUDE_DEOBFUSCATOR_H
80