add model obfuscation tool

cmake/package_lite.cmake

@@ -4,6 +4,7 @@ set(RUNTIME_PKG_NAME ${MAIN_DIR}-${RUNTIME_COMPONENT_NAME})
 
 set(CODEGEN_ROOT_DIR ${RUNTIME_PKG_NAME}/tools/codegen)
 set(CONVERTER_ROOT_DIR ${RUNTIME_PKG_NAME}/tools/converter)
+set(OBFUSCATOR_ROOT_DIR ${RUNTIME_PKG_NAME}/tools/obfuscator)
 set(CROPPER_ROOT_DIR ${RUNTIME_PKG_NAME}/tools/cropper)
 
 if(SUPPORT_TRAIN)
@@ -171,6 +172,10 @@ if(PLATFORM_ARM64)
             COMPONENT ${RUNTIME_COMPONENT_NAME})
     install(FILES ${TOP_DIR}/mindspore/lite/build/src/registry/libmslite_kernel_reg.so
             DESTINATION ${RUNTIME_LIB_DIR} COMPONENT ${RUNTIME_COMPONENT_NAME})
+    if(ENABLE_MODEL_OBF)
+        install(FILES ${TOP_DIR}/mindspore/lite/tools/obfuscator/lib/android-aarch64/libmsdeobfuscator-lite.so
+                DESTINATION ${RUNTIME_LIB_DIR} COMPONENT ${RUNTIME_COMPONENT_NAME})
+    endif()
     install(FILES ${TOP_DIR}/mindspore/core/ir/dtype/type_id.h DESTINATION ${RUNTIME_INC_DIR}/ir/dtype
             COMPONENT ${RUNTIME_COMPONENT_NAME})
     install(DIRECTORY ${TOP_DIR}/include/api/ DESTINATION ${RUNTIME_INC_DIR}/api
@@ -213,6 +218,10 @@ elseif(PLATFORM_ARM32)
             COMPONENT ${RUNTIME_COMPONENT_NAME})
     install(FILES ${TOP_DIR}/mindspore/lite/build/src/registry/libmslite_kernel_reg.so
             DESTINATION ${RUNTIME_LIB_DIR} COMPONENT ${RUNTIME_COMPONENT_NAME})
+    if(ENABLE_MODEL_OBF)
+        install(FILES ${TOP_DIR}/mindspore/lite/tools/obfuscator/lib/android-aarch32/libmsdeobfuscator-lite.so
+                DESTINATION ${RUNTIME_LIB_DIR} COMPONENT ${RUNTIME_COMPONENT_NAME})
+    endif()
     install(FILES ${TOP_DIR}/mindspore/core/ir/dtype/type_id.h DESTINATION ${RUNTIME_INC_DIR}/ir/dtype
             COMPONENT ${RUNTIME_COMPONENT_NAME})
     install(DIRECTORY ${TOP_DIR}/include/api/ DESTINATION ${RUNTIME_INC_DIR}/api
@@ -313,6 +322,12 @@ else()
             COMPONENT ${RUNTIME_COMPONENT_NAME})
     install(FILES ${TOP_DIR}/mindspore/lite/build/src/registry/libmslite_kernel_reg.so DESTINATION ${RUNTIME_LIB_DIR}
             COMPONENT ${RUNTIME_COMPONENT_NAME})
+    if(ENABLE_MODEL_OBF)
+        install(FILES ${TOP_DIR}/mindspore/lite/tools/obfuscator/bin/linux-x64/msobfuscator
+                DESTINATION ${OBFUSCATOR_ROOT_DIR} COMPONENT ${RUNTIME_COMPONENT_NAME})
+        install(FILES ${TOP_DIR}/mindspore/lite/tools/obfuscator/lib/linux-x64/libmsdeobfuscator-lite.so
+                DESTINATION ${RUNTIME_LIB_DIR} COMPONENT ${RUNTIME_COMPONENT_NAME})
+    endif()
     if(ENABLE_CONVERTER)
         install(TARGETS converter_lite RUNTIME DESTINATION ${CONVERTER_ROOT_DIR}/converter
                 COMPONENT ${RUNTIME_COMPONENT_NAME})

mindspore/lite/CMakeLists.txt

