add functions for pickling/unpickling a decryption object

5 files changed, 250 insertions, 0 deletions

diff --git a/include/olm/pk.h b/include/olm/pk.h
index a91a80d..d2e031e 100644
--- a/include/olm/pk.h
+++ b/include/olm/pk.h
@@ -122,6 +122,36 @@ size_t olm_pk_generate_key(
     void * random, size_t random_length
 );
 
+/** Returns the number of bytes needed to store a decryption object. */
+size_t olm_pickle_pk_decryption_length(
+    OlmPkDecryption * decryption
+);
+
+/** Stores decryption object as a base64 string. Encrypts the object using the
+ * supplied key. Returns the length of the pickled object on success.
+ * Returns olm_error() on failure. If the pickle output buffer
+ * is smaller than olm_pickle_account_length() then
+ * olm_pk_decryption_last_error() will be "OUTPUT_BUFFER_TOO_SMALL" */
+size_t olm_pickle_pk_decryption(
+    OlmPkDecryption * decryption,
+    void const * key, size_t key_length,
+    void *pickled, size_t pickled_length
+);
+
+/** Loads a decryption object from a pickled base64 string. The associated
+ * public key will be written to the pubkey buffer. Decrypts the object using
+ * the supplied key. Returns olm_error() on failure. If the key doesn't
+ * match the one used to encrypt the account then olm_pk_decryption_last_error()
+ * will be "BAD_ACCOUNT_KEY". If the base64 couldn't be decoded then
+ * olm_pk_decryption_last_error() will be "INVALID_BASE64". The input pickled
+ * buffer is destroyed */
+size_t olm_unpickle_pk_decryption(
+    OlmPkDecryption * decryption,
+    void const * key, size_t key_length,
+    void *pickled, size_t pickled_length,
+    void *pubkey, size_t pubkey_length
+);
+
 /** Get the length of the plaintext that will correspond to a ciphertext of the
  * given length. */
 size_t olm_pk_max_plaintext_length(
diff --git a/javascript/olm_pk.js b/javascript/olm_pk.js
index 580d854..2542707 100644
--- a/javascript/olm_pk.js
+++ b/javascript/olm_pk.js
@@ -135,6 +135,35 @@ PkDecryption.prototype['generate_key'] = restore_stack(function () {
     return Pointer_stringify(pubkey_buffer);
 });
 
+PkDecryption.prototype['pickle'] = restore_stack(function (key) {
+    var key_array = array_from_string(key);
+    var pickle_length = pk_decryption_method(
+        Module['_olm_pickle_pk_decryption_length']
+    )(this.ptr);
+    var key_buffer = stack(key_array);
+    var pickle_buffer = stack(pickle_length + NULL_BYTE_PADDING_LENGTH);
+    pk_decryption_method(Module['_olm_pickle_pk_decryption'])(
+        this.ptr, key_buffer, key_array.length, pickle_buffer, pickle_length
+    );
+    return Pointer_stringify(pickle_buffer);
+});
+
+PkDecryption.prototype['unpickle'] = restore_stack(function (key, pickle) {
+    var key_array = array_from_string(key);
+    var key_buffer = stack(key_array);
+    var pickle_array = array_from_string(pickle);
+    var pickle_buffer = stack(pickle_array);
+    var ephemeral_length = pk_decryption_method(
+        Module["_olm_pk_key_length"]
+    )();
+    var ephemeral_buffer = stack(ephemeral_length + NULL_BYTE_PADDING_LENGTH);
+    pk_decryption_method(Module['_olm_unpickle_pk_decryption'])(
+        this.ptr, key_buffer, key_array.length, pickle_buffer,
+        pickle_array.length, ephemeral_buffer, ephemeral_length
+    );
+    return Pointer_stringify(ephemeral_buffer);
+});
+
 PkDecryption.prototype['decrypt'] = restore_stack(function (
     ephemeral_key, mac, ciphertext
 ) {
diff --git a/javascript/test/pk.spec.js b/javascript/test/pk.spec.js
index 9eec47e..aec90ac 100644
--- a/javascript/test/pk.spec.js
+++ b/javascript/test/pk.spec.js
@@ -61,4 +61,22 @@ describe("pk", function() {
         console.log(TEST_TEXT, "->", decrypted);
         expect(decrypted).toEqual(TEST_TEXT);
     });
+
+    it('should pickle and unpickle', function () {
+        var TEST_TEXT = 'têst1';
+        var pubkey = decryption.generate_key();
+        encryption.set_recipient_key(pubkey);
+        var encrypted = encryption.encrypt(TEST_TEXT);
+
+        var PICKLE_KEY = 'secret_key';
+        var pickle = decryption.pickle(PICKLE_KEY);
+
+        var new_decryption = new Olm.PkDecryption();
+        var new_pubkey = new_decryption.unpickle(PICKLE_KEY, pickle);
+        expect(new_pubkey).toEqual(pubkey);
+        var decrypted = new_decryption.decrypt(encrypted.ephemeral, encrypted.mac, encrypted.ciphertext);
+        console.log(TEST_TEXT, "->", decrypted);
+        expect(decrypted).toEqual(TEST_TEXT);
+        new_decryption.free();
+    });
 });
diff --git a/src/pk.cpp b/src/pk.cpp
index b4edf15..1925d7e 100644
--- a/src/pk.cpp
+++ b/src/pk.cpp
@@ -19,6 +19,8 @@
 #include "olm/error.h"
 #include "olm/memory.hh"
 #include "olm/base64.hh"
+#include "olm/pickle_encoding.h"
+#include "olm/pickle.hh"
 
 extern "C" {
 
@@ -203,6 +205,106 @@ size_t olm_pk_generate_key(
     return 0;
 }
 
+namespace {
+    static const std::uint32_t PK_DECRYPTION_PICKLE_VERSION = 1;
+
+    static std::size_t pickle_length(
+        OlmPkDecryption const & value
+    ) {
+        std::size_t length = 0;
+        length += olm::pickle_length(PK_DECRYPTION_PICKLE_VERSION);
+        length += olm::pickle_length(value.key_pair);
+        return length;
+    }
+
+
+    static std::uint8_t * pickle(
+        std::uint8_t * pos,
+        OlmPkDecryption const & value
+    ) {
+        pos = olm::pickle(pos, PK_DECRYPTION_PICKLE_VERSION);
+        pos = olm::pickle(pos, value.key_pair);
+        return pos;
+    }
+
+
+    static std::uint8_t const * unpickle(
+        std::uint8_t const * pos, std::uint8_t const * end,
+        OlmPkDecryption & value
+    ) {
+        uint32_t pickle_version;
+        pos = olm::unpickle(pos, end, pickle_version);
+
+        switch (pickle_version) {
+        case 1:
+            break;
+
+        default:
+            value.last_error = OlmErrorCode::OLM_UNKNOWN_PICKLE_VERSION;
+            return end;
+        }
+
+        pos = olm::unpickle(pos, end, value.key_pair);
+        return pos;
+    }
+}
+
+size_t olm_pickle_pk_decryption_length(
+    OlmPkDecryption * decryption
+) {
+    return _olm_enc_output_length(pickle_length(*decryption));
+}
+
+size_t olm_pickle_pk_decryption(
+    OlmPkDecryption * decryption,
+    void const * key, size_t key_length,
+    void *pickled, size_t pickled_length
+) {
+    OlmPkDecryption & object = *decryption;
+    std::size_t raw_length = pickle_length(object);
+    if (pickled_length < _olm_enc_output_length(raw_length)) {
+        object.last_error = OlmErrorCode::OLM_OUTPUT_BUFFER_TOO_SMALL;
+        return std::size_t(-1);
+    }
+    pickle(_olm_enc_output_pos(reinterpret_cast<std::uint8_t *>(pickled), raw_length), object);
+    return _olm_enc_output(reinterpret_cast<std::uint8_t const *>(key), key_length, reinterpret_cast<std::uint8_t *>(pickled), raw_length);
+}
+
+size_t olm_unpickle_pk_decryption(
+    OlmPkDecryption * decryption,
+    void const * key, size_t key_length,
+    void *pickled, size_t pickled_length,
+    void *pubkey, size_t pubkey_length
+) {
+    OlmPkDecryption & object = *decryption;
+    if (pubkey != NULL && pubkey_length < olm_pk_key_length()) {
+        object.last_error = OlmErrorCode::OLM_OUTPUT_BUFFER_TOO_SMALL;
+        return std::size_t(-1);
+    }
+    std::uint8_t * const pos = reinterpret_cast<std::uint8_t *>(pickled);
+    std::size_t raw_length = _olm_enc_input(
+        reinterpret_cast<std::uint8_t const *>(key), key_length, pos, pickled_length, &object.last_error
+    );
+    if (raw_length == std::size_t(-1)) {
+        return std::size_t(-1);
+    }
+    std::uint8_t * const end = pos + raw_length;
+    /* On success unpickle will return (pos + raw_length). If unpickling
+     * terminates too soon then it will return a pointer before
+     * (pos + raw_length). On error unpickle will return (pos + raw_length + 1).
+     */
+    if (end != unpickle(pos, end + 1, object)) {
+        if (object.last_error == OlmErrorCode::OLM_SUCCESS) {
+            object.last_error = OlmErrorCode::OLM_CORRUPTED_PICKLE;
+        }
+        return std::size_t(-1);
+    }
+    if (pubkey != NULL) {
+        olm::encode_base64((const uint8_t *)object.key_pair.public_key.public_key, CURVE25519_KEY_LENGTH, (uint8_t *)pubkey);
+    }
+    return pickled_length;
+}
+
 size_t olm_pk_max_plaintext_length(
     OlmPkDecryption * decryption,
     size_t ciphertext_length
diff --git a/tests/test_pk.cpp b/tests/test_pk.cpp
index dcb412a..ab1f477 100644
--- a/tests/test_pk.cpp
+++ b/tests/test_pk.cpp
@@ -87,4 +87,75 @@ free(plaintext_buffer);
 
 }
 
+{ /* Encryption Test Case 1 */
+
+TestCase test_case("Public Key Decryption pickling");
+
+std::uint8_t decryption_buffer[olm_pk_decryption_size()];
+OlmPkDecryption *decryption = olm_pk_decryption(decryption_buffer);
+
+std::uint8_t alice_private[32] = {
+    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
+};
+
+const std::uint8_t *alice_public = (std::uint8_t *) "hSDwCYkwp1R0i33ctD73Wg2/Og0mOBr066SpjqqbTmoK";
+
+std::uint8_t pubkey[olm_pk_key_length()];
+
+olm_pk_generate_key(
+    decryption,
+    pubkey, sizeof(pubkey),
+    alice_private, sizeof(alice_private)
+);
+
+const uint8_t *PICKLE_KEY=(uint8_t *)"secret_key";
+std::uint8_t pickle_buffer[olm_pickle_pk_decryption_length(decryption)];
+const uint8_t *expected_pickle = (uint8_t *) "qx37WTQrjZLz5tId/uBX9B3/okqAbV1ofl9UnHKno1eipByCpXleAAlAZoJgYnCDOQZDQWzo3luTSfkF9pU1mOILCbbouubs6TVeDyPfgGD9i86J8irHjA";
+
+olm_pickle_pk_decryption(
+    decryption,
+    PICKLE_KEY, strlen((char *)PICKLE_KEY),
+    pickle_buffer, sizeof(pickle_buffer)
+);
+assert_equals(expected_pickle, pickle_buffer, olm_pickle_pk_decryption_length(decryption));
+
+olm_clear_pk_decryption(decryption);
+
+memset(pubkey, 0, olm_pk_key_length());
+
+olm_unpickle_pk_decryption(
+    decryption,
+    PICKLE_KEY, strlen((char *)PICKLE_KEY),
+    pickle_buffer, sizeof(pickle_buffer),
+    pubkey, sizeof(pubkey)
+);
+
+assert_equals(alice_public, pubkey, olm_pk_key_length());
+
+char *ciphertext = strdup("ntk49j/KozVFtSqJXhCejg");
+const char *mac = "zpzU6BkZcNI";
+const char *ephemeral_key = "3p7bfXt9wbTTW2HC7OQ1Nz+DQ8hbeGdNrfx+FG+IK08";
+
+size_t max_plaintext_length = olm_pk_max_plaintext_length(decryption, strlen(ciphertext));
+std::uint8_t *plaintext_buffer = (std::uint8_t *) malloc(max_plaintext_length);
+
+olm_pk_decrypt(
+    decryption,
+    ephemeral_key, strlen(ephemeral_key),
+    mac, strlen(mac),
+    ciphertext, strlen(ciphertext),
+    plaintext_buffer, max_plaintext_length
+);
+
+const std::uint8_t *plaintext = (std::uint8_t *) "This is a test";
+
+assert_equals(plaintext, plaintext_buffer, strlen((const char *)plaintext));
+
+free(ciphertext);
+free(plaintext_buffer);
+
+}
 }