Implement pickling/unpickling for outbound group sessions

2 files changed, 100 insertions, 2 deletions

diff --git a/src/megolm.c b/src/megolm.c
index 58fe725..7567894 100644
--- a/src/megolm.c
+++ b/src/megolm.c
@@ -20,6 +20,7 @@
 
 #include "olm/cipher.h"
 #include "olm/crypto.h"
+#include "olm/pickle.h"
 
 const struct _olm_cipher *megolm_cipher() {
     static const uint8_t CIPHER_KDF_INFO[] = "MEGOLM_KEYS";
@@ -59,12 +60,32 @@ static void rehash_part(
 
 
 
-void megolm_init(Megolm *megolm, uint8_t const *random_data, uint32_t counter)
-{
+void megolm_init(Megolm *megolm, uint8_t const *random_data, uint32_t counter) {
     megolm->counter = counter;
     memcpy(megolm->data, random_data, MEGOLM_RATCHET_LENGTH);
 }
 
+size_t megolm_pickle_length(const Megolm *megolm) {
+    size_t length = 0;
+    length += _olm_pickle_bytes_length(megolm_get_data(megolm), MEGOLM_RATCHET_LENGTH);
+    length += _olm_pickle_uint32_length(megolm->counter);
+    return length;
+
+}
+
+uint8_t * megolm_pickle(const Megolm *megolm,  uint8_t *pos) {
+    pos = _olm_pickle_bytes(pos, megolm_get_data(megolm), MEGOLM_RATCHET_LENGTH);
+    pos = _olm_pickle_uint32(pos, megolm->counter);
+    return pos;
+}
+
+const uint8_t * megolm_unpickle(Megolm *megolm, const uint8_t *pos,
+                                const uint8_t *end) {
+    pos = _olm_unpickle_bytes(pos, end, (uint8_t *)(megolm->data),
+                             MEGOLM_RATCHET_LENGTH);
+    pos = _olm_unpickle_uint32(pos, end, &megolm->counter);
+    return pos;
+}
 
 /* simplistic implementation for a single step */
 void megolm_advance(Megolm *megolm) {
diff --git a/src/outbound_group_session.c b/src/outbound_group_session.c
index a23f684..8dc1cd1 100644
--- a/src/outbound_group_session.c
+++ b/src/outbound_group_session.c
@@ -23,10 +23,13 @@
 #include "olm/error.h"
 #include "olm/megolm.h"
 #include "olm/message.h"
+#include "olm/pickle.h"
+#include "olm/pickle_encoding.h"
 
 #define OLM_PROTOCOL_VERSION     3
 #define SESSION_ID_RANDOM_BYTES  4
 #define GROUP_SESSION_ID_LENGTH (sizeof(struct timeval) + SESSION_ID_RANDOM_BYTES)
+#define PICKLE_VERSION           1
 
 struct OlmOutboundGroupSession {
     /** the Megolm ratchet providing the encryption keys */
@@ -64,6 +67,80 @@ size_t olm_clear_outbound_group_session(
     return sizeof(OlmOutboundGroupSession);
 }
 
+static size_t raw_pickle_length(
+    const OlmOutboundGroupSession *session
+) {
+    size_t length = 0;
+    length += _olm_pickle_uint32_length(PICKLE_VERSION);
+    length += megolm_pickle_length(&(session->ratchet));
+    length += _olm_pickle_bytes_length(session->session_id,
+                                       GROUP_SESSION_ID_LENGTH);
+    return length;
+}
+
+size_t olm_pickle_outbound_group_session_length(
+    const OlmOutboundGroupSession *session
+) {
+    return _olm_enc_output_length(raw_pickle_length(session));
+}
+
+size_t olm_pickle_outbound_group_session(
+    OlmOutboundGroupSession *session,
+    void const * key, size_t key_length,
+    void * pickled, size_t pickled_length
+) {
+    size_t raw_length = raw_pickle_length(session);
+    uint8_t *pos;
+
+    if (pickled_length < _olm_enc_output_length(raw_length)) {
+        session->last_error = OLM_OUTPUT_BUFFER_TOO_SMALL;
+        return (size_t)-1;
+    }
+
+    pos = _olm_enc_output_pos(pickled, raw_length);
+    pos = _olm_pickle_uint32(pos, PICKLE_VERSION);
+    pos = megolm_pickle(&(session->ratchet), pos);
+    pos = _olm_pickle_bytes(pos, session->session_id, GROUP_SESSION_ID_LENGTH);
+
+    return _olm_enc_output(key, key_length, pickled, raw_length);
+}
+
+size_t olm_unpickle_outbound_group_session(
+    OlmOutboundGroupSession *session,
+    void const * key, size_t key_length,
+    void * pickled, size_t pickled_length
+) {
+    const uint8_t *pos;
+    const uint8_t *end;
+    uint32_t pickle_version;
+
+    size_t raw_length = _olm_enc_input(
+        key, key_length, pickled, pickled_length, &(session->last_error)
+    );
+    if (raw_length == (size_t)-1) {
+        return raw_length;
+    }
+
+    pos = pickled;
+    end = pos + raw_length;
+    pos = _olm_unpickle_uint32(pos, end, &pickle_version);
+    if (pickle_version != PICKLE_VERSION) {
+        session->last_error = OLM_UNKNOWN_PICKLE_VERSION;
+        return (size_t)-1;
+    }
+    pos = megolm_unpickle(&(session->ratchet), pos, end);
+    pos = _olm_unpickle_bytes(pos, end, session->session_id, GROUP_SESSION_ID_LENGTH);
+
+    if (end != pos) {
+        /* We had the wrong number of bytes in the input. */
+        session->last_error = OLM_CORRUPTED_PICKLE;
+        return (size_t)-1;
+    }
+
+    return pickled_length;
+}
+
+
 size_t olm_init_outbound_group_session_random_length(
     const OlmOutboundGroupSession *session
 ) {