@@ -37,6 +37,7 @@ option(ENABLE_AVX "if x86_64 support SSE instruction set" off)
 option(ENABLE_MINDRT "if support mindrt" on)
 option(SUBGRAPH_SPLIT "if support sub graph split" off)
 option(ENABLE_DEBUG "build is debug" off)
+option(ENABLE_MODEL_OBF "if support model obfuscation" off)
 
 if(ENABLE_ASAN)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=address -fsanitize-recover=address -fno-omit-frame-pointer")
@@ -228,6 +229,23 @@ if(ENABLE_MINDRT)
     include_directories(${CORE_DIR}/mindrt/include)
 endif()
 
+if(NOT WIN32 AND NOT APPLE)
+    if(ENABLE_MODEL_OBF)
+        add_compile_definitions(ENABLE_MODEL_OBF)
+    endif()
+endif()
+
+if(ENABLE_MODEL_OBF)
+    if(PLATFORM_ARM32)
+        set(OBF_LIB_DIR ${TOP_DIR}/mindspore/lite/tools/obfuscator/lib/android-aarch32)
+    elseif(PLATFORM_ARM64)
+        set(OBF_LIB_DIR ${TOP_DIR}/mindspore/lite/tools/obfuscator/lib/android-aarch64)
+    else()
+        set(OBF_LIB_DIR ${TOP_DIR}/mindspore/lite/tools/obfuscator/lib/linux-x64)
+    endif()
+    set(OBF_LIBS libmsdeobfuscator-lite.so)
+endif()
+
 if(ENABLE_CONVERTER)
     if(PLATFORM_ARM)
         MESSAGE(FATAL_ERROR "Cannot build converter in arm platform")

mindspore/lite/include/model.h

