From eeddfa288debe5f84b5cf889ab74025cac5b2805 Mon Sep 17 00:00:00 2001
From: xsanm <kkurowski.kk@gmail.com>
Date: Wed, 12 Feb 2025 13:21:04 +0100
Subject: [PATCH] [CryptoModule] handle not-existing prekeys in Olm account

Summary:
Fixing the issue described in [ENG-10142](https://linear.app/comm/issue/ENG-10142/user-with-id-125229-has-broken-key-bundle-on-prod).

First, I wanted to implement a different solution, where instead of checking for empty prekeys we check if the signature is correct, but this will be a lot more complex:
1. Parse Prekey which format in `CryptoModule` is `{"curve25519":{"AAAAAQ":"Md8kV9rF5iTUi9IF290UIzpB8WdxmxUosjP1LFyXeSI"}}`;
2. Decode it somehow, Olm uses its own base64 encoder which is not part of the API, we have our own implementation in `Base64.h` but making it work will take some time.
3. Parse public singing key in CryptoModule.

The risk with this solution is that in theory, not generated prekey might have an indeterminate value, not just an empty encoded string. However, Olm uses only stack variables with constant size, e.g. `uint8_t public_key[CURVE25519_KEY_LENGTH];` so I think the risk is minimal. If someone thinks we should implement the safe one, I am okay with this too but this is just a lot more time-consuming.

This solution should work for users that already have broken key bundles in Identity and also for very old users that will log to Identity after logging.

Test Plan:
1. Patched my local Olm to not generate prekeys on account creation and added a bunch of logs.
2. Checked in AWS - my device has malformed keys as users listed in the task.
3. Patch this diff and build an app, `PrekeysHandler` called `validateAndUploadPrekeys` which healed the case, I am able to send DM to my mobile app and new keys are uploaded.
4. Comment out `PrekeysHandler`, log out and restore, I see the correct prekey uploaded for my device which makes sure that `validateAndGetPrekeys` works correctly too.

Reviewers: tomek, bartek, ashoat

Reviewed By: tomek, ashoat

Differential Revision: https://phab.comm.dev/D14354
---
 .../CommonCpp/CryptoTools/CryptoModule.cpp    | 23 +++++++++++++++++++
 .../cpp/CommonCpp/CryptoTools/CryptoModule.h  |  1 +
 2 files changed, 24 insertions(+)

diff --git a/native/cpp/CommonCpp/CryptoTools/CryptoModule.cpp b/native/cpp/CommonCpp/CryptoTools/CryptoModule.cpp
index 2ed7bfa80d..7c9ba9b164 100644
--- a/native/cpp/CommonCpp/CryptoTools/CryptoModule.cpp
+++ b/native/cpp/CommonCpp/CryptoTools/CryptoModule.cpp
@@ -121,6 +121,24 @@ bool CryptoModule::prekeyExistsAndOlderThan(uint64_t threshold) {
   return currentTime - lastPrekeyPublishTime >= threshold;
 }
 
+bool CryptoModule::prekeyDoesntExist() {
+  // Prekey generation when creating an Olm Account was added to clients
+  // after the initial launch of Olm. Because of that, there is a high
+  // chance there are users who have the account without a generated prekey.
+
+  // When prekey or signature is empty it contains bytes with only 0 values.
+  // Because of Base64 encoding, 0 is encoded as char `A` and empty
+  // prekey/signature is a string built from only `A`'s.
+  const std::string emptyPrekey(KEYSIZE, 'A');
+  const std::string emptySignature(SIGNATURESIZE, 'A');
+
+  std::string prekey = this->getPrekey();
+  std::string signature = this->getPrekeySignature();
+
+  return prekey.find(emptyPrekey) != std::string::npos ||
+      signature == emptySignature;
+}
+
 Keys CryptoModule::keysFromStrings(
     const std::string &identityKeys,
     const std::string &oneTimeKeys) {
@@ -455,6 +473,11 @@ std::optional<std::string> CryptoModule::validatePrekey() {
   static const uint64_t maxOldPrekeyAge = 2 * 60;
   std::optional<std::string> maybeNewPrekey;
 
+  bool prekeyDoesntExist = this->prekeyDoesntExist();
+  if (prekeyDoesntExist) {
+    return this->generateAndGetPrekey();
+  }
+
   bool shouldRotatePrekey =
       this->prekeyExistsAndOlderThan(maxPrekeyPublishTime);
   if (shouldRotatePrekey) {
diff --git a/native/cpp/CommonCpp/CryptoTools/CryptoModule.h b/native/cpp/CommonCpp/CryptoTools/CryptoModule.h
index b6066afd4c..2b61f448c4 100644
--- a/native/cpp/CommonCpp/CryptoTools/CryptoModule.h
+++ b/native/cpp/CommonCpp/CryptoTools/CryptoModule.h
@@ -29,6 +29,7 @@ class CryptoModule {
   // returns number of published keys
   size_t publishOneTimeKeys();
   bool prekeyExistsAndOlderThan(uint64_t threshold);
+  bool prekeyDoesntExist();
   OlmBuffer pickleAccount(const std::string &secretKey);
 
 public: