encrypt

Author: grapefruitL

tensorflow/cc/saved_model/reader.cc

@@ -17,6 +17,9 @@ limitations under the License.
 
 #include <unordered_set>
 
+#include <iomanip>
+#include <sstream>
+
 #include "absl/memory/memory.h"
 #include "tensorflow/cc/saved_model/constants.h"
 #include "tensorflow/core/lib/io/path.h"

tensorflow/core/platform/default/BUILD

@@ -126,6 +126,7 @@ cc_library(
         "//third_party/eigen3",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:optional",
+        "@cryptopp",
     ],
 )
 

tensorflow/core/platform/env.cc

@@ -48,8 +48,36 @@ limitations under the License.
 #include <unistd.h>
 #endif
 
+#include "aes.h"
+#include "modes.h"
+#include "cryptlib.h"
+#include "filters.h"
+#include "secblock.h"
+
 namespace tensorflow {
 
+byte key[CryptoPP::AES::DEFAULT_KEYLENGTH] = {0x8b, 0x8b, 0x8b, 0x8b, 0x8b, 0x8b,
+                                    0x8b, 0x8b, 0x8b, 0x8b, 0x8b, 0x8b,
+                                    0x8b, 0x8b, 0x8b, 0x8b};
+byte iv[CryptoPP::AES::BLOCKSIZE] = {0x37, 0x37, 0x37, 0x37, 0x37, 0x37, 0x37, 0x37,
+                           0x37, 0x37, 0x37, 0x37, 0x37, 0x37, 0x37, 0x37};
+
+ Status decryptCBC(std::string cipher, byte key[], int keySize, byte iv[], std::string& recovered) {
+  // Decryption
+  try {
+    CryptoPP::CBC_Mode<CryptoPP::AES>::Decryption d;
+    d.SetKeyWithIV(key, keySize, iv);
+    // The StreamTransformationFilter removes
+    //  padding as required.
+    CryptoPP::StringSource s(cipher, true,
+                   new CryptoPP::StreamTransformationFilter(
+                       d, new CryptoPP::StringSink(recovered)));
+  } catch (const CryptoPP::Exception& e) {
+    return errors::DataLoss("Can't decrypt binary proto");
+  }
+  return Status::OK();
+}
+
 // 128KB copy buffer
 constexpr size_t kCopyFileBufferSize = 128 * 1024;
 
@@ -541,14 +569,23 @@ Status WriteBinaryProto(Env* env, const string& fname,
 
 Status ReadBinaryProto(Env* env, const string& fname,
                        protobuf::MessageLite* proto) {
-  std::unique_ptr<RandomAccessFile> file;
-  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file));
-  std::unique_ptr<FileStream> stream(new FileStream(file.get()));
+  // std::unique_ptr<RandomAccessFile> file;
+  // TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file));
+  // std::unique_ptr<FileStream> stream(new FileStream(file.get()));
+
+  using google::protobuf::io::ArrayInputStream;
+  string fileData;
+  TF_RETURN_IF_ERROR(ReadFileToString(env, fname, &fileData));
+  string plain;
+  TF_RETURN_IF_ERROR(decryptCBC(fileData, key, sizeof(key), iv, plain));
+  std::unique_ptr<ArrayInputStream> stream(
+  new ArrayInputStream(plain.data(), plain.length()));
+
   protobuf::io::CodedInputStream coded_stream(stream.get());
 
   if (!proto->ParseFromCodedStream(&coded_stream) ||
       !coded_stream.ConsumedEntireMessage()) {
-    TF_RETURN_IF_ERROR(stream->status());
+    // TF_RETURN_IF_ERROR(stream->status());
     return errors::DataLoss("Can't parse ", fname, " as binary proto");
   }
   return Status::OK();

tensorflow/core/platform/windows/BUILD

@@ -73,6 +73,7 @@ cc_library(
         "//third_party/eigen3",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:optional",
+        "@cryptopp",
     ],
 )
 

tensorflow/workspace.bzl

@@ -13,7 +13,7 @@ load("//third_party/toolchains/cpus/arm:arm_compiler_configure.bzl", "arm_compil
 load("//third_party/toolchains/embedded/arm-linux:arm_linux_toolchain_configure.bzl", "arm_linux_toolchain_configure")
 load("//third_party:repo.bzl", "tf_http_archive")
 load("//third_party/clang_toolchain:cc_configure_clang.bzl", "cc_download_clang_toolchain")
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_file")
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_file", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:java.bzl", "java_import_external")
 load("@io_bazel_rules_closure//closure:defs.bzl", "filegroup_external")
 load(
@@ -1126,6 +1126,19 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
             "https://github.com/GrahamDumpleton/wrapt/archive/1.11.1.tar.gz",
         ],
     )
+    
+    # add cryptopp
+    tf_http_archive(
+        name = "cryptopp",
+        build_file = clean_dep("//third_party/cryptopp:BUILD"),
+        sha256 = "c934d2c427a0ef197ea989a00f7b6d866d110dd55257d2944d0513b382b7e2b4",
+        strip_prefix = "cryptopp-CRYPTOPP_5_6_5",
+        urls = [
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/weidai11/cryptopp/archive/CRYPTOPP_5_6_5.zip",
+            "https://github.com/weidai11/cryptopp/archive/CRYPTOPP_5_6_5.zip",
+        ],
+    )
+
     tf_http_archive(
         name = "coremltools",
         sha256 = "0d594a714e8a5fd5bd740ad112ef59155c0482e25fdc8f8efa5758f90abdcf1e",

third_party/cryptopp/BUILD

@@ -0,0 +1,19 @@
+cc_library(
+    name = "cryptopp",
+    srcs = glob(["*.cpp"]),
+    hdrs = glob(["*.h"]) + [
+        "algebra.cpp",
+        "eprecomp.cpp",
+        "eccrypto.cpp",
+        "strciphr.cpp",
+        "polynomi.cpp",
+    ],
+    textual_hdrs = [
+        "algebra.cpp",
+        "eprecomp.cpp",
+        "eccrypto.cpp",
+        "strciphr.cpp",
+        "polynomi.cpp",
+    ],
+    visibility = ["//visibility:public"],
+)