Enable exporting inbound group session keys

A pair of functions which allow you to export the megolm keys for an inbound
group session, so that an application can save/restore them.

5 files changed, 185 insertions, 20 deletions

diff --git a/include/olm/inbound_group_session.h b/include/olm/inbound_group_session.h
index f8a0bc3..d47af8a 100644
--- a/include/olm/inbound_group_session.h
+++ b/include/olm/inbound_group_session.h
@@ -165,6 +165,37 @@ size_t olm_inbound_group_session_id(
     uint8_t * id, size_t id_length
 );
 
+/**
+ * Get the first message index we know how to decrypt.
+ */
+uint32_t olm_inbound_group_session_first_known_index(
+    const OlmInboundGroupSession *session
+);
+
+/**
+ * Get the number of bytes returned by olm_export_inbound_group_session()
+ */
+size_t olm_export_inbound_group_session_length(
+    const OlmInboundGroupSession *session
+);
+
+/**
+ * Export the base64-encoded ratchet key for this session, at the given index,
+ * in a format which can be used by olm_import_inbound_group_session
+ *
+ * Returns the length of the ratchet key on success or olm_error() on
+ * failure. On failure last_error will be set with an error code. The
+ * last_error will be:
+ *   * OUTPUT_BUFFER_TOO_SMALL if the buffer was too small
+ *   * OLM_UNKNOWN_MESSAGE_INDEX  if we do not have a session key corresponding to the
+ *     given index (ie, it was sent before the session key was shared with
+ *     us)
+ */
+size_t olm_export_inbound_group_session(
+    OlmInboundGroupSession *session,
+    uint8_t * key, size_t key_length, uint32_t message_index
+);
+
 
 #ifdef __cplusplus
 } // extern "C"
diff --git a/javascript/olm_inbound_group_session.js b/javascript/olm_inbound_group_session.js
index 2e4727f..8721a10 100644
--- a/javascript/olm_inbound_group_session.js
+++ b/javascript/olm_inbound_group_session.js
@@ -123,4 +123,21 @@ InboundGroupSession.prototype['session_id'] = restore_stack(function() {
     return Pointer_stringify(session_id);
 });
 
+InboundGroupSession.prototype['first_known_index'] = restore_stack(function() {
+    return inbound_group_session_method(
+        Module['_olm_inbound_group_session_first_known_index']
+    )(this.ptr);
+});
+
+InboundGroupSession.prototype['export_session'] = restore_stack(function(message_index) {
+    var key_length = inbound_group_session_method(
+        Module['_olm_export_inbound_group_session_length']
+    )(this.ptr);
+    var key = stack(key_length + NULL_BYTE_PADDING_LENGTH);
+    outbound_group_session_method(Module['_olm_export_inbound_group_session'])(
+        this.ptr, key, key_length, message_index
+    );
+    return Pointer_stringify(key);
+});
+
 olm_exports['InboundGroupSession'] = InboundGroupSession;
diff --git a/python/olm/__main__.py b/python/olm/__main__.py
index eb76301..f195591 100755
--- a/python/olm/__main__.py
+++ b/python/olm/__main__.py
@@ -277,6 +277,27 @@ def build_arg_parser():
                                type=argparse.FileType('wb'), nargs='?',
                                default=sys.stdout)
     group_decrypt.set_defaults(func=do_group_decrypt)
+
+
+    export_inbound_group = commands.add_parser(
+        "export_inbound_group",
+        help="Export the keys for an inbound group session",
+    )
+    export_inbound_group.add_argument(
+        "session_file", help="Local inbound group session file",
+    )
+    export_inbound_group.add_argument(
+        "export_file", help="File to export to (default stdout)",
+        type=argparse.FileType('w'), nargs='?',
+        default=sys.stdout,
+    )
+    export_inbound_group.add_argument(
+        "--message_index",
+        help="Index to export session at. Defaults to the earliest known index",
+        type=int,
+    )
+    export_inbound_group.set_defaults(func=do_export_inbound_group)
+
     return parser
 
 def do_outbound_group(args):
@@ -333,6 +354,15 @@ def do_group_decrypt(args):
         f.write(session.pickle(args.key))
     args.plaintext_file.write(plaintext)
 
