SecureStorage encryption functions

GitOrigin-RevId: 8df89a0fb672d83f9f743aac1bd779dd0635c70e

CMakeLists.txt

@@ -400,6 +400,7 @@ set(TDLIB_SOURCE
   td/telegram/SecretChatActor.cpp
   td/telegram/SecretChatDb.cpp
   td/telegram/SecretChatsManager.cpp
+  td/telegram/SecureStorage.cpp
   td/telegram/SequenceDispatcher.cpp
   td/telegram/StateManager.cpp
   td/telegram/StickersManager.cpp
@@ -519,6 +520,7 @@ set(TDLIB_SOURCE
   td/telegram/SecretChatDb.h
   td/telegram/SecretChatsManager.h
   td/telegram/SecretInputMedia.h
+  td/telegram/SecureStorage.h
   td/telegram/SequenceDispatcher.h
   td/telegram/StateManager.h
   td/telegram/StickersManager.h

td/telegram/SecureStorage.cpp

@@ -0,0 +1,316 @@
+//
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2018
+//
+// Distributed under the Boost Software License, Version 1.0. (See accompanying
+// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
+//
+#include "td/telegram/SecureStorage.h"
+#include "td/utils/misc.h"
+namespace td {
+namespace secure_storage {
+// Helpers
+AesCbcState calc_aes_cbc_state(Slice seed) {
+  UInt<512> hash;
+  auto hash_slice = as_slice(hash);
+  sha512(seed, hash_slice);
+  UInt256 key;
+  as_slice(key).copy_from(hash_slice.substr(0, 32));
+  UInt128 iv;
+  as_slice(iv).copy_from(hash_slice.substr(32, 16));
+  return AesCbcState{key, iv};
+}
+
+template <class F>
+Status data_view_for_each(DataView &data, F &&f) {
+  const int64 step = 128 << 10;
+  for (int64 i = 0, size = data.size(); i < size; i += step) {
+    TRY_RESULT(bytes, data.pread(i, std::min(step, size - i)));
+    TRY_STATUS(f(std::move(bytes)));
+  }
+  return Status::OK();
+}
+
+Result<ValueHash> calc_value_hash(DataView &data_view) {
+  Sha256State state;
+  sha256_init(&state);
+  data_view_for_each(data_view, [&state](BufferSlice bytes) {
+    sha256_update(bytes.as_slice(), &state);
+    return Status::OK();
+  });
+  UInt256 res;
+  sha256_final(&state, as_slice(res));
+  return ValueHash{res};
+}
+
+BufferSlice gen_random_prefix(int64 data_size) {
+  BufferSlice buff(((32 + 15 + data_size) & -16) - data_size);
+  Random::secure_bytes(buff.as_slice());
+  buff.as_slice()[0] = narrow_cast<uint8>(buff.size());
+  CHECK((buff.size() + data_size) % 16 == 0);
+  return buff;
+}
+
+FileDataView::FileDataView(FileFd &fd, int64 size) : fd_(fd), size_(size) {
+}
+
+int64 FileDataView::size() {
+  return size_;
+}
+
+Result<BufferSlice> FileDataView::pread(int64 offset, int64 size) {
+  auto slice = BufferSlice(size);
+  TRY_RESULT(actual_size, fd_.pread(slice.as_slice(), offset));
+  if (static_cast<int64>(actual_size) != size) {
+    return Status::Error("Not enough data in file");
+  }
+  return std::move(slice);
+}
+
+BufferSliceDataView::BufferSliceDataView(BufferSlice buffer_slice) : buffer_slice_(std::move(buffer_slice)) {
+}
+int64 BufferSliceDataView::size() {
+  return buffer_slice_.size();
+}
+Result<BufferSlice> BufferSliceDataView::pread(int64 offset, int64 size) {
+  auto end_offset = size + offset;
+  if (this->size() < end_offset) {
+    return Status::Error("Not enough data in BufferSlice");
+  }
+  return BufferSlice(buffer_slice_.as_slice().substr(offset, size));
+}
+
+ConcatDataView::ConcatDataView(DataView &left, DataView &right) : left_(left), right_(right) {
+}
+int64 ConcatDataView::size() {
+  return left_.size() + right_.size();
+}
+Result<BufferSlice> ConcatDataView::pread(int64 offset, int64 size) {
+  auto end_offset = size + offset;
+  if (this->size() < end_offset) {
+    return Status::Error("Not enough data in ConcatDataView");
+  }
+
+  auto substr = [](DataView &slice, int64 offset, int64 size) -> Result<BufferSlice> {
+    auto l = std::max(int64{0}, offset);
+    auto r = std::min(slice.size(), offset + size);
+    if (l >= r) {
+      return BufferSlice();
+    }
+    return slice.pread(l, r - l);
+  };
+
+  TRY_RESULT(a, substr(left_, offset, size));
+  TRY_RESULT(b, substr(right_, offset - left_.size(), size));
+
+  if (a.empty()) {
+    return std::move(b);
+  }
+  if (b.empty()) {
+    return std::move(a);
+  }
+
+  BufferSlice res(a.size() + b.size());
+  res.as_slice().copy_from(a.as_slice());
+  res.as_slice().substr(a.size()).copy_from(b.as_slice());
+  return std::move(res);
+}
+
+// Password
+Password::Password(std::string password) : password_(std::move(password)) {
+}
+Slice Password::as_slice() const {
+  return password_;
+}
+
+// Secret
+namespace {
+uint8 secret_checksum(Slice secret) {
+  uint8 sum = 0;
+  for (uint8 c : secret) {
+    sum += c;
+  }
+  return static_cast<uint8>(239 - sum);
+}
+}  // namespace
+
+Result<Secret> Secret::create(Slice secret) {
+  if (secret.size() != 32) {
+    return Status::Error("wrong secret size");
+  }
+  if (secret_checksum(secret) != 0) {
+    return Status::Error("Wrong cheksum");
+  }
+  UInt256 res;
+  td::as_slice(res).copy_from(secret);
+  return Secret{res};
+}
+
+Secret Secret::create_new() {
+  UInt256 secret;
+  auto secret_slice = td::as_slice(secret);
+  Random::secure_bytes(secret_slice);
+  auto checksum_diff = secret_checksum(secret_slice);
+  secret_slice.ubegin()[0] += checksum_diff;
+  return create(secret_slice).move_as_ok();
+}
+
+Slice Secret::as_slice() const {
+  using td::as_slice;
+  return as_slice(secret_);
+}
+
+EncryptedSecret Secret::encrypt(Slice key) {
+  auto aes_cbc_state = calc_aes_cbc_state(key);
+  UInt256 res;
+  aes_cbc_state.encrypt(as_slice(), td::as_slice(res));
+  return EncryptedSecret::create(td::as_slice(res)).move_as_ok();
+}
+
+Secret::Secret(UInt256 secret) : secret_(secret) {
+}
+
+//EncryptedSecret
+Result<EncryptedSecret> EncryptedSecret::create(Slice encrypted_secret) {
+  if (encrypted_secret.size() != 32) {
+    return Status::Error("Wrong encrypte secret size");
+  }
+  UInt256 res;
+  td::as_slice(res).copy_from(encrypted_secret);
+  return EncryptedSecret{res};
+}
+Result<Secret> EncryptedSecret::decrypt(Slice key) {
+  auto aes_cbc_state = calc_aes_cbc_state(key);
+  UInt256 res;
+  aes_cbc_state.decrypt(td::as_slice(encrypted_secret_), td::as_slice(res));
+  return Secret::create(td::as_slice(res));
+}
+
+Slice EncryptedSecret::as_slice() const {
+  return td::as_slice(encrypted_secret_);
+}
+
+EncryptedSecret::EncryptedSecret(UInt256 encrypted_secret) : encrypted_secret_(encrypted_secret) {
+}
+
+// Decryption
+Decryptor::Decryptor(AesCbcState aes_cbc_state) : aes_cbc_state_(std::move(aes_cbc_state)) {
+  sha256_init(&sha256_state_);
+}
+Result<BufferSlice> Decryptor::append(BufferSlice data) {
+  if (data.empty()) {
+    return BufferSlice();
+  }
+  if (data.size() % 16 != 0) {
+    return Status::Error("Part size should be divisible by 16");
+  }
+  aes_cbc_state_.decrypt(data.as_slice(), data.as_slice());
+  sha256_update(data.as_slice(), &sha256_state_);
+  if (!skipped_prefix_) {
+    to_skip_ = data.as_slice().ubegin()[0];
+    size_t to_skip = std::min(to_skip_, data.size());
+    skipped_prefix_ = true;
+    data = data.from_slice(data.as_slice().remove_prefix(to_skip));
+  }
+  return std::move(data);
+}
+Result<ValueHash> Decryptor::finish() {
+  if (!skipped_prefix_) {
+    return Status::Error("No data was given");
+  }
+  if (to_skip_ < 32) {
+    return Status::Error("Too small random prefix");
+  }
+  UInt256 res;
+  sha256_final(&sha256_state_, as_slice(res));
+  return ValueHash{res};
+}
+
+// Encryptor
+Encryptor::Encryptor(AesCbcState aes_cbc_state, DataView &data_view)
+    : aes_cbc_state_(std::move(aes_cbc_state)), data_view_(data_view) {
+}
+
+int64 Encryptor::size() {
+  return data_view_.size();
+}
+
+Result<BufferSlice> Encryptor::pread(int64 offset, int64 size) {
+  if (offset != current_offset_) {
+    return Status::Error("Arbitrary offset is not supported");
+  }
+  if (size % 16 != 0) {
+    return Status::Error("Part size should be divisible by 16");
+  }
+  TRY_RESULT(part, data_view_.pread(offset, size));
+  aes_cbc_state_.encrypt(part.as_slice(), part.as_slice());
+  current_offset_ += size;
+  return std::move(part);
+}
+
+Result<EncryptedValue> encrypt_value(const Secret &secret, Slice data) {
+  auto random_prefix_view = BufferSliceDataView(gen_random_prefix(data.size()));
+  auto data_view = BufferSliceDataView(BufferSlice(data));
+  auto full_view = ConcatDataView(random_prefix_view, data_view);
+
+  TRY_RESULT(hash, calc_value_hash(full_view));
+
+  auto aes_cbc_state = calc_aes_cbc_state(PSLICE() << secret.as_slice() << hash.as_slice());
+  Encryptor encryptor(aes_cbc_state, full_view);
+  TRY_RESULT(encrypted_data, encryptor.pread(0, encryptor.size()));
+  return EncryptedValue{std::move(encrypted_data), std::move(hash)};
+}
+
+Result<BufferSlice> decrypt_value(const Secret &secret, const ValueHash &hash, Slice data) {
+  auto aes_cbc_state = calc_aes_cbc_state(PSLICE() << secret.as_slice() << hash.as_slice());
+  Decryptor decryptor(aes_cbc_state);
+  TRY_RESULT(decrypted_value, decryptor.append(BufferSlice(data)));
+  TRY_RESULT(got_hash, decryptor.finish());
+  if (got_hash.as_slice() != hash.as_slice()) {
+    return Status::Error("Hash mismatch");
+  }
+  return std::move(decrypted_value);
+}
+
+Result<ValueHash> encrypt_file(const Secret &secret, std::string src, std::string dest) {
+  TRY_RESULT(src_file, FileFd::open(src, FileFd::Flags::Read));
+  TRY_RESULT(dest_file, FileFd::open(dest, FileFd::Flags::Truncate | FileFd::Flags::Write | FileFd::Create));
+  auto src_file_size = src_file.get_size();
+
+  auto random_prefix_view = BufferSliceDataView(gen_random_prefix(src_file_size));
+  auto data_view = FileDataView(src_file, src_file_size);
+  auto full_view = ConcatDataView(random_prefix_view, data_view);
+
+  TRY_RESULT(hash, calc_value_hash(full_view));
+
+  auto aes_cbc_state = calc_aes_cbc_state(PSLICE() << secret.as_slice() << hash.as_slice());
+  Encryptor encryptor(aes_cbc_state, full_view);
+  TRY_STATUS(
+      data_view_for_each(encryptor, [&dest_file](BufferSlice bytes) { return dest_file.write(bytes.as_slice()); }));
+  return std::move(hash);
+}
+
+Status decrypt_file(const Secret &secret, const ValueHash &hash, std::string src, std::string dest) {
+  TRY_RESULT(src_file, FileFd::open(src, FileFd::Flags::Read));
+  TRY_RESULT(dest_file, FileFd::open(dest, FileFd::Flags::Truncate | FileFd::Flags::Write | FileFd::Create));
+  auto src_file_size = src_file.get_size();
+
+  auto src_file_view = FileDataView(src_file, src_file_size);
+
+  auto aes_cbc_state = calc_aes_cbc_state(PSLICE() << secret.as_slice() << hash.as_slice());
+  Decryptor decryptor(aes_cbc_state);
+  TRY_STATUS(data_view_for_each(src_file_view, [&decryptor, &dest_file](BufferSlice bytes) {
+    TRY_RESULT(decrypted_bytes, decryptor.append(std::move(bytes)));
+    TRY_STATUS(dest_file.write(decrypted_bytes.as_slice()));
+    return Status::OK();
+  }));
+
+  TRY_RESULT(got_hash, decryptor.finish());
+
+  if (hash.as_slice() != got_hash.as_slice()) {
+    return Status::Error("Hash mismatch");
+  }
+
+  return Status::OK();
+}
+}  // namespace secure_storage
+}  // namespace td

td/telegram/SecureStorage.h

@@ -0,0 +1,204 @@
+//
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2018
+//
+// Distributed under the Boost Software License, Version 1.0. (See accompanying
+// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
+//
+#include "td/utils/Random.h"
+#include "td/utils/buffer.h"
+#include "td/utils/crypto.h"
+#include "td/utils/format.h"
+#include "td/utils/int_types.h"
+#include "td/utils/Status.h"
+#include "td/utils/optional.h"
+#include "td/utils/port/FileFd.h"
+
+namespace td {
+// secureValueSendingStatePending = SecureValueSendingState;
+// secureValueSendingStateFailed = SecureValueSendingState;
+//
+// secureValue key:string sending_state:SecureValueSendingState unencrypted_text:string data:string files:vector<File> = SecureValue;
+//
+// inputSecureValue key:string unencrypted_text:string data:string files:vector<InputFile> = InputSecureValue;
+//
+// updateSecureValue value:SecureValue = Update;
+//
+// setSecureValue password:string value:InputSecureValue = SecureValue;
+// getSecureValue password:string key:string = SecureValue;
+//
+
+// Types
+// Password
+// Secret - 32 bytes with sum == 239
+// EncryptedSecret - encrypted secret
+// ValueHash - 32 bytes, sha256 from value
+//
+// ValueFull = ValueText? ValueData? ValueFile* = [Value]
+// Value = ValueText | ValueData | ValueFile
+//
+// ValueMeta = random_prefix, secret, hash
+//
+// Helpers
+//   calc_aes_cbc_state :: ValueSecret -> ValueHash -> AesCbcState
+//
+// Encryption.
+// To encrypt data:
+//   RandomPrefix, ValueSecret, Value:
+//   calc_value_hash :: RandomPrefix -> Value -> ValueHash
+//   do_encrypt :: RandomPrefix -> Value -> AesCbcState -> EncryptedValue // async
+//   encrypt :: (ValueSecret, RandomPrefix, Value) -> (EncryptedValue, ValueHash)
+//
+//
+// To decrypt data:
+//   ValueSecret, ValueHash, EncryptedValue
+//   do_decrypt :: EncryptedValue -> AesCbcState -> (RandomPrefix, Value, ValueHash) // async
+//   decrypt :: (ValueSecret, ValueHash, EncryptedValue) -> Value
+//
+// To encrypt FullValue:
+//   ValueSecret, [(RandomPrefix, Value)]
+//   (ValueSecret, [(RandomPrefix, Value)]) -> [(ValueSecret, RandomPrefix, Value)]
+//   [(ValueSecret, RandomPrefix, Value)] -> [(EncryptedValue, ValueHash)]
+//
+
+namespace secure_storage {
+// Helpers
+class ValueHash {
+ public:
+  ValueHash(UInt256 hash) : hash_(hash) {
+  }
+  Slice as_slice() const {
+    return td::as_slice(hash_);
+  }
+
+ private:
+  UInt256 hash_;
+};
+
+class DataView {
+ public:
+  virtual int64 size() = 0;
+  virtual Result<BufferSlice> pread(int64 offset, int64 size) = 0;
+  virtual ~DataView() = default;
+};
+
+class FileDataView : public DataView {
+ public:
+  FileDataView(FileFd &fd, int64 size);
+
+  int64 size() override;
+  Result<BufferSlice> pread(int64 offset, int64 size) override;
+
+ private:
+  FileFd &fd_;
+  int64 size_;
+};
+
+class BufferSliceDataView : public DataView {
+ public:
+  BufferSliceDataView(BufferSlice buffer_slice);
+  int64 size() override;
+  Result<BufferSlice> pread(int64 offset, int64 size) override;
+
+ private:
+  BufferSlice buffer_slice_;
+};
+
+class ConcatDataView : public DataView {
+ public:
+  ConcatDataView(DataView &left, DataView &right);
+  int64 size() override;
+  Result<BufferSlice> pread(int64 offset, int64 size) override;
+
+ private:
+  DataView &left_;
+  DataView &right_;
+};
+
+AesCbcState calc_aes_cbc_state(Slice seed);
+Result<ValueHash> calc_value_hash(DataView &data_view);
+BufferSlice gen_random_prefix(int64 data_size);
+
+class Password {
+ public:
+  Password(std::string password);
+  Slice as_slice() const;
+
+ private:
+  std::string password_;
+};
+
+class EncryptedSecret;
+
+class Secret {
+ public:
+  static Result<Secret> create(Slice secret);
+  static Secret create_new();
+
+  Slice as_slice() const;
+  EncryptedSecret encrypt(Slice key);
+
+ private:
+  Secret(UInt256 secret);
+  UInt256 secret_;
+};
+
+class EncryptedSecret {
+ public:
+  static Result<EncryptedSecret> create(Slice encrypted_secret);
+  Result<Secret> decrypt(Slice key);
+  Slice as_slice() const;
+
+ private:
+  EncryptedSecret(UInt256 encrypted_secret);
+  UInt256 encrypted_secret_;
+};
+
+// Decryption
+class Decryptor {
+ public:
+  Decryptor(AesCbcState aes_cbc_state);
+  Result<BufferSlice> append(BufferSlice data);
+  Result<ValueHash> finish();
+
+ private:
+  AesCbcState aes_cbc_state_;
+  Sha256State sha256_state_;
+  bool skipped_prefix_{false};
+  size_t to_skip_{0};
+};
+
+// Encryption
+class Encryptor : public DataView {
+ public:
+  Encryptor(AesCbcState aes_cbc_state, DataView &data_view);
+  int64 size() override;
+  Result<BufferSlice> pread(int64 offset, int64 size) override;
+
+ private:
+  AesCbcState aes_cbc_state_;
+  int64 current_offset_{0};
+  DataView &data_view_;
+};
+
+// Main functions
+
+//   decrypt :: (ValueSecret, ValueHash, EncryptedValue) -> Value
+//   encrypt :: (ValueSecret, RandomPrefix, Value) -> (EncryptedValue, ValueHash)
+
+struct EncryptedValue {
+  BufferSlice data;
+  ValueHash hash;
+};
+struct EncryptedFile {
+  std::string path;
+  ValueHash hash;
+};
+
+Result<EncryptedValue> encrypt_value(const Secret &secret, Slice data);
+Result<ValueHash> encrypt_file(const Secret &secret, std::string src, std::string dest);
+
+Result<BufferSlice> decrypt_value(const Secret &secret, const ValueHash &hash, Slice data);
+Status decrypt_file(const Secret &secret, const ValueHash &hash, std::string src, std::string dest);
+
+}  // namespace secure_storage
+}  // namespace td

test/CMakeLists.txt

@@ -7,6 +7,7 @@ set(TD_TEST_SOURCE
   ${CMAKE_CURRENT_SOURCE_DIR}/mtproto.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/message_entities.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/secret.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/secure_storage.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/string_cleaning.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/TestsRunner.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/tests_runner.cpp

test/secure_storage.cpp

@@ -0,0 +1,67 @@
+//
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2018
+//
+// Distributed under the Boost Software License, Version 1.0. (See accompanying
+// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
+//
+#include "td/utils/tests.h"
+#include "td/utils/filesystem.h"
+#include "td/utils/port/path.h"
+
+#include "td/telegram/SecureStorage.h"
+
+using namespace td;
+
+TEST(SecureStorage, secret) {
+  using namespace td::secure_storage;
+  auto secret = Secret::create_new();
+  std::string key = "cucumber";
+  auto encrypted_secret = secret.encrypt(key);
+  ASSERT_TRUE(encrypted_secret.as_slice() != secret.as_slice());
+  auto decrypted_secret = encrypted_secret.decrypt(key).ok();
+  ASSERT_TRUE(secret.as_slice() == decrypted_secret.as_slice());
+  ASSERT_TRUE(encrypted_secret.decrypt("notcucumber").is_error());
+}
+
+TEST(SecureStorage, simple) {
+  using namespace td::secure_storage;
+
+  BufferSlice value("Small tale about cucumbers");
+  auto value_secret = Secret::create_new();
+
+  auto value_view = BufferSliceDataView(value.copy());
+  BufferSlice prefix = gen_random_prefix(value_view.size());
+  auto prefix_view = BufferSliceDataView(std::move(prefix));
+  auto full_value_view = ConcatDataView(prefix_view, value_view);
+  auto hash = calc_value_hash(full_value_view).move_as_ok();
+
+  Encryptor encryptor(calc_aes_cbc_state(PSLICE() << value_secret.as_slice() << hash.as_slice()), full_value_view);
+  auto encrypted_value = encryptor.pread(0, encryptor.size()).move_as_ok();
+
+  Decryptor decryptor(calc_aes_cbc_state(PSLICE() << value_secret.as_slice() << hash.as_slice()));
+  auto res = decryptor.append(encrypted_value.copy()).move_as_ok();
+  auto decrypted_hash = decryptor.finish().ok();
+  ASSERT_TRUE(decrypted_hash.as_slice() == hash.as_slice());
+  ASSERT_TRUE(res.as_slice() == value.as_slice());
+
+  {
+    auto encrypted_value = encrypt_value(value_secret, value.as_slice()).move_as_ok();
+    auto decrypted_value =
+        decrypt_value(value_secret, encrypted_value.hash, encrypted_value.data.as_slice()).move_as_ok();
+    ASSERT_TRUE(decrypted_value.as_slice() == value.as_slice());
+  }
+
+  {
+    std::string value_path = "value.txt";
+    std::string encrypted_path = "encrypted.txt";
+    std::string decrypted_path = "decrypted.txt";
+    td::unlink(value_path).ignore();
+    td::unlink(encrypted_path).ignore();
+    td::unlink(decrypted_path).ignore();
+    std::string value(100000, 'a');
+    td::write_file(value_path, value);
+    auto hash = encrypt_file(value_secret, value_path, encrypted_path).move_as_ok();
+    decrypt_file(value_secret, hash, encrypted_path, decrypted_path).ensure();
+    ASSERT_TRUE(td::read_file(decrypted_path).move_as_ok().as_slice() == value);
+  }
+}