fix writing keys to file (weird assert bug)

OpenHD · Dec 14, 2023 · 69e83d3 · 69e83d3
1 parent f368fd1
commit 69e83d3
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 29 deletions.
diff --git a/executables/unit_test.cpp b/executables/unit_test.cpp
@@ -115,7 +115,9 @@ static void test_encrypt_decrypt_validate(const bool use_key_from_file,bool mess
   const std::string KEY_FILENAME="../example_key/txrx.key";
   wb::KeyPairTxRx keyPairTxRx{};
   if(use_key_from_file){
-    keyPairTxRx=wb::read_keypair_from_file(KEY_FILENAME);
+    auto tmp=wb::read_keypair_from_file(KEY_FILENAME);
+    assert(tmp.has_value());
+    keyPairTxRx=tmp.value();
   }else{
     const auto before=std::chrono::steady_clock::now();
     keyPairTxRx=wb::generate_keypair_from_bind_phrase("openhd");
@@ -176,6 +178,13 @@ static void test_encrypt_decrypt_validate(const bool use_key_from_file,bool mess
   }
   fmt::print("Test {} with {} passed\n",TEST_TYPE,TEST_KEY_TYPE);
 }
+static void test_encryption_serialize(){
+  auto keypair1=wb::generate_keypair_from_bind_phrase("openhd");
+  auto raw=wb::KeyPairTxRx::as_raw(keypair1);
+  auto serialized_deserialized=wb::KeyPairTxRx::from_raw(raw);
+  assert(keypair1==serialized_deserialized);
+  fmt::print("Serialize / Deserialize test passed\n");
+}
 
 
 int main(int argc, char *argv[]) {
@@ -210,6 +219,7 @@ int main(int argc, char *argv[]) {
 	}
 	if (test_mode == 0 || test_mode == 2) {
 	  std::cout << "Testing Encryption"<<std::endl;
+          test_encryption_serialize();
           test_encrypt_decrypt_validate(false, false);
           test_encrypt_decrypt_validate(false, true);
           test_encrypt_decrypt_validate(true, false);

diff --git a/executables/wfb_keygen.cpp b/executables/wfb_keygen.cpp
@@ -53,5 +53,12 @@ int main(int argc, char *const *argv) {
     keyPairTxRx=wb::generate_keypair_random();
   }
   //auto keypair=wb::generate_keypair_from_bind_phrase("openhd");
-  return wb::write_keypair_to_file(keyPairTxRx,"txrx.key");
+  auto res= wb::write_keypair_to_file(keyPairTxRx, "txrx.key");
+  if(res){
+    std::cout<<"Wrote keypair to file"<<std::endl;
+    return 0;
+  }else{
+    std::cout<<"Cannot write keypair to file"<<std::endl;
+    return -1;
+  }
 }
diff --git a/src/Encryption.cpp b/src/Encryption.cpp
@@ -50,35 +50,38 @@ wb::KeyPairTxRx wb::generate_keypair_from_bind_phrase(
   return ret;
 }
 
-int wb::write_keypair_to_file(const wb::KeyPairTxRx& keypair_txrx,
+bool wb::write_keypair_to_file(const wb::KeyPairTxRx& keypair_txrx,
                               const std::string& filename) {
   FILE *fp;
   if ((fp = fopen(filename.c_str(), "w")) == nullptr) {
     std::cerr<<"Unable to save "<<filename<<std::endl;
-    assert(false);
-    return 1;
+    return false;
+  }
+  const auto raw=KeyPairTxRx::as_raw(keypair_txrx);
+  auto res= fwrite(raw.data(),raw.size(),1,fp);
+  if(res!=1){
+    std::cerr<<"Cannot write to file"<<std::endl;
+    fclose(fp);
+    return false;
   }
-  assert(fwrite(keypair_txrx.key_1.secret_key.data(), crypto_box_SECRETKEYBYTES, 1, fp)==1);
-  assert(fwrite(keypair_txrx.key_1.public_key.data(), crypto_box_PUBLICKEYBYTES, 1, fp)==1);
-  assert(fwrite(keypair_txrx.key_2.secret_key.data(), crypto_box_SECRETKEYBYTES, 1, fp)==1);
-  assert(fwrite(keypair_txrx.key_2.public_key.data(), crypto_box_PUBLICKEYBYTES, 1, fp)==1);
-  fclose(fp);
-  return 0;
+  return true;
 }
 
-wb::KeyPairTxRx wb::read_keypair_from_file(const std::string& filename) {
+std::optional<wb::KeyPairTxRx> wb::read_keypair_from_file(const std::string& filename) {
   KeyPairTxRx ret{};
   FILE *fp;
   if ((fp = fopen(filename.c_str(), "r")) == nullptr) {
     std::cerr<<fmt::format("Unable to open {}: {}", filename.c_str(), strerror(errno))<<std::endl;
-    assert(false);
+    return std::nullopt;
   }
-  assert(fread(ret.key_1.secret_key.data(), crypto_box_SECRETKEYBYTES, 1, fp)==1);
-  assert(fread(ret.key_1.public_key.data(), crypto_box_PUBLICKEYBYTES, 1, fp)==1);
-  assert(fread(ret.key_2.secret_key.data(), crypto_box_SECRETKEYBYTES, 1, fp)==1);
-  assert(fread(ret.key_2.public_key.data(), crypto_box_PUBLICKEYBYTES, 1, fp)==1);
-  fclose(fp);
-  return ret;
+  std::array<uint8_t,KEYPAIR_RAW_SIZE> raw{};
+  auto res= fread(raw.data(),raw.size(),1,fp);
+  if(res!=1){
+    std::cerr<<"Cannot read keypair file"<<std::endl;
+    fclose(fp);
+    return std::nullopt;
+  }
+  return KeyPairTxRx::from_raw(raw);
 }
 
 std::array<uint8_t, crypto_aead_chacha20poly1305_KEYBYTES> wb::create_onetimeauth_subkey(
@@ -198,3 +201,16 @@ wb::Decryptor::authenticate_and_decrypt_buff(const uint64_t& nonce,
   }
   return nullptr;
 }
+
+wb::KeyPairTxRx wb::KeyPairTxRx::from_raw(
+    const std::array<uint8_t, KEYPAIR_RAW_SIZE>& raw) {
+    wb::KeyPairTxRx ret{};
+    memcpy(&ret,raw.data(),KEYPAIR_RAW_SIZE);
+    return ret;
+}
+std::array<uint8_t, wb::KEYPAIR_RAW_SIZE> wb::KeyPairTxRx::as_raw(
+    const wb::KeyPairTxRx& keypair) {
+    std::array<uint8_t, 32 * 4> ret{};
+    memcpy(ret.data(), &keypair,KEYPAIR_RAW_SIZE);
+    return ret;
+}
diff --git a/src/Encryption.h b/src/Encryption.h
@@ -2,12 +2,13 @@
 #ifndef ENCRYPTION_HPP
 #define ENCRYPTION_HPP
 
-#include <vector>
+#include <sodium.h>
+
 #include <array>
-#include <string>
 #include <memory>
-
-#include <sodium.h>
+#include <optional>
+#include <string>
+#include <vector>
 
 // Namespace that can be used to add encryption+packet validation
 // (Or packet validation only to save CPU resources)
@@ -27,11 +28,16 @@ namespace wb{
 struct Key {
   std::array<uint8_t,crypto_box_PUBLICKEYBYTES> public_key;
   std::array<uint8_t,crypto_box_SECRETKEYBYTES> secret_key;
-};
-
+  int operator==(const Key& other)const{
+    return std::equal(std::begin(public_key), std::end(public_key), std::begin(other.public_key)) &&
+            std::equal(std::begin(secret_key), std::end(secret_key), std::begin(other.secret_key));
+  }
+}__attribute__ ((packed));;
+static_assert(sizeof(Key)==crypto_box_PUBLICKEYBYTES+crypto_box_SECRETKEYBYTES);
 // A wb keypair are 2 keys, one for transmitting, one for receiving
 // (Since both ground and air unit talk bidirectional)
 // We use a different key for the down-link / uplink, respective
+static constexpr const int KEYPAIR_RAW_SIZE=32*4;
 struct KeyPairTxRx {
   Key key_1;
   Key key_2;
@@ -41,7 +47,14 @@ struct KeyPairTxRx {
   Key get_rx_key(bool is_air){
       return is_air ? key_2 : key_1;
   }
-};
+  int operator==(const KeyPairTxRx& other)const{
+      return key_1==other.key_1 && key_2==other.key_2;
+  }
+  static std::array<uint8_t,KEYPAIR_RAW_SIZE> as_raw(const KeyPairTxRx& keypair);
+  static KeyPairTxRx from_raw(const std::array<uint8_t,KEYPAIR_RAW_SIZE>& raw);
+}__attribute__ ((packed));
+static_assert(sizeof(KeyPairTxRx)==2*sizeof(Key));
+static_assert(sizeof(KeyPairTxRx)==KEYPAIR_RAW_SIZE);
 
 /**
  * Generates a new keypair. Non-deterministic, 100% secure.
@@ -67,13 +80,12 @@ KeyPairTxRx generate_keypair_from_bind_phrase(const std::string& bind_phrase=DEF
 /**
  * Saves the KeyPairTxRx as a raw file
  */
-int write_keypair_to_file(const KeyPairTxRx& keypair_txrx,const std::string& filename);
+bool write_keypair_to_file(const KeyPairTxRx& keypair_txrx,const std::string& filename);
 
 /**
  * Reads a raw KeyPairTxRx from a raw file previusly generated.
  */
-KeyPairTxRx read_keypair_from_file(const std::string& filename);
-
+std::optional<KeyPairTxRx> read_keypair_from_file(const std::string& filename);
 
 /**
  * https://libsodium.gitbook.io/doc/key_derivation