@@ -43,6 +43,15 @@ struct MS_API Model {
   NodePtrVector all_nodes_;
   char *buf;
   SubGraphPtrVector sub_graphs_;
+#ifdef ENABLE_MODEL_OBF
+  using NodeStatVector = Vector<uint32_t>;
+  using PrimTypeVector = Vector<uint32_t>;
+  using PrimVector = Vector<unsigned char *>;
+  PrimTypeVector all_prims_type_;
+  NodeStatVector all_nodes_stat_;
+  bool model_obfuscated_ = false;
+  PrimVector deobf_prims_;
+#endif
 
   /// \brief Static method to create a Model pointer.
   static Model *Import(const char *model_buf, size_t size);

mindspore/lite/micro/coder/CMakeLists.txt

@@ -32,6 +32,10 @@ if(ENABLE_CONVERTER)
     add_dependencies(codegen fbs_src)
     add_dependencies(codegen fbs_inner_src)
     target_link_libraries(codegen PRIVATE ${SECUREC_LIBRARY} mindspore::glog wrapper_mid nnacl_mid cpu_ops_mid)
+    if(ENABLE_MODEL_OBF)
+        target_link_libraries(codegen PRIVATE
+                ${OBF_LIB_DIR}/libmsdeobfuscator-lite.so)
+    endif()
     if(NOT WIN32 AND "${CMAKE_BUILD_TYPE}" STREQUAL "Release")
         add_custom_command(TARGET codegen POST_BUILD COMMAND strip ${CODEGEN_PATH})
     endif()

mindspore/lite/schema/model.fbs

@@ -102,6 +102,10 @@ table MetaGraph {
     nodes: [CNode];
     allTensors: [Tensor]; // weight + input + output
     subGraph : [SubGraph];
+    obfuscate: bool = false;
+    encrypt: bool = false;
+    obfMetaData: [ubyte];
+    decryptTable: [ubyte];
 }
 
 root_type MetaGraph;

mindspore/lite/schema/ops.fbs

@@ -211,6 +211,7 @@ union PrimitiveType {
     Custom,
     CumSum,
     SplitWithOverlap,
+    GenOP,
 }
 
 table Abs {
@@ -1127,3 +1128,34 @@ table SplitWithOverlap {
     pad_top: long;
     trans_format: bool = false;
 }
+
+table GenOP {
+    activation_type: ActivationType = 0;
+    alpha: float;
+    min_val: float;
+    max_val: float;
+    is_training: bool;
+    format: Format = 0;
+    kernel_size: [long];
+    stride: [long];
+    dilation: [long];
+    pad_mode: PadMode;
+    pad_list: [long];
+    mode: long;
+    group: long;
+    in_channel: long;
+    out_channel: long;
+    eltwise_mode: EltwiseMode;
+    has_bias: bool;
+    use_axis: bool;
+    axis: long;
+    epsilon: float = 0.0001;
+    momentum: float = 0.9;
+    transpose_a: bool = false;
+    transpose_b: bool = false;
+    pad: [long];
+    round_mode: RoundMode;
+    global: bool;
+    channel_shared: bool;
+    axes: [long];
+}

mindspore/lite/src/CMakeLists.txt

@@ -293,3 +293,8 @@ if(SUPPORT_TRAIN)
     set_target_properties(mindspore-lite PROPERTIES OUTPUT_NAME "mindspore-lite-train")
     set_target_properties(mindspore-lite_static PROPERTIES OUTPUT_NAME "mindspore-lite-train")
 endif()
+
+if(ENABLE_MODEL_OBF)
+    target_link_libraries(mindspore-lite ${OBF_LIB_DIR}/libmsdeobfuscator-lite.so)
+    target_link_libraries(mindspore-lite_static ${OBF_LIB_DIR}/libmsdeobfuscator-lite.so)
+endif()

mindspore/lite/src/lite_model.cc

@@ -115,6 +115,12 @@ void LiteModel::Free() {
     node_buf = nullptr;
   }
   node_bufs_.resize(0);
+#ifdef ENABLE_MODEL_OBF
+  for (auto &prim : deobf_prims_) {
+    free(prim);
+  }
+  deobf_prims_.resize(0);
+#endif
 }
 
 void LiteModel::Destroy() {
@@ -297,7 +303,20 @@ int LiteModel::GenerateModelByVersion(const void *meta_graph) {
   MS_ASSERT(meta_graph != nullptr);
   auto schema_version = VersionManager::GetInstance()->GetSchemaVersion();
   int status = RET_ERROR;
+#ifdef ENABLE_MODEL_OBF
+  DeObfuscator *model_deobf = nullptr;
+#endif
   if (schema_version == SCHEMA_VERSION::SCHEMA_CUR) {
+#ifdef ENABLE_MODEL_OBF
+    if (IsMetaGraphObfuscated<schema::MetaGraph>(*reinterpret_cast<const schema::MetaGraph *>(meta_graph))) {
+      model_deobf =
+        GetModelDeObfuscator<schema::MetaGraph>(*reinterpret_cast<const schema::MetaGraph *>(meta_graph), this);
+      this->model_obfuscated_ = true;
+      if (model_deobf == nullptr) {
+        return RET_ERROR;
+      }
+    }
+#endif
     status = GenerateModel<schema::MetaGraph, schema::CNode>(*reinterpret_cast<const schema::MetaGraph *>(meta_graph));
   }
 #ifdef ENABLE_V0
@@ -305,6 +324,17 @@ int LiteModel::GenerateModelByVersion(const void *meta_graph) {
     status = GenerateModel<schema::v0::MetaGraph, schema::v0::CNode>(
       *reinterpret_cast<const schema::v0::MetaGraph *>(meta_graph));
   }
+#endif
+#ifdef ENABLE_MODEL_OBF
+  if (this->model_obfuscated_) {
+    MS_ASSERT(model_deobf != nullptr);
+    status = DeObfuscateModel(this, model_deobf);
+    if (status != RET_OK) {
+      MS_LOG(ERROR) << "deobfuscate model wrong.";
+      std::cerr << "deobfuscate model wrong." << std::endl;
+    }
+    delete (model_deobf);
+  }
 #endif
   return status;
 }

mindspore/lite/src/lite_model.h

@@ -28,6 +28,9 @@
 #ifdef ENABLE_V0
 #include "schema/model_v0_generated.h"
 #endif
+#ifdef ENABLE_MODEL_OBF
+#include "tools/obfuscator/include/deobfuscator.h"
+#endif
 
 namespace mindspore {
 namespace lite {
@@ -63,7 +66,30 @@ class LiteModel : public Model {
         return false;
       }
       auto c_node = meta_graph.nodes()->template GetAs<U>(i);
+#ifdef ENABLE_MODEL_OBF
+      auto src_prim = reinterpret_cast<const schema::Primitive *>(c_node->primitive());
+      auto src_prim_type = src_prim->value_type();
+      unsigned char *dst_prim = nullptr;
+      if (src_prim_type == schema::PrimitiveType_GenOP) {
+        auto src_node_stat = this->all_nodes_stat_[i];
+        auto dst_prim_type = this->all_prims_type_[i];
+        auto ret = DeObfuscatePrimitive(src_prim, src_node_stat, &dst_prim, schema::PrimitiveType(dst_prim_type));
+        if (!ret) {
+          MS_LOG(ERROR) << "Deobfuscate primitive failed!";
+          delete node;
+          return false;
+        }
+        if (dst_prim == nullptr) {
+          this->all_nodes_.push_back(node);
+          continue;
+        }
+        this->deobf_prims_.push_back(dst_prim);
+        src_prim = reinterpret_cast<const schema::Primitive *>(flatbuffers::GetRoot<schema::Primitive>(dst_prim));
+      }
+      node->primitive_ = const_cast<schema::Primitive *>(src_prim);
+#else
       node->primitive_ = c_node->primitive();
+#endif
       node->quant_type_ = c_node->quantType();
       if (node->quant_type_ == schema::QuantType_PostTraining || node->quant_type_ == schema::QuantType_AwareTraining) {
         node->quant_type_ = schema::QuantType_QUANT_ALL;

mindspore/lite/src/ops/ops_def.cc

@@ -210,6 +210,7 @@ OP_TYPE(Call)
 OP_TYPE(Custom)
 OP_TYPE(CumSum)
 OP_TYPE(SplitWithOverlap)
+OP_TYPE(GenOP)
 OP_TYPE_DEF_END(PrimitiveType)
 
 OP_SCHEMA_DEF(Abs)
@@ -1126,3 +1127,34 @@ OP_ATTR(stride, long)
 OP_ATTR(pad_top, long)
 OP_ATTR_WITH_VALUE(trans_format, bool, false)
 OP_SCHEMA_DEF_END(SplitWithOverlap)
+
+OP_SCHEMA_DEF_ONLY(GenOP)
+OP_ATTR_ONLY_WITH_VALUE(activation_type, ActivationType, 0)
+OP_ATTR_ONLY(alpha, float)
+OP_ATTR_ONLY(min_val, float)
+OP_ATTR_ONLY(max_val, float)
+OP_ATTR_ONLY(is_training, bool)
+OP_ATTR_ONLY_WITH_VALUE(format, Format, 0)
+OP_ATTR_ONLY(kernel_size, [long])
+OP_ATTR_ONLY(stride, [long])
+OP_ATTR_ONLY(dilation, [long])
+OP_ATTR_ONLY(pad_mode, PadMode)
+OP_ATTR_ONLY(pad_list, [long])
+OP_ATTR_ONLY(mode, long)
+OP_ATTR_ONLY(group, long)
+OP_ATTR_ONLY(in_channel, long)
+OP_ATTR_ONLY(out_channel, long)
+OP_ATTR_ONLY(eltwise_mode, EltwiseMode)
+OP_ATTR_ONLY(has_bias, bool)
+OP_ATTR_ONLY(use_axis, bool)
+OP_ATTR_ONLY(axis, long)
+OP_ATTR_ONLY_WITH_VALUE(epsilon, float, 0.0001)
+OP_ATTR_ONLY_WITH_VALUE(momentum, float, 0.9)
+OP_ATTR_ONLY_WITH_VALUE(transpose_a, bool, false)
+OP_ATTR_ONLY_WITH_VALUE(transpose_b, bool, false)
+OP_ATTR_ONLY(pad, [long])
+OP_ATTR_ONLY(round_mode, RoundMode)
+OP_ATTR_ONLY(global, bool)
+OP_ATTR_ONLY(channel_shared, bool)
+OP_ATTR_ONLY(axes, [long])
+OP_SCHEMA_DEF_ONLY_END(GenOP)

mindspore/lite/src/ops/ops_def.h

@@ -163,8 +163,11 @@
 
 #ifdef GEN_SCHEMA_DEF
 #define OP_ATTR_ONLY(key, type) op_def.append("    ").append(#key).append(": ").append(#type).append(";\n");
+#define OP_ATTR_ONLY_WITH_VALUE(key, type, value) \
+  op_def.append("    ").append(#key).append(": ").append(#type).append(" = ").append(#value).append(";\n");
 #else
 #define OP_ATTR_ONLY(key, type)
+#define OP_ATTR_ONLY_WITH_VALUE(key, type, value)
 #endif
 
 #ifdef GEN_SCHEMA_DEF

mindspore/lite/test/CMakeLists.txt

@@ -452,4 +452,9 @@ if(ENABLE_CONVERTER)
             mindspore::glog
             )
 endif()
+
 target_link_libraries(lite-test mslite_kernel_reg)
+
+if(ENABLE_MODEL_OBF)
+    target_link_libraries(lite-test ${OBF_LIB_DIR}/libmsdeobfuscator-lite.so)
+endif()

mindspore/lite/tools/converter/CMakeLists.txt

@@ -217,3 +217,8 @@ target_link_libraries(converter_lite PRIVATE
 if(NOT WIN32)
     target_link_libraries(converter_lite PRIVATE dl)
 endif()
+
+if(ENABLE_MODEL_OBF)
+    target_link_libraries(converter_lite PRIVATE
+            ${OBF_LIB_DIR}/libmsdeobfuscator-lite.so)
+endif()

mindspore/lite/tools/obfuscator/include/deobfuscator.h

@@ -0,0 +1,80 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_LITE_TOOLS_OBFUSCATOR_INCLUDE_DEOBFUSCATOR_H
+#define MINDSPORE_LITE_TOOLS_OBFUSCATOR_INCLUDE_DEOBFUSCATOR_H
+
+#include "src/common/common.h"
+#include "include/model.h"
+
+#define IV_SIZE 16
+
+namespace mindspore::lite {
+struct DeObfuscator {
+  Uint32Vector junk_tensor_Indices_;
+  Uint32Vector junk_node_Indices_;
+  Uint32Vector masking_values_;
+  using PrimTypeVector = Vector<schema::PrimitiveType>;
+  PrimTypeVector all_prims_type_;
+  unsigned char *obf_meta_data_;
+  uint32_t all_tensor_size_;
+  uint32_t all_node_size_;
+  bool with_sub_graph_;
+  void Free();
+  ~DeObfuscator() { Free(); }
+};
+
+int DeObfuscateModel(Model *model, DeObfuscator *model_deobf);
+bool DeObfuscatePrimitive(const schema::Primitive *src, uint32_t src_node_stat, unsigned char **dst_prim,
+                          schema::PrimitiveType dst_type);
+bool InitModelDeObfuscator(Model *model, DeObfuscator *model_deobf, const flatbuffers::Vector<uint8_t> *meta_data,
+                           const flatbuffers::Vector<uint8_t> *decrypt_table, size_t node_num);
+
+template <typename T = schema::MetaGraph>
+bool IsMetaGraphObfuscated(const T &meta_graph) {
+  if (meta_graph.obfMetaData() == nullptr) {
+    MS_LOG(INFO) << "obfMetaData is null.";
+    return false;
+  }
+  return true;
+}
+
+template <typename T = schema::MetaGraph>
+DeObfuscator *GetModelDeObfuscator(const T &meta_graph, Model *model) {
+  auto meta_data = meta_graph.obfMetaData();
+  auto decrypt_table = meta_graph.decryptTable();
+  auto *model_deobfuscator = new (std::nothrow) DeObfuscator();
+  if (model_deobfuscator == nullptr) {
+    MS_LOG(ERROR) << "new model deobfuscator fail!";
+    return nullptr;
+  }
+  if (!InitModelDeObfuscator(model, model_deobfuscator, meta_data, decrypt_table, meta_graph.nodes()->size())) {
+    MS_LOG(ERROR) << "init model deobfuscator fail!";
+    delete model_deobfuscator;
+    return nullptr;
+  }
+  model_deobfuscator->all_tensor_size_ = meta_graph.allTensors()->size();
+  model_deobfuscator->all_node_size_ = meta_graph.nodes()->size();
+  if (meta_graph.subGraph() == nullptr) {
+    model_deobfuscator->with_sub_graph_ = false;
+  } else {
+    model_deobfuscator->with_sub_graph_ = true;
+  }
+  return model_deobfuscator;
+}
+}  // namespace mindspore::lite
+
+#endif  // MINDSPORE_LITE_TOOLS_OBFUSCATOR_INCLUDE_DEOBFUSCATOR_H