Work with PkDecryption keys by their private keys

Change interface to allow the app to get the private part of the
key and instantiate a decryption object from just the private part
of the key.

Changes the function generating a key from random bytes to be
initialising a key with a private key (because it's exactly the
same thing). Exports & imports private key parts as ArrayBuffer at
JS level rather than base64 assuming we are moving that way in
general.

7 files changed, 112 insertions, 21 deletions

diff --git a/include/olm/error.h b/include/olm/error.h
index 9d44a94..ee2187c 100644
--- a/include/olm/error.h
+++ b/include/olm/error.h
@@ -51,6 +51,13 @@ enum OlmErrorCode {
      */
     OLM_BAD_SIGNATURE = 14,
 
+    OLM_INPUT_BUFFER_TOO_SMALL = 15,
+
+    // Not an error code, just here to pad out the enum past 16 because
+    // otherwise the compiler warns about a redunant check. If you're
+    // adding an error code, replace this one!
+    OLM_ERROR_NOT_INVENTED_YET = 16,
+
     /* remember to update the list of string constants in error.c when updating
      * this list. */
 };
diff --git a/include/olm/pk.h b/include/olm/pk.h
index 1f3f9ff..5e779ce 100644
--- a/include/olm/pk.h
+++ b/include/olm/pk.h
@@ -76,7 +76,7 @@ size_t olm_pk_encrypt_random_length(
  * ciphertext, mac, or ephemeral_key buffers were too small then
  * olm_pk_encryption_last_error() will be "OUTPUT_BUFFER_TOO_SMALL". If there
  * weren't enough random bytes then olm_pk_encryption_last_error() will be
- * "NOT_ENOUGH_RANDOM". */
+ * "OLM_INPUT_BUFFER_TOO_SMALL". */
 size_t olm_pk_encrypt(
     OlmPkEncryption *encryption,
     void const * plaintext, size_t plaintext_length,
@@ -108,18 +108,24 @@ size_t olm_clear_pk_decryption(
     OlmPkDecryption *decryption
 );
 
-/** The number of random bytes needed to generate a new key. */
-size_t olm_pk_generate_key_random_length(void);
-
-/** Generate a new key to use for decrypting messages. The associated public
- * key will be written to the pubkey buffer. Returns olm_error() on failure. If
- * the pubkey buffer is too small then olm_pk_decryption_last_error() will be
- * "OUTPUT_BUFFER_TOO_SMALL". If there weren't enough random bytes then
- * olm_pk_decryption_last_error() will be "NOT_ENOUGH_RANDOM". */
-size_t olm_pk_generate_key(
+/** Get the number of bytes required to store an olm private key
+ */
+size_t olm_pk_private_key_length();
+
+/** Initialise the key from the private part of a key as returned by
+ * olm_pk_get_private_key(). The associated public key will be written to the
+ * pubkey buffer. Returns olm_error() on failure. If the pubkey buffer is too
+ * small then olm_pk_decryption_last_error() will be "OUTPUT_BUFFER_TOO_SMALL".
+ * If the private key was not long enough then olm_pk_decryption_last_error()
+ * will be "OLM_INPUT_BUFFER_TOO_SMALL".
+ *
+ * Note that the pubkey is a base64 encoded string, but the private key is
+ * an unencoded byte array
+ */
+size_t olm_pk_key_from_private(
     OlmPkDecryption * decryption,
     void * pubkey, size_t pubkey_length,
-    void * random, size_t random_length
+    void * privkey, size_t privkey_length
 );
 
 /** Returns the number of bytes needed to store a decryption object. */
@@ -171,6 +177,19 @@ size_t olm_pk_decrypt(
     void * plaintext, size_t max_plaintext_length
 );
 
+/**
+ * Get the private key for an OlmDecryption object as an unencoded byte array
+ * private_key must be a pointer to a buffer of at least
+ * olm_pk_private_key_length() bytes and this length must be passed in
+ * private_key_length. If the given buffer is too small, returns olm_error()
+ * and olm_pk_encryption_last_error() will be "OUTPUT_BUFFER_TOO_SMALL".
+ * Returns the number of bytes written.
+ */
+size_t olm_pk_get_private_key(
+    OlmPkDecryption * decryption,
+    void *private_key, size_t private_key_length
+);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/javascript/olm_pk.js b/javascript/olm_pk.js
index 25e0fee..2212470 100644
--- a/javascript/olm_pk.js
+++ b/javascript/olm_pk.js
@@ -118,16 +118,32 @@ PkDecryption.prototype['free'] = function() {
     free(this.ptr);
 }
 
+PkDecryption.prototype['init_with_private_key'] = restore_stack(function (private_key) {
+    var private_key_buffer = stack(private_key.length);
+    Module['HEAPU8'].set(private_key, private_key_buffer);
+
+    var pubkey_length = pk_decryption_method(
+        Module['_olm_pk_key_length']
+    )();
+    var pubkey_buffer = stack(pubkey_length + NULL_BYTE_PADDING_LENGTH);
+    pk_decryption_method(Module['_olm_pk_key_from_private'])(
+        this.ptr,
+        pubkey_buffer, pubkey_length,
+        private_key_buffer, private_key.length
+    );
+    return Pointer_stringify(pubkey_buffer);
+});
+
 PkDecryption.prototype['generate_key'] = restore_stack(function () {
     var random_length = pk_decryption_method(
-        Module['_olm_pk_generate_key_random_length']
+        Module['_olm_pk_private_key_length']
     )();
     var random_buffer = random_stack(random_length);
     var pubkey_length = pk_encryption_method(
         Module['_olm_pk_key_length']
     )();
     var pubkey_buffer = stack(pubkey_length + NULL_BYTE_PADDING_LENGTH);
-    pk_decryption_method(Module['_olm_pk_generate_key'])(
+    pk_decryption_method(Module['_olm_pk_key_from_private'])(
         this.ptr,
         pubkey_buffer, pubkey_length,
         random_buffer, random_length
@@ -135,6 +151,18 @@ PkDecryption.prototype['generate_key'] = restore_stack(function () {
     return Pointer_stringify(pubkey_buffer);
 });
 
+PkDecryption.prototype['get_private_key'] = restore_stack(function () {
+    var privkey_length = pk_encryption_method(
+        Module['_olm_pk_private_key_length']
+    )();
+    var privkey_buffer  = stack(privkey_length);
+    pk_decryption_method(Module['_olm_pk_get_private_key'])(
+        this.ptr,
+        privkey_buffer, privkey_length
+    );
+    return new Uint8Array(Module['HEAPU8'].buffer, privkey_buffer, privkey_length);
+});
+
 PkDecryption.prototype['pickle'] = restore_stack(function (key) {
     var key_array = array_from_string(key);
     var pickle_length = pk_decryption_method(
diff --git a/javascript/test/pk.spec.js b/javascript/test/pk.spec.js
index 007882f..d155cf5 100644
--- a/javascript/test/pk.spec.js
+++ b/javascript/test/pk.spec.js
@@ -49,6 +49,20 @@ describe("pk", function() {
         }
     });
 
+    it('should import & export keys from private parts', function () {
+        var alice_private = new Uint8Array([
+            0x77, 0x07, 0x6D, 0x0A, 0x73, 0x18, 0xA5, 0x7D,
+            0x3C, 0x16, 0xC1, 0x72, 0x51, 0xB2, 0x66, 0x45,
+            0xDF, 0x4C, 0x2F, 0x87, 0xEB, 0xC0, 0x99, 0x2A,
+            0xB1, 0x77, 0xFB, 0xA5, 0x1D, 0xB9, 0x2C, 0x2A
+        ]);
+        var alice_public = decryption.init_with_private_key(alice_private);
+        expect(alice_public).toEqual("hSDwCYkwp1R0i33ctD73Wg2/Og0mOBr066SpjqqbTmo");
+
+        var alice_private_out = decryption.get_private_key();
+        expect(alice_private_out).toEqual(alice_private);
+    });
+
     it('should encrypt and decrypt', function () {
         var TEST_TEXT='têst1';
         var pubkey = decryption.generate_key();
diff --git a/src/error.c b/src/error.c
index f541a93..5147b5c 100644
--- a/src/error.c
+++ b/src/error.c
@@ -31,6 +31,7 @@ static const char * ERRORS[] = {
     "UNKNOWN_MESSAGE_INDEX",
     "BAD_LEGACY_ACCOUNT_PICKLE",
     "BAD_SIGNATURE",
+    "OLM_INPUT_BUFFER_TOO_SMALL",
 };
 
 const char * _olm_error_to_string(enum OlmErrorCode error)
diff --git a/src/pk.cpp b/src/pk.cpp
index e646dc4..20ab991 100644
--- a/src/pk.cpp
+++ b/src/pk.cpp
@@ -176,7 +176,7 @@ size_t olm_clear_pk_decryption(
     return sizeof(OlmPkDecryption);
 }
 
-size_t olm_pk_generate_key_random_length(void) {
+size_t olm_pk_private_key_length(void) {
     return CURVE25519_KEY_LENGTH;
 }
 
@@ -184,23 +184,23 @@ size_t olm_pk_key_length(void) {
     return olm::encode_base64_length(CURVE25519_KEY_LENGTH);
 }
 
-size_t olm_pk_generate_key(
+size_t olm_pk_key_from_private(
     OlmPkDecryption * decryption,
     void * pubkey, size_t pubkey_length,
-    void * random, size_t random_length
+    void * privkey, size_t privkey_length
 ) {
     if (pubkey_length < olm_pk_key_length()) {
         decryption->last_error =
             OlmErrorCode::OLM_OUTPUT_BUFFER_TOO_SMALL;
         return std::size_t(-1);
     }
-    if (random_length < olm_pk_generate_key_random_length()) {
+    if (privkey_length < olm_pk_private_key_length()) {
         decryption->last_error =
-            OlmErrorCode::OLM_NOT_ENOUGH_RANDOM;
+            OlmErrorCode::OLM_INPUT_BUFFER_TOO_SMALL;
         return std::size_t(-1);
     }
 
-    _olm_crypto_curve25519_generate_key((uint8_t *) random, &decryption->key_pair);
+    _olm_crypto_curve25519_generate_key((uint8_t *) privkey, &decryption->key_pair);
     olm::encode_base64((const uint8_t *)decryption->key_pair.public_key.public_key, CURVE25519_KEY_LENGTH, (uint8_t *)pubkey);
     return 0;
 }
@@ -352,4 +352,21 @@ size_t olm_pk_decrypt(
     }
 }
 
+size_t olm_pk_get_private_key(
+    OlmPkDecryption * decryption,
+    void *private_key, size_t private_key_length
+) {
+    if (private_key_length < olm_pk_private_key_length()) {
+        decryption->last_error =
+            OlmErrorCode::OLM_OUTPUT_BUFFER_TOO_SMALL;
+        return std::size_t(-1);
+    }
+    std::memcpy(
+        private_key,
+        decryption->key_pair.private_key.private_key,
+        olm_pk_private_key_length()
+    );
+    return olm_pk_private_key_length();
+}
+
 }
diff --git a/tests/test_pk.cpp b/tests/test_pk.cpp
index ee12603..42cc8c9 100644
--- a/tests/test_pk.cpp
+++ b/tests/test_pk.cpp
@@ -36,7 +36,7 @@ const std::uint8_t *bob_public = (std::uint8_t *) "3p7bfXt9wbTTW2HC7OQ1Nz+DQ8hbe
 
 std::uint8_t pubkey[::olm_pk_key_length()];
 
-olm_pk_generate_key(
+olm_pk_key_from_private(
     decryption,
     pubkey, sizeof(pubkey),
     alice_private, sizeof(alice_private)
@@ -44,6 +44,11 @@ olm_pk_generate_key(
 
 assert_equals(alice_public, pubkey, olm_pk_key_length());
 
+uint8_t *alice_private_back_out = (uint8_t *)malloc(olm_pk_private_key_length());
+olm_pk_get_private_key(decryption, alice_private_back_out, olm_pk_private_key_length());
+assert_equals(alice_private, alice_private_back_out, olm_pk_private_key_length());
+free(alice_private_back_out);
+
 std::uint8_t encryption_buffer[olm_pk_encryption_size()];
 OlmPkEncryption *encryption = olm_pk_encryption(encryption_buffer);
 
@@ -105,7 +110,7 @@ const std::uint8_t *alice_public = (std::uint8_t *) "hSDwCYkwp1R0i33ctD73Wg2/Og0
 
 std::uint8_t pubkey[olm_pk_key_length()];
 
-olm_pk_generate_key(
+olm_pk_key_from_private(
     decryption,
     pubkey, sizeof(pubkey),
     alice_private, sizeof(alice_private)