paddle_stream_decrypt.cpp 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. // Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #include <string.h>
  15. #include <iostream>
  16. #include <string>
  17. #include <fstream>
  18. #include <memory>
  19. #include <iterator>
  20. #include <algorithm>
  21. #include "encryption/include/paddle_stream_decrypt.h"
  22. #include "encryption/include/model_code.h"
  23. #include "encryption/util/include/crypto/aes_gcm.h"
  24. #include "encryption/util/include/io_utils.h"
  25. #include "encryption/util/include/log.h"
  26. #include "encryption/util/include/constant/constant_model.h"
  27. #include "encryption/util/include/system_utils.h"
  28. #include "encryption/util/include/crypto/base64.h"
  29. int paddle_check_stream_encrypted(std::istream &cipher_stream) {
  30. return util::SystemUtils::check_file_encrypted(cipher_stream);
  31. }
  32. int decrypt_stream(std::istream &cipher_stream,
  33. std::ostream &plain_stream,
  34. const std::string &key_base64) {
  35. int ret = paddle_check_stream_encrypted(cipher_stream);
  36. if (ret != CODE_OK) {
  37. LOGD("[M]check file encrypted failed, code: %d", ret);
  38. return ret;
  39. }
  40. std::string key_str =
  41. baidu::base::base64::base64_decode(key_base64.c_str());
  42. int ret_check = util::SystemUtils::check_key_match(key_str, cipher_stream);
  43. if (ret_check != CODE_OK) {
  44. LOGD("[M]check key failed in decrypt_file, code: %d", ret_check);
  45. return CODE_KEY_NOT_MATCH;
  46. }
  47. std::string aes_key = key_str.substr(0, AES_GCM_KEY_LENGTH);
  48. std::string aes_iv = key_str.substr(16, AES_GCM_IV_LENGTH);
  49. cipher_stream.seekg(0, std::ios::beg);
  50. cipher_stream.seekg(0, std::ios::end);
  51. int data_len = cipher_stream.tellg();
  52. cipher_stream.seekg(0, std::ios::beg);
  53. size_t pos = constant::MAGIC_NUMBER_LEN +
  54. constant::VERSION_LEN + constant::TAG_LEN;
  55. size_t cipher_len = data_len - pos;
  56. std::unique_ptr<unsigned char[]> model_cipher(
  57. new unsigned char[cipher_len]);
  58. cipher_stream.seekg(pos); // skip header
  59. cipher_stream.read(reinterpret_cast<char *>(model_cipher.get()),
  60. cipher_len);
  61. size_t plain_len = data_len - AES_GCM_TAG_LENGTH - pos;
  62. std::unique_ptr<unsigned char[]> model_plain(new unsigned char[plain_len]);
  63. int ret_decrypt_file = util::crypto::AesGcm::decrypt_aes_gcm(
  64. model_cipher.get(),
  65. cipher_len,
  66. reinterpret_cast<const unsigned char*>(aes_key.c_str()),
  67. reinterpret_cast<const unsigned char*>(aes_iv.c_str()),
  68. model_plain.get(),
  69. reinterpret_cast<int&>(plain_len));
  70. if (ret_decrypt_file != CODE_OK) {
  71. LOGD("[M]decrypt file failed, decrypt ret = %d", ret_decrypt_file);
  72. return ret_decrypt_file;
  73. }
  74. plain_stream.write(reinterpret_cast<const char*>(model_plain.get()),
  75. plain_len);
  76. return CODE_OK;
  77. }