+def do_export_inbound_group(args):
+    session = InboundGroupSession()
+    session.unpickle(args.key, read_base64_file(args.session_file))
+    index = args.message_index
+    if index is None:
+        # default to first known index
+        index = session.first_known_index()
+    args.export_file.write(session.export_session(index))
+
 if __name__ == '__main__':
     parser = build_arg_parser()
     args = parser.parse_args()
diff --git a/python/olm/inbound_group_session.py b/python/olm/inbound_group_session.py
index 27a569c..67906b2 100644
--- a/python/olm/inbound_group_session.py
+++ b/python/olm/inbound_group_session.py
@@ -49,6 +49,13 @@ inbound_group_session_function(
 inbound_group_session_function(lib.olm_inbound_group_session_id_length)
 inbound_group_session_function(lib.olm_inbound_group_session_id, c_void_p, c_size_t)
 
+lib.olm_inbound_group_session_first_known_index.argtypes = (c_void_p,)
+lib.olm_inbound_group_session_first_known_index.restypes = c_uint32
+
+inbound_group_session_function(lib.olm_export_inbound_group_session_length)
+inbound_group_session_function(lib.olm_export_inbound_group_session, c_void_p, c_size_t, c_uint32)
+
+
 class InboundGroupSession(object):
     def __init__(self):
         self.buf = create_string_buffer(lib.olm_inbound_group_session_size())
@@ -95,5 +102,15 @@ class InboundGroupSession(object):
     def session_id(self):
         id_length = lib.olm_inbound_group_session_id_length(self.ptr)
         id_buffer = create_string_buffer(id_length)
-        lib.olm_inbound_group_session_id(self.ptr, id_buffer, id_length);
+        lib.olm_inbound_group_session_id(self.ptr, id_buffer, id_length)
         return id_buffer.raw
+
+    def first_known_index(self):
+        return lib.olm_inbound_group_session_first_known_index(self.ptr)
+
+    def export_session(self, message_index):
+        length = lib.olm_export_inbound_group_session_length(self.ptr)
+        buffer = create_string_buffer(length)
+        lib.olm_export_inbound_group_session(self.ptr, buffer, length,
+                                             message_index)
+        return buffer.raw
diff --git a/src/inbound_group_session.c b/src/inbound_group_session.c
index a54e55f..122c26c 100644
--- a/src/inbound_group_session.c
+++ b/src/inbound_group_session.c
@@ -32,6 +32,7 @@
 #define GROUP_SESSION_ID_LENGTH  ED25519_PUBLIC_KEY_LENGTH
 #define PICKLE_VERSION           1
 #define SESSION_KEY_VERSION      2
+#define SESSION_EXPORT_VERSION   1
 
 struct OlmInboundGroupSession {
     /** our earliest known ratchet value */
@@ -71,6 +72,9 @@ size_t olm_clear_inbound_group_session(
     return sizeof(OlmInboundGroupSession);
 }
 
+#define SESSION_EXPORT_RAW_LENGTH \
+    (1 + 4 + MEGOLM_RATCHET_LENGTH + ED25519_PUBLIC_KEY_LENGTH)
+
 #define SESSION_KEY_RAW_LENGTH \
     (1 + 4 + MEGOLM_RATCHET_LENGTH + ED25519_PUBLIC_KEY_LENGTH\
         + ED25519_SIGNATURE_LENGTH)
@@ -258,6 +262,32 @@ size_t olm_group_decrypt_max_plaintext_length(
 }
 
 /**
+ * get a copy of the megolm ratchet, advanced
+ * to the relevant index. Returns 0 on success, -1 on error
+ */
+static size_t _get_megolm(
+    OlmInboundGroupSession *session, uint32_t message_index, Megolm *result
+) {
+    /* pick a megolm instance to use. If we're at or beyond the latest ratchet
+     * value, use that */
+    if ((message_index - session->latest_ratchet.counter) < (1U << 31)) {
+        megolm_advance_to(&session->latest_ratchet, message_index);
+        *result = session->latest_ratchet;
+        return 0;
+    } else if ((message_index - session->initial_ratchet.counter) >= (1U << 31)) {
+        /* the counter is before our intial ratchet - we can't decode this. */
+        session->last_error = OLM_UNKNOWN_MESSAGE_INDEX;
+        return (size_t)-1;
+    } else {
+        /* otherwise, start from the initial megolm. Take a copy so that we
+         * don't overwrite the initial megolm */
+        *result = session->initial_ratchet;
+        megolm_advance_to(result, message_index);
+        return 0;
+    }
+}
+
+/**
  * decrypt an un-base64-ed message
  */
 static size_t _decrypt(
@@ -268,8 +298,7 @@ static size_t _decrypt(
 ) {
     struct _OlmDecodeGroupMessageResults decoded_results;
     size_t max_length, r;
-    Megolm *megolm;
-    Megolm tmp_megolm;
+    Megolm megolm;
 
     _olm_decode_group_message(
         message, message_length,
@@ -316,33 +345,21 @@ static size_t _decrypt(
         return (size_t)-1;
     }
 
-    /* pick a megolm instance to use. If we're at or beyond the latest ratchet
-     * value, use that */
-    if ((decoded_results.message_index - session->latest_ratchet.counter) < (1U << 31)) {
-        megolm = &session->latest_ratchet;
-    } else if ((decoded_results.message_index - session->initial_ratchet.counter) >= (1U << 31)) {
-        /* the counter is before our intial ratchet - we can't decode this. */
-        session->last_error = OLM_UNKNOWN_MESSAGE_INDEX;
-        return (size_t)-1;
-    } else {
-        /* otherwise, start from the initial megolm. Take a copy so that we
-         * don't overwrite the initial megolm */
-        tmp_megolm = session->initial_ratchet;
-        megolm = &tmp_megolm;
+    r = _get_megolm(session, decoded_results.message_index, &megolm);
+    if (r == (size_t)-1) {
+        return r;
     }
 
-    megolm_advance_to(megolm, decoded_results.message_index);
-
     /* now try checking the mac, and decrypting */
     r = megolm_cipher->ops->decrypt(
         megolm_cipher,
-        megolm_get_data(megolm), MEGOLM_RATCHET_LENGTH,
+        megolm_get_data(&megolm), MEGOLM_RATCHET_LENGTH,
         message, message_length,
         decoded_results.ciphertext, decoded_results.ciphertext_length,
         plaintext, max_plaintext_length
     );
 
-    _olm_unset(&tmp_megolm, sizeof(tmp_megolm));
+    _olm_unset(&megolm, sizeof(megolm));
     if (r == (size_t)-1) {
         session->last_error = OLM_BAD_MESSAGE_MAC;
         return r;
@@ -391,3 +408,56 @@ size_t olm_inbound_group_session_id(
         session->signing_key.public_key, GROUP_SESSION_ID_LENGTH, id
     );
 }
+
+uint32_t olm_inbound_group_session_first_known_index(
+    const OlmInboundGroupSession *session
+) {
+    return session->initial_ratchet.counter;
+}
+
+size_t olm_export_inbound_group_session_length(
+    const OlmInboundGroupSession *session
+) {
+    return _olm_encode_base64_length(SESSION_EXPORT_RAW_LENGTH);
+}
+
+size_t olm_export_inbound_group_session(
+    OlmInboundGroupSession *session,
+    uint8_t * key, size_t key_length, uint32_t message_index
+) {
+    uint8_t *raw;
+    uint8_t *ptr;
+    Megolm megolm;
+    size_t r;
+    size_t encoded_length = olm_export_inbound_group_session_length(session);
+
+    if (key_length < encoded_length) {
+        session->last_error = OLM_OUTPUT_BUFFER_TOO_SMALL;
+        return (size_t)-1;
+    }
+
+    r = _get_megolm(session, message_index, &megolm);
+    if (r == (size_t)-1) {
+        return r;
+    }
+
+    /* put the raw data at the end of the output buffer. */
+    raw = ptr = key + encoded_length - SESSION_EXPORT_RAW_LENGTH;
+    *ptr++ = SESSION_EXPORT_VERSION;
+
+    // Encode message index as a big endian 32-bit number.
+    for (unsigned i = 0; i < 4; i++) {
+        *ptr++ = 0xFF & (message_index >> 24); message_index <<= 8;
+    }
+
+    memcpy(ptr, megolm_get_data(&megolm), MEGOLM_RATCHET_LENGTH);
+    ptr += MEGOLM_RATCHET_LENGTH;
+
+    memcpy(
+        ptr, session->signing_key.public_key,
+        ED25519_PUBLIC_KEY_LENGTH
+    );
+    ptr += ED25519_PUBLIC_KEY_LENGTH;
+
+    return _olm_encode_base64(raw, SESSION_EXPORT_RAW_LENGTH, key);
+}