1 files changed, 102 insertions, 0 deletions
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