e2e: handle unencrypted prefix of packet in library

5 files changed, 78 insertions, 36 deletions

diff --git a/tde2e/td/e2e/Call.cpp b/tde2e/td/e2e/Call.cpp
index f249c1f29..dfbb4659e 100644
--- a/tde2e/td/e2e/Call.cpp
+++ b/tde2e/td/e2e/Call.cpp
@@ -38,6 +38,15 @@ CallVerificationChain::State CallVerificationChain::get_state() const {
   return state_;
 }
 
+template <typename... Args>
+std::string concat(Args... args) {
+  const size_t total_size = (args.size() + ...);
+  std::string result;
+  result.reserve(total_size);
+  (result.append(args.data(), args.size()), ...);
+  return result;
+}
+
 template <class F>
 struct LambdaStorer {
   const F &store_;
@@ -138,7 +147,7 @@ std::string CallVerificationChain::to_short_string(e2e::object_ptr<e2e::e2e_chai
   return sb.as_cslice().str();
 }
 
-td::Status CallVerificationChain::process_broadcast(std::string message,
+td::Status CallVerificationChain::process_broadcast(td::Slice message,
                                                     e2e::object_ptr<e2e::e2e_chain_GroupBroadcast> broadcast) {
   td::Status status;
   td::UInt256 broadcast_chain_hash{};
@@ -274,8 +283,19 @@ void CallEncryption::forget_shared_key(td::int32 epoch, td::UInt256 epoch_hash)
   epochs_to_forget_.emplace(td::Timestamp::in(FORGET_EPOCH_DELAY), epoch);
 }
 
-td::Result<std::string> CallEncryption::decrypt(td::int64 user_id, td::int32 channel_id, td::Slice encrypted_data) {
+td::Result<std::string> CallEncryption::decrypt(td::int64 user_id, td::int32 channel_id, td::Slice packet) {
   sync();
+  if (packet.size() < 4) {
+    return td::Status::Error("Packet too small");
+  }
+  td::uint32 unencrypted_prefix_size = td::as<td::uint32>(packet.data() + packet.size() - 4);
+  packet.remove_suffix(4);
+  if (unencrypted_prefix_size > packet.size() || unencrypted_prefix_size >= (1 << 16)) {
+    return td::Status::Error("Unencrypted prefix size is too large");
+  }
+  auto unencrypted_prefix = packet.substr(0, unencrypted_prefix_size);
+  auto encrypted_data = packet.substr(unencrypted_prefix_size);
+
   if (user_id == user_id_) {
     return td::Status::Error("Packet is encrypted by us");
   }
@@ -321,17 +341,25 @@ td::Result<std::string> CallEncryption::decrypt(td::int64 user_id, td::int32 cha
         auto &epoch_info = it2->second;
         TRY_RESULT(one_time_secret,
                    MessageEncryption::decrypt_header(encrypted_header, encrypted_packet, epoch_info.secret_));
-        return decrypt_packet_with_secret(user_id, channel_id, unencrypted_header, encrypted_packet, one_time_secret,
-                                          epoch_info.group_state_);
+        return decrypt_packet_with_secret(user_id, channel_id, unencrypted_header, unencrypted_prefix, encrypted_packet,
+                                          one_time_secret, epoch_info.group_state_);
       }
     }
   }
   return Error(E::Decrypt_UnknownEpoch);
 }
 
-td::Result<std::string> CallEncryption::encrypt(td::int32 channel_id, td::Slice decrypted_data) {
+td::Result<std::string> CallEncryption::encrypt(td::int32 channel_id, td::Slice data,
+                                                size_t unencrypted_header_length) {
   sync();
 
+  if (unencrypted_header_length > data.size() || unencrypted_header_length >= (1 << 16)) {
+    return td::Status::Error("Unencrypted header length is too large");
+  }
+  TRY_RESULT(unencrypted_header_length_u32, td::narrow_cast_safe<td::uint32>(unencrypted_header_length));
+  auto unencrypted_prefix = data.substr(0, unencrypted_header_length);
+  auto decrypted_data = data.substr(unencrypted_header_length);
+
   // use all active epochs
   if (epochs_.empty()) {
     return Error(E::Encrypt_UnknownEpoch);
@@ -348,7 +376,8 @@ td::Result<std::string> CallEncryption::encrypt(td::int32 channel_id, td::Slice
 
   td::SecureString one_time_secret(32, 0);
   td::Random::secure_bytes(one_time_secret.as_mutable_slice());
-  TRY_RESULT(encrypted_packet, encrypt_packet_with_secret(channel_id, header_a, decrypted_data, one_time_secret));
+  TRY_RESULT(encrypted_packet, encrypt_packet_with_secret(channel_id, header_a + unencrypted_prefix.str(),
+                                                          decrypted_data, one_time_secret));
 
   std::vector<td::SecureString> encrypted_headers;
   for (auto &[epoch_i, epoch] : epochs_) {
@@ -363,14 +392,16 @@ td::Result<std::string> CallEncryption::encrypt(td::int32 channel_id, td::Slice
     }
   });
 
-  //LOG(ERROR) << decrypted_data.size() << " -> " << header_a.size() << " + " << header_b.size() << " + " << encrypted_packet.size();
-  return header_a + header_b + encrypted_packet;
+  std::string trailer(4, '\0');
+  td::as<td::uint32>(trailer.data()) = unencrypted_header_length_u32;
+
+  //LOG(ERROR) << decrypted_data.size() << " -> " << unencrypted_prefix.size() << " + " << header_a.size() << " + " << header_b.size() << " + " << encrypted_packet.size() << " + " << trailer.size();
+  return concat(unencrypted_prefix, header_a, header_b, encrypted_packet, trailer);
 }
 
-std::string add_magic(td::int32 magic, td::Slice header) {
-  std::string res(4 + header.size(), '\0');
+std::string make_magic(td::int32 magic) {
+  std::string res(4, '\0');
   td::as<td::int32>(res.data()) = magic;
-  td::MutableSlice(res).substr(4).copy_from(header);
   return res;
 }
 
@@ -393,16 +424,17 @@ td::Result<std::string> CallEncryption::encrypt_packet_with_secret(td::int32 cha
   // TODO: there is too much copies happening here. Almost all of them could be avoided
   td::UInt256 large_msg_id{};
   auto encrypted_payload = MessageEncryption::encrypt_data(
-      payload, one_time_secret, add_magic(td::e2e_api::e2e_callPacket::ID, unencrypted_part), &large_msg_id);
-  auto to_sign = add_magic(td::e2e_api::e2e_callPacketLargeMsgId::ID, large_msg_id.as_slice());
+      payload, one_time_secret, concat(make_magic(td::e2e_api::e2e_callPacket::ID), unencrypted_part), &large_msg_id);
+  auto to_sign = concat(make_magic(td::e2e_api::e2e_callPacketLargeMsgId::ID), large_msg_id.as_slice());
 
   TRY_RESULT(signature, private_key_.sign(to_sign));
   return encrypted_payload.as_slice().str() + signature.to_slice().str();
 }
 
 td::Result<std::string> CallEncryption::decrypt_packet_with_secret(
-    td::int64 expected_user_id, td::int32 expected_channel_id, td::Slice unencrypted_header, td::Slice encrypted_packet,
-    td::Slice one_time_secret, const GroupStateRef &group_state) {
+    td::int64 expected_user_id, td::int32 expected_channel_id, td::Slice unencrypted_header,
+    td::Slice unencrypted_prefix, td::Slice encrypted_packet, td::Slice one_time_secret,
+    const GroupStateRef &group_state) {
   TRY_RESULT(participant, group_state->get_participant(expected_user_id));
   if (encrypted_packet.size() < 64) {
     return td::Status::Error("Not enough encryption data");
@@ -411,13 +443,14 @@ td::Result<std::string> CallEncryption::decrypt_packet_with_secret(
   encrypted_packet.remove_suffix(64);
 
   td::UInt256 large_msg_id{};
-  TRY_RESULT(payload_str, MessageEncryption::decrypt_data(
-                              encrypted_packet, one_time_secret,
-                              add_magic(td::e2e_api::e2e_callPacket::ID, unencrypted_header), &large_msg_id));
+  TRY_RESULT(payload_str, MessageEncryption::decrypt_data(encrypted_packet, one_time_secret,
+                                                          concat(make_magic(td::e2e_api::e2e_callPacket::ID),
+                                                                 unencrypted_header, unencrypted_prefix),
+                                                          &large_msg_id));
   // we know that this is packet created by some participant
 
   auto payload = td::Slice(payload_str);
-  auto to_verify = add_magic(td::e2e_api::e2e_callPacketLargeMsgId::ID, large_msg_id.as_slice());
+  auto to_verify = concat(make_magic(td::e2e_api::e2e_callPacketLargeMsgId::ID), large_msg_id.as_slice());
   TRY_STATUS(participant.public_key.verify(to_verify, signature));
 
   td::TlParser parser(payload);
@@ -436,7 +469,7 @@ td::Result<std::string> CallEncryption::decrypt_packet_with_secret(
   }
   TRY_STATUS(check_not_seen(participant.public_key, channel_id, seqno));
   mark_as_seen(participant.public_key, channel_id, seqno);
-  return result;
+  return concat(unencrypted_prefix, result);
 }
 
 td::Status CallEncryption::check_not_seen(const PublicKey &public_key, td::int32 channel_id, td::uint32 seqno) {
diff --git a/tde2e/td/e2e/Call.h b/tde2e/td/e2e/Call.h
index ab04b7272..4e55d0d4a 100644
--- a/tde2e/td/e2e/Call.h
+++ b/tde2e/td/e2e/Call.h
@@ -55,7 +55,7 @@ struct CallVerificationChain {
   friend td::StringBuilder &operator<<(td::StringBuilder &sb, const CallVerificationChain &chain);
 
  private:
-  td::Status process_broadcast(std::string message, e2e::object_ptr<e2e::e2e_chain_GroupBroadcast> broadcast);
+  td::Status process_broadcast(td::Slice message, e2e::object_ptr<e2e::e2e_chain_GroupBroadcast> broadcast);
   td::Status process_broadcast(e2e::e2e_chain_groupBroadcastNonceCommit &nonce_commit);
   td::Status process_broadcast(e2e::e2e_chain_groupBroadcastNonceReveal &nonce_reveal);
 
@@ -91,8 +91,8 @@ class CallEncryption {
   td::Status add_shared_key(td::int32 epoch, td::UInt256 epoch_hash, td::SecureString key, GroupStateRef group_state);
   void forget_shared_key(td::int32 epoch, td::UInt256 epoch_hash);
 
-  td::Result<std::string> decrypt(td::int64 expected_user_id, td::int32 expected_channel_id, td::Slice encrypted_data);
-  td::Result<std::string> encrypt(td::int32 channel_id, td::Slice decrypted_data);
+  td::Result<std::string> decrypt(td::int64 expected_user_id, td::int32 expected_channel_id, td::Slice packet);
+  td::Result<std::string> encrypt(td::int32 channel_id, td::Slice data, size_t unencrypted_header_length);
 
  private:
   static constexpr double FORGET_EPOCH_DELAY = 10;
@@ -128,8 +128,9 @@ class CallEncryption {
   td::Result<std::string> encrypt_packet_with_secret(td::int32 channel_id, td::Slice header, td::Slice packet,
                                                      td::Slice one_time_secret);
   td::Result<std::string> decrypt_packet_with_secret(td::int64 expected_user_id, td::int32 expected_channel_id,
-                                                     td::Slice unencrypted_packet, td::Slice encrypted_packet,
-                                                     td::Slice one_time_secret, const GroupStateRef &group_state);
+                                                     td::Slice unencrypted_header, td::Slice unencrypted_prefix,
+                                                     td::Slice encrypted_packet, td::Slice one_time_secret,
+                                                     const GroupStateRef &group_state);
   td::Status check_not_seen(const PublicKey &public_key, td::int32 channel_id, td::uint32 seqno);
   void mark_as_seen(const PublicKey &public_key, td::int32 channel_id, td::uint32 seqno);
   static td::Status validate_channel_id(td::int32 channel_id);
@@ -189,9 +190,9 @@ struct Call {
     TRY_STATUS(get_status());
     return call_encryption_.decrypt(user_id, channel_id, encrypted_data);
   }
-  td::Result<std::string> encrypt(td::int32 channel_id, td::Slice decrypted_data) {
+  td::Result<std::string> encrypt(td::int32 channel_id, td::Slice decrypted_data, size_t unencrypted_prefix_size) {
     TRY_STATUS(get_status());
-    return call_encryption_.encrypt(channel_id, decrypted_data);
+    return call_encryption_.encrypt(channel_id, decrypted_data, unencrypted_prefix_size);
   }
 
   td::Result<std::vector<std::string>> pull_outbound_messages() {
diff --git a/tde2e/td/e2e/e2e_api.cpp b/tde2e/td/e2e/e2e_api.cpp
index f87aa05bd..55f9d414a 100644
--- a/tde2e/td/e2e/e2e_api.cpp
+++ b/tde2e/td/e2e/e2e_api.cpp
@@ -467,9 +467,9 @@ class KeyChain {
     TRY_RESULT(shared_key, call_ref->shared_key());
     return shared_key.as_slice().str();
   }
-  td::Result<api::Bytes> call_encrypt(api::CallId call_id, api::CallChannelId channel_id, td::Slice message) {
+  td::Result<api::Bytes> call_encrypt(api::CallId call_id, api::CallChannelId channel_id, td::Slice message, size_t unencrypted_prefix_size) {
     TRY_RESULT(call_ref, to_call_ref(call_id));
-    return call_ref->encrypt(channel_id, message);
+    return call_ref->encrypt(channel_id, message, unencrypted_prefix_size);
   }
   td::Result<api::SecureBytes> call_decrypt(api::CallId call_id, api::UserId user_id, api::CallChannelId channel_id,
                                             td::Slice message) {
@@ -792,8 +792,8 @@ Result<Bytes> call_create_change_state_block(CallId call_id, const CallState &ne
 Result<SecureBytes> call_export_shared_key(CallId call_id) {
   return get_default_keychain().call_export_shared_key(call_id);
 }
-Result<Bytes> call_encrypt(CallId call_id, CallChannelId channel_id, SecureSlice message) {
-  return get_default_keychain().call_encrypt(call_id, channel_id, to_slice(message));
+Result<Bytes> call_encrypt(CallId call_id, CallChannelId channel_id, SecureSlice message, size_t unencrypted_prefix_size) {
+  return get_default_keychain().call_encrypt(call_id, channel_id, to_slice(message), unencrypted_prefix_size);
 }
 Result<SecureBytes> call_decrypt(CallId call_id, UserId user_id, CallChannelId channel_id, Slice message) {
   return get_default_keychain().call_decrypt(call_id, user_id, channel_id, to_slice(message));
diff --git a/tde2e/td/e2e/e2e_api.h b/tde2e/td/e2e/e2e_api.h
index 1c25c2228..a44a3195e 100644
--- a/tde2e/td/e2e/e2e_api.h
+++ b/tde2e/td/e2e/e2e_api.h
@@ -303,7 +303,7 @@ Result<std::string> call_describe_block(Slice block);
 Result<std::string> call_describe_message(Slice message);
 
 Result<Bytes> call_create_change_state_block(CallId call_id, const CallState &new_state);
-Result<Bytes> call_encrypt(CallId call_id, CallChannelId channel_id, SecureSlice message);
+Result<Bytes> call_encrypt(CallId call_id, CallChannelId channel_id, SecureSlice message, size_t unencrypted_prefix_size);
 Result<SecureBytes> call_decrypt(CallId call_id, UserId user_id, CallChannelId channel_id, Slice message);
 
 Result<int> call_get_height(CallId call_id);
diff --git a/tde2e/test/e2e.cpp b/tde2e/test/e2e.cpp
index 808a776f7..43dac6c58 100644
--- a/tde2e/test/e2e.cpp
+++ b/tde2e/test/e2e.cpp
@@ -786,8 +786,8 @@ TEST(Call, Basic_API) {
   ASSERT_EQ(call_get_verification_state(call2).value().emoji_hash.value(),
             call_get_verification_state(call3).value().emoji_hash.value());
 
-  auto e = call_encrypt(call2, 1, "hello").value();
-  auto e2 = call_encrypt(call2, 1, "hello").value();
+  auto e = call_encrypt(call2, 1, "hello", 0).value();
+  auto e2 = call_encrypt(call2, 1, "hello", 0).value();
   CHECK(e != "hello");
   LOG(ERROR) << e.size();
   ASSERT_TRUE(!call_decrypt(call2, 2, 1, e).is_ok());
@@ -796,13 +796,21 @@ TEST(Call, Basic_API) {
   ASSERT_EQ("hello", call_decrypt(call3, 2, 1, e).value());
   ASSERT_TRUE(!call_decrypt(call3, 2, 1, e).is_ok());
 
+  {
+    auto hel_x = call_encrypt(call2, 1, "hello world", 3).value();
+    ASSERT_TRUE(td::begins_with(hel_x, "hel"));
+    ASSERT_TRUE(!td::begins_with(hel_x, "hello wo"));
+    auto hello = call_decrypt(call3, 2, 1, hel_x).value();
+    ASSERT_EQ("hello world", hello);
+  }
+
   auto block3 = F(call_create_change_state_block(
                       call2, CallState{0, {CallParticipant{2, pkey2, 3}, CallParticipant{3, pkey3, 3}}}))
                     .value();
   call_apply_block(call3, block3).value();
   ASSERT_TRUE(!call_decrypt(call3, 2, 1, e).is_ok());
   ASSERT_EQ("hello", call_decrypt(call3, 2, 1, e2).value());
-  ASSERT_TRUE(call_decrypt(call2, 3, 1, call_encrypt(call3, 1, "bye").value()).is_ok());
+  ASSERT_TRUE(call_decrypt(call2, 3, 1, call_encrypt(call3, 1, "bye", 0).value()).is_ok());
 
   LOG(ERROR) << call_describe(call1).value();
   LOG(ERROR) << call_describe(call2).value();
@@ -879,7 +887,7 @@ class CallEncryptionBench final : public td::Benchmark {
 
   void run(int n) final {
     for (int i = 0; i < n; i++) {
-      auto encrypted = e1_->encrypt(1, msg_).move_as_ok();
+      auto encrypted = e1_->encrypt(1, msg_, 0).move_as_ok();
       CHECK(msg_ == e2_->decrypt(1, 1, encrypted).move_as_ok());
     }
   }