From ac6a83da500295bd0957fe08d6d429a925e21fcc Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Tue, 23 Jul 2019 03:26:26 +0300
Subject: [PATCH] Make Sha256State safe.

GitOrigin-RevId: 9b81cc1915b8ca1bc85219b863a6e8dcab814b37
---
 tdutils/td/utils/crypto.cpp | 30 +++++++++++++++++++++++++++---
 tdutils/td/utils/crypto.h   |  9 +++++++--
 tdutils/test/crypto.cpp     |  4 +++-
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/tdutils/td/utils/crypto.cpp b/tdutils/td/utils/crypto.cpp
index 7e27623d..91a3c439 100644
--- a/tdutils/td/utils/crypto.cpp
+++ b/tdutils/td/utils/crypto.cpp
@@ -382,20 +382,42 @@ struct Sha256StateImpl {
 };
 
 Sha256State::Sha256State() = default;
-Sha256State::Sha256State(Sha256State &&from) = default;
-Sha256State &Sha256State::operator=(Sha256State &&from) = default;
-Sha256State::~Sha256State() = default;
+
+Sha256State::Sha256State(Sha256State &&other) {
+  impl = std::move(other.impl);
+  is_inited = other.is_inited;
+  other.is_inited = false;
+}
+
+Sha256State &Sha256State::operator=(Sha256State &&other) {
+  Sha256State copy(std::move(other));
+  using std::swap;
+  swap(impl, copy.impl);
+  swap(is_inited, copy.is_inited);
+  return *this;
+}
+
+Sha256State::~Sha256State() {
+  if (is_inited) {
+    char result[32];
+    extract(MutableSlice{result, 32});
+    CHECK(!is_inited);
+  }
+}
 
 void sha256_init(Sha256State *state) {
   if (!state->impl) {
     state->impl = make_unique<Sha256StateImpl>();
   }
+  CHECK(!state->is_inited);
   int err = SHA256_Init(&state->impl->ctx);
   LOG_IF(FATAL, err != 1);
+  state->is_inited = true;
 }
 
 void sha256_update(Slice data, Sha256State *state) {
   CHECK(state->impl);
+  CHECK(state->is_inited);
   int err = SHA256_Update(&state->impl->ctx, data.ubegin(), data.size());
   LOG_IF(FATAL, err != 1);
 }
@@ -403,8 +425,10 @@ void sha256_update(Slice data, Sha256State *state) {
 void sha256_final(Sha256State *state, MutableSlice output, bool destroy) {
   CHECK(output.size() >= 32);
   CHECK(state->impl);
+  CHECK(state->is_inited);
   int err = SHA256_Final(output.ubegin(), &state->impl->ctx);
   LOG_IF(FATAL, err != 1);
+  state->is_inited = false;
   if (destroy) {
     state->impl.reset();
   }
diff --git a/tdutils/td/utils/crypto.h b/tdutils/td/utils/crypto.h
index 719105f3..f5a4a10d 100644
--- a/tdutils/td/utils/crypto.h
+++ b/tdutils/td/utils/crypto.h
@@ -78,9 +78,12 @@ void sha256_final(Sha256State *state, MutableSlice output, bool destroy = true);
 
 struct Sha256State {
   Sha256State();
-  Sha256State(Sha256State &&from);
-  Sha256State &operator=(Sha256State &&from);
+  Sha256State(const Sha256State &other) = delete;
+  Sha256State &operator=(const Sha256State &other) = delete;
+  Sha256State(Sha256State &&other);
+  Sha256State &operator=(Sha256State &&other);
   ~Sha256State();
+
   void init() {
     sha256_init(this);
   }
@@ -90,7 +93,9 @@ struct Sha256State {
   void extract(MutableSlice dest) {
     sha256_final(this, dest, false);
   }
+
   unique_ptr<Sha256StateImpl> impl;
+  bool is_inited = false;
 };
 
 void md5(Slice input, MutableSlice output);
diff --git a/tdutils/test/crypto.cpp b/tdutils/test/crypto.cpp
index e3da6df8..2f475567 100644
--- a/tdutils/test/crypto.cpp
+++ b/tdutils/test/crypto.cpp
@@ -72,10 +72,12 @@ TEST(Crypto, Sha256State) {
 
     td::Sha256State state;
     state.init();
+    td::Sha256State state2 = std::move(state);
     auto v = td::rand_split(s);
     for (auto &x : v) {
-      state.feed(x);
+      state2.feed(x);
     }
+    state = std::move(state2);
     td::UInt256 result;
     state.extract(as_slice(result));
     ASSERT_TRUE(baseline == result);