Encode the account keys as a signed JSON object

9 files changed, 444 insertions, 73 deletions

diff --git a/include/olm/account.hh b/include/olm/account.hh
index d3bde0d..1a4778b 100644
--- a/include/olm/account.hh
+++ b/include/olm/account.hh
@@ -24,7 +24,12 @@
 namespace olm {
 
 
-struct LocalKey {
+struct IdentityKeys {
+    Ed25519KeyPair ed25519_key;
+    Curve25519KeyPair curve25519_key;
+};
+
+struct OneTimeKey {
     std::uint32_t id;
     Curve25519KeyPair key;
 };
@@ -34,8 +39,8 @@ static std::size_t const MAX_ONE_TIME_KEYS = 100;
 
 
 struct Account {
-    LocalKey identity_key;
-    List<LocalKey, MAX_ONE_TIME_KEYS> one_time_keys;
+    IdentityKeys identity_keys;
+    List<OneTimeKey, MAX_ONE_TIME_KEYS> one_time_keys;
     ErrorCode last_error;
 
     /** Number of random bytes needed to create a new account */
@@ -47,7 +52,54 @@ struct Account {
         uint8_t const * random, std::size_t random_length
     );
 
-    LocalKey const * lookup_key(
+    /** Number of bytes needed to output the identity keys for this account */
+    std::size_t get_identity_json_length(
+        std::size_t user_id_length,
+        std::size_t device_id_length,
+        std::uint64_t valid_after_ts,
+        std::uint64_t valid_until_ts
+    );
+
+    /** Output the identity keys for this account as JSON in the following
+     * format.
+     *
+     *  14 "{\"algorithms\":"
+     *  30 "[\"m.olm.curve25519-aes-sha256\""
+     *  15 "],\"device_id\":\""
+     *   ? <device identifier>
+     *  22 "\",\"keys\":{\"curve25519:"
+     *   4 <base64 characters>
+     *   3 "\":\""
+     *  43 <base64 characters>
+     *  11 "\",\"ed25519:"
+     *   4 <base64 characters>
+     *   3 "\":\""
+     *  43 <base64 characters>
+     *  14 "\"},\"user_id\":\""
+     *   ? <user identifier>
+     *  19 "\",\"valid_after_ts\":"
+     *   ? <digits>
+     *  18 ",\"valid_until_ts\":"
+     *   ? <digits>
+     *  16 ",\"signatures\":{\""
+     *   ? <user identifier>
+     *   1 "/"
+     *   ? <device identifier>
+     *  12 "\":{\"ed25519:"
+     *   4 <base64 characters>
+     *   3 "\":\""
+     *  86 <base64 characters>
+     *   4 "\"}}}"
+     */
+    std::size_t get_identity_json(
+        std::uint8_t const * user_id, std::size_t user_id_length,
+        std::uint8_t const * device_id, std::size_t device_id_length,
+        std::uint64_t valid_after_ts,
+        std::uint64_t valid_until_ts,
+        std::uint8_t * identity_keys, std::size_t identity_keys_length
+    );
+
+    OneTimeKey const * lookup_key(
         std::uint32_t id
     );
 
diff --git a/include/olm/olm.hh b/include/olm/olm.hh
index fca35c4..ab71689 100644
--- a/include/olm/olm.hh
+++ b/include/olm/olm.hh
@@ -16,6 +16,7 @@
 #define OLM_HH_
 
 #include <stddef.h>
+#include <stdint.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -131,18 +132,23 @@ size_t olm_create_account(
 
 /** The size of the output buffer needed to hold the identity keys */
 size_t olm_account_identity_keys_length(
-    OlmAccount * account
+    OlmAccount * account,
+    size_t user_id_length,
+    size_t device_id_length,
+    uint64_t valid_after_ts,
+    uint64_t valid_until_ts
 );
 
 /** Writes the public parts of the identity keys for the account into the
- * identity_keys output buffer. The output is formatted as though it was
- * created with sprintf(output, "[[%10d,\"%43s\"]\n]", key_id, key_base64).
- * The output can either be parsed as fixed width using the above format or by
- * a JSON parser. Returns olm_error() on failure. If the identity_keys
- * buffer was too small then olm_account_last_error() will be
+ * identity_keys output buffer. Returns olm_error() on failure. If the
+ * identity_keys buffer was too small then olm_account_last_error() will be
  * "OUTPUT_BUFFER_TOO_SMALL". */
 size_t olm_account_identity_keys(
     OlmAccount * account,
+    void const * user_id, size_t user_id_length,
+    void const * device_id, size_t device_id_length,
+    uint64_t valid_after_ts,
+    uint64_t valid_until_ts,
     void * identity_keys, size_t identity_key_length
 );
 
diff --git a/include/olm/pickle.hh b/include/olm/pickle.hh
index 1676e23..7a2bd1b 100644
--- a/include/olm/pickle.hh
+++ b/include/olm/pickle.hh
@@ -170,6 +170,39 @@ std::uint8_t const * unpickle(
 );
 
 
+std::size_t pickle_length(
+    const Ed25519PublicKey & value
+);
+
+
+std::uint8_t * pickle(
+    std::uint8_t * pos,
+    const Ed25519PublicKey & value
+);
+
+
+std::uint8_t const * unpickle(
+    std::uint8_t const * pos, std::uint8_t const * end,
+    Ed25519PublicKey & value
+);
+
+
+std::size_t pickle_length(
+    const Ed25519KeyPair & value
+);
+
+
+std::uint8_t * pickle(
+    std::uint8_t * pos,
+    const Ed25519KeyPair & value
+);
+
+
+std::uint8_t const * unpickle(
+    std::uint8_t const * pos, std::uint8_t const * end,
+    Ed25519KeyPair & value
+);
+
 } // namespace olm
 
 
diff --git a/olm.py b/olm.py
index a481ac7..a77ae63 100755
--- a/olm.py
+++ b/olm.py
@@ -45,8 +45,15 @@ account_function(
 )
 account_function(lib.olm_create_account_random_length)
 account_function(lib.olm_create_account, c_void_p, c_size_t)
-account_function(lib.olm_account_identity_keys_length)
-account_function(lib.olm_account_identity_keys, c_void_p, c_size_t)
+account_function(
+    lib.olm_account_identity_keys_length,
+    c_size_t, c_size_t, c_uint64, c_uint64
+)
+account_function(
+    lib.olm_account_identity_keys,
+    c_void_p, c_size_t, c_void_p, c_size_t, c_uint64, c_uint64,
+    c_void_p, c_size_t
+)
 account_function(lib.olm_account_one_time_keys_length)
 account_function(lib.olm_account_one_time_keys, c_void_p, c_size_t)
 
@@ -81,10 +88,20 @@ class Account(object):
             self.ptr, key_buffer, len(key), pickle_buffer, len(pickle)
         )
 
-    def identity_keys(self):
-        out_length = lib.olm_account_identity_keys_length(self.ptr)
+    def identity_keys(self, user_id, device_id, valid_after, valid_until):
+        out_length = lib.olm_account_identity_keys_length(
+            self.ptr, len(user_id), len(device_id), valid_after, valid_until
+        )
+        user_id_buffer = create_string_buffer(user_id)
+        device_id_buffer = create_string_buffer(device_id)
         out_buffer = create_string_buffer(out_length)
-        lib.olm_account_identity_keys(self.ptr, out_buffer, out_length)
+        lib.olm_account_identity_keys(
+            self.ptr,
+            user_id_buffer, len(user_id),
+            device_id_buffer, len(device_id),
+            valid_after, valid_until,
+            out_buffer, out_length
+        )
         return json.loads(out_buffer.raw)
 
     def one_time_keys(self):
@@ -280,33 +297,31 @@ if __name__ == '__main__':
     create_account.set_defaults(func=do_create_account)
 
     keys = commands.add_parser("keys", help="List public keys for an account")
-    keys.add_argument("account_file", help="Local account_file")
+    keys.add_argument("--user-id", default="A User ID")
+    keys.add_argument("--device-id", default="A Device ID")
+    keys.add_argument("--valid-after", default=0, type=int)
+    keys.add_argument("--valid-until", default=0, type=int)
+    keys.add_argument("account_file", help="Local account file")
 
     def do_keys(args):
         account = Account()
         with open(args.account_file, "rb") as f:
             account.unpickle(args.key, f.read())
-        (r_id, id_key), (signed_id, signed_key) = account.identity_keys()
-        ot_keys = account.one_time_keys()
         result1 = {
-            "identityKey": str(id_key),
-            "signedKey": {
-                "keyId": signed_id,
-                "publicKey": str(signed_key),
-            },
-            "lastResortKey": {
-                "keyId": ot_keys[0][0],
-                "publicKey": str(ot_keys[0][1]),
-            },
+            "device_keys": account.identity_keys(
+                args.user_id, args.device_id,
+                args.valid_after, args.valid_until,
+            )
         }
+        ot_keys = account.one_time_keys()
         result2 = {
-            "oneTimeKeys": [{
+            "one_time_keys": [{
                 "keyId": k[0],
                 "publicKey": str(k[1]),
             } for k in ot_keys[1:]]
         }
         try:
-            yaml.dump(result1, sys.stdout, default_flow_style=False)
+            yaml.safe_dump(result1, sys.stdout, default_flow_style=False)
             yaml.dump(result2, sys.stdout, default_flow_style=False)
         except:
             pass
diff --git a/src/account.cpp b/src/account.cpp
index 198c5d3..1ca3b13 100644
--- a/src/account.cpp
+++ b/src/account.cpp
@@ -13,13 +13,14 @@
  * limitations under the License.
  */
 #include "olm/account.hh"
+#include "olm/base64.hh"
 #include "olm/pickle.hh"
 
 
-olm::LocalKey const * olm::Account::lookup_key(
+olm::OneTimeKey const * olm::Account::lookup_key(
     std::uint32_t id
 ) {
-    for (olm::LocalKey const & key : one_time_keys) {
+    for (olm::OneTimeKey const & key : one_time_keys) {
         if (key.id == id) return &key;
     }
     return 0;
@@ -28,7 +29,7 @@ olm::LocalKey const * olm::Account::lookup_key(
 std::size_t olm::Account::remove_key(
     std::uint32_t id
 ) {
-    LocalKey * i;
+    OneTimeKey * i;
     for (i = one_time_keys.begin(); i != one_time_keys.end(); ++i) {
         if (i->id == id) {
             one_time_keys.erase(i);
@@ -52,14 +53,13 @@ std::size_t olm::Account::new_account(
 
     unsigned id = 0;
 
-    identity_key.id = ++id;
-    olm::curve25519_generate_key(random, identity_key.key);
+    olm::ed25519_generate_key(random, identity_keys.ed25519_key);
     random += 32;
-
+    olm::curve25519_generate_key(random, identity_keys.curve25519_key);
     random += 32;
 
     for (unsigned i = 0; i < 10; ++i) {
-        LocalKey & key = *one_time_keys.insert(one_time_keys.end());
+        OneTimeKey & key = *one_time_keys.insert(one_time_keys.end());
         key.id = ++id;
         olm::curve25519_generate_key(random, key.key);
         random += 32;
@@ -68,12 +68,200 @@ std::size_t olm::Account::new_account(
     return 0;
 }
 
+namespace {
+
+static const uint8_t IDENTITY_JSON_PART_0[] =
+    "{\"algorithms\":"
+    "[\"m.olm.curve25519-aes-sha256\""
+    "],\"device_id\":\"";
+static const uint8_t IDENTITY_JSON_PART_1[] = "\",\"keys\":{\"curve25519:";
+static const uint8_t IDENTITY_JSON_PART_2[] = "\":\"";
+static const uint8_t IDENTITY_JSON_PART_3[] = "\",\"ed25519:";
+static const uint8_t IDENTITY_JSON_PART_4[] = "\":\"";
+static const uint8_t IDENTITY_JSON_PART_5[] = "\"},\"user_id\":\"";
+static const uint8_t IDENTITY_JSON_PART_6[] = "\",\"valid_after_ts\":";
+static const uint8_t IDENTITY_JSON_PART_7[] = ",\"valid_until_ts\":";
+static const uint8_t IDENTITY_JSON_PART_8[] = ",\"signatures\":{\"";
+static const uint8_t IDENTITY_JSON_PART_9[] = "/";
+static const uint8_t IDENTITY_JSON_PART_A[] = "\":{\"ed25519:";
+static const uint8_t IDENTITY_JSON_PART_B[] = "\":\"";
+static const uint8_t IDENTITY_JSON_PART_C[] = "\"}}}";
+
+std::size_t count_digits(
+    std::uint64_t value
+) {
+    std::size_t digits = 0;
+    do {
+        digits++;
+        value /= 10;
+    } while (value);
+    return digits;
+}
+
+template<typename T>
+std::uint8_t * write_string(
+    std::uint8_t * pos,
+    T const & value
+) {
+    std::memcpy(pos, value, sizeof(T) - 1);
+    return pos + (sizeof(T) - 1);
+}
+
+std::uint8_t * write_string(
+    std::uint8_t * pos,
+    std::uint8_t const * value, std::size_t value_length
+) {
+    std::memcpy(pos, value, value_length);
+    return pos + value_length;
+}
+
+std::uint8_t * write_digits(
+    std::uint8_t * pos,
+    std::uint64_t value
+) {
+    size_t digits = count_digits(value);
+    pos += digits;
+    do {
+        *(--pos) = '0' + (value % 10);
+        value /= 10;
+    } while (value);
+    return pos + digits;
+}
+
+}
+
+
+std::size_t olm::Account::get_identity_json_length(
+    std::size_t user_id_length,
+    std::size_t device_id_length,
+    std::uint64_t valid_after_ts,
+    std::uint64_t valid_until_ts
+) {
+    std::size_t length = 0;
+    length += sizeof(IDENTITY_JSON_PART_0) - 1;
+    length += device_id_length;
+    length += sizeof(IDENTITY_JSON_PART_1) - 1;
+    length += 4;
+    length += sizeof(IDENTITY_JSON_PART_2) - 1;
+    length += 43;
+    length += sizeof(IDENTITY_JSON_PART_3) - 1;
+    length += 4;
+    length += sizeof(IDENTITY_JSON_PART_4) - 1;
+    length += 43;
+    length += sizeof(IDENTITY_JSON_PART_5) - 1;
+    length += user_id_length;
+    length += sizeof(IDENTITY_JSON_PART_6) - 1;
+    length += count_digits(valid_after_ts);
+    length += sizeof(IDENTITY_JSON_PART_7) - 1;
+    length += count_digits(valid_until_ts);
+    length += sizeof(IDENTITY_JSON_PART_8) - 1;
+    length += user_id_length;
+    length += sizeof(IDENTITY_JSON_PART_9) - 1;
+    length += device_id_length;
+    length += sizeof(IDENTITY_JSON_PART_A) - 1;
+    length += 4;
+    length += sizeof(IDENTITY_JSON_PART_B) - 1;
+    length += 86;
+    length += sizeof(IDENTITY_JSON_PART_C) - 1;
+    return length;
+}
+
+
+std::size_t olm::Account::get_identity_json(
+    std::uint8_t const * user_id, std::size_t user_id_length,
+    std::uint8_t const * device_id, std::size_t device_id_length,
+    std::uint64_t valid_until_ts,
+    std::uint64_t valid_after_ts,
+    std::uint8_t * identity_json, std::size_t identity_json_length
+) {
+
+    std::uint8_t * pos = identity_json;
+    std::uint8_t signature[64];
+    size_t expected_length = get_identity_json_length(
+            user_id_length, device_id_length, valid_after_ts, valid_until_ts
+    );
+
+    if (identity_json_length < expected_length) {
+        last_error = olm::ErrorCode::OUTPUT_BUFFER_TOO_SMALL;
+        return std::size_t(-1);
+    }
+
+    pos = write_string(pos, IDENTITY_JSON_PART_0);
+    pos = write_string(pos, device_id, device_id_length);
+    pos = write_string(pos, IDENTITY_JSON_PART_1);
+    encode_base64(identity_keys.curve25519_key.public_key, 3, pos);
+    pos += 4;
+    pos = write_string(pos, IDENTITY_JSON_PART_2);
+    encode_base64(identity_keys.curve25519_key.public_key, 32, pos);
+    pos += 43;
+    pos = write_string(pos, IDENTITY_JSON_PART_3);
+    encode_base64(identity_keys.ed25519_key.public_key, 3, pos);
+    pos += 4;
+    pos = write_string(pos, IDENTITY_JSON_PART_4);
+    encode_base64(identity_keys.ed25519_key.public_key, 32, pos);
+    pos += 43;
+    pos = write_string(pos, IDENTITY_JSON_PART_5);
+    pos = write_string(pos, user_id, user_id_length);
+    pos = write_string(pos, IDENTITY_JSON_PART_6);
+    pos = write_digits(pos, valid_after_ts);
+    pos = write_string(pos, IDENTITY_JSON_PART_7);
+    pos = write_digits(pos, valid_until_ts);
+    *pos = '}';
+    // Sign the JSON up to written up to this point.
+    ed25519_sign(
+        identity_keys.ed25519_key,
+        identity_json, 1 + pos - identity_json,
+        signature
+    );
+    // Append the signature to the end of the JSON.
+    pos = write_string(pos, IDENTITY_JSON_PART_8);
+    pos = write_string(pos, user_id, user_id_length);
+    pos = write_string(pos, IDENTITY_JSON_PART_9);
+    pos = write_string(pos, device_id, device_id_length);
+    pos = write_string(pos, IDENTITY_JSON_PART_A);
+    encode_base64(identity_keys.ed25519_key.public_key, 3, pos);
+    pos += 4;
+    pos = write_string(pos, IDENTITY_JSON_PART_B);
+    encode_base64(signature, 64, pos);
+    pos += 86;
+    pos = write_string(pos, IDENTITY_JSON_PART_C);
+    return pos - identity_json;
+}
 
 namespace olm {
 
+static std::size_t pickle_length(
+    olm::IdentityKeys const & value
+) {
+    size_t length = 0;
+    length += olm::pickle_length(value.ed25519_key);
+    length += olm::pickle_length(value.curve25519_key);
+    return length;
+}
+
+
+static std::uint8_t * pickle(
+    std::uint8_t * pos,
+    olm::IdentityKeys const & value
+) {
+    pos = olm::pickle(pos, value.ed25519_key);
+    pos = olm::pickle(pos, value.curve25519_key);
+    return pos;
+}
+
+
+static std::uint8_t const * unpickle(
+    std::uint8_t const * pos, std::uint8_t const * end,
+    olm::IdentityKeys & value
+) {
+    pos = olm::unpickle(pos, end, value.ed25519_key);
+    pos = olm::unpickle(pos, end, value.curve25519_key);
+    return pos;
+}
+
 
 static std::size_t pickle_length(
-    olm::LocalKey const & value
+    olm::OneTimeKey const & value
 ) {
     return olm::pickle_length(value.id) + olm::pickle_length(value.key);
 }
@@ -81,7 +269,7 @@ static std::size_t pickle_length(
 
 static std::uint8_t * pickle(
     std::uint8_t * pos,
-    olm::LocalKey const & value
+    olm::OneTimeKey const & value
 ) {
     pos = olm::pickle(pos, value.id);
     pos = olm::pickle(pos, value.key);
@@ -91,7 +279,7 @@ static std::uint8_t * pickle(
 
 static std::uint8_t const * unpickle(
     std::uint8_t const * pos, std::uint8_t const * end,
-    olm::LocalKey & value
+    olm::OneTimeKey & value
 ) {
     pos = olm::unpickle(pos, end, value.id);
     pos = olm::unpickle(pos, end, value.key);
@@ -105,7 +293,7 @@ std::size_t olm::pickle_length(
     olm::Account const & value
 ) {
     std::size_t length = 0;
-    length += olm::pickle_length(value.identity_key);
+    length += olm::pickle_length(value.identity_keys);
     length += olm::pickle_length(value.one_time_keys);
     return length;
 }
@@ -115,7 +303,7 @@ std::uint8_t * olm::pickle(
     std::uint8_t * pos,
     olm::Account const & value
 ) {
-    pos = olm::pickle(pos, value.identity_key);
+    pos = olm::pickle(pos, value.identity_keys);
     pos = olm::pickle(pos, value.one_time_keys);
     return pos;
 }
@@ -125,7 +313,7 @@ std::uint8_t const * olm::unpickle(
     std::uint8_t const * pos, std::uint8_t const * end,
     olm::Account & value
 ) {
-    pos = olm::unpickle(pos, end, value.identity_key);
+    pos = olm::unpickle(pos, end, value.identity_keys);
     pos = olm::unpickle(pos, end, value.one_time_keys);
     return pos;
 }
diff --git a/src/olm.cpp b/src/olm.cpp
index 3aab4e2..a9735a5 100644
--- a/src/olm.cpp
+++ b/src/olm.cpp
@@ -325,8 +325,8 @@ namespace {
 static const std::size_t OUTPUT_KEY_LENGTH = 2 + 10 + 2 +
         olm::encode_base64_length(32) + 3;
 
-void output_key(
-    olm::LocalKey const & key,
+void output_one_time_key(
+    olm::OneTimeKey const & key,
     std::uint8_t sep,
     std::uint8_t * output
 ) {
@@ -353,27 +353,35 @@ void output_key(
 
 
 size_t olm_account_identity_keys_length(
-    OlmAccount * account
-) {
-    return OUTPUT_KEY_LENGTH + 1;
+    OlmAccount * account,
+    size_t user_id_length,
+    size_t device_id_length,
+    uint64_t valid_after_ts,
+    uint64_t valid_until_ts
+) {
+    return from_c(account)->get_identity_json_length(
+        user_id_length,
+        device_id_length,
+        valid_after_ts,
+        valid_until_ts
+    );
 }
 
-
 size_t olm_account_identity_keys(
     OlmAccount * account,
+    void const * user_id, size_t user_id_length,
+    void const * device_id, size_t device_id_length,
+    uint64_t valid_after_ts,
+    uint64_t valid_until_ts,
     void * identity_keys, size_t identity_key_length
 ) {
-    std::size_t length = olm_account_identity_keys_length(account);
-    if (identity_key_length < length) {
-        from_c(account)->last_error =
-            olm::ErrorCode::OUTPUT_BUFFER_TOO_SMALL;
-        return size_t(-1);
-    }
-    std::uint8_t * output = from_c(identity_keys);
-    output_key(from_c(account)->identity_key, '[', output);
-    output += OUTPUT_KEY_LENGTH;
-    output[0] = ']';
-    return length;
+    return from_c(account)->get_identity_json(
+        from_c(user_id), user_id_length,
+        from_c(device_id), device_id_length,
+        valid_after_ts,
+        valid_until_ts,
+        from_c(identity_keys), identity_key_length
+    );
 }
 
 
@@ -381,7 +389,7 @@ size_t olm_account_one_time_keys_length(
     OlmAccount * account
 ) {
     size_t count = from_c(account)->one_time_keys.size();
-    return OUTPUT_KEY_LENGTH * (count + 1) + 1;
+    return OUTPUT_KEY_LENGTH * count + 1;
 }
 
 
@@ -397,9 +405,8 @@ size_t olm_account_one_time_keys(
     }
     std::uint8_t * output = from_c(identity_keys);
     std::uint8_t sep = '[';
-    output += OUTPUT_KEY_LENGTH;
     for (auto const & key : from_c(account)->one_time_keys) {
-        output_key(key, sep, output);
+        output_one_time_key(key, sep, output);
         output += OUTPUT_KEY_LENGTH;
         sep = ',';
     }
diff --git a/src/pickle.cpp b/src/pickle.cpp
index 3aa64ed..25580d2 100644
--- a/src/pickle.cpp
+++ b/src/pickle.cpp
@@ -78,3 +78,67 @@ std::uint8_t const * olm::unpickle(
     );
     return pos;
 }
+
+std::size_t olm::pickle_length(
+    const olm::Ed25519PublicKey & value
+) {
+    return sizeof(value.public_key);
+}
+
+
+std::uint8_t * olm::pickle(
+    std::uint8_t * pos,
+    const olm::Ed25519PublicKey & value
+) {
+    pos = olm::pickle_bytes(
+        pos, value.public_key, sizeof(value.public_key)
+    );
+    return pos;
+}
+
+
+std::uint8_t const * olm::unpickle(
+    std::uint8_t const * pos, std::uint8_t const * end,
+    olm::Ed25519PublicKey & value
+) {
+    pos = olm::unpickle_bytes(
+        pos, end, value.public_key, sizeof(value.public_key)
+    );
+    return pos;
+
+}
+
+
+std::size_t olm::pickle_length(
+    const olm::Ed25519KeyPair & value
+) {
+    return sizeof(value.public_key) + sizeof(value.private_key);
+}
+
+
+std::uint8_t * olm::pickle(
+    std::uint8_t * pos,
+    const olm::Ed25519KeyPair & value
+) {
+    pos = olm::pickle_bytes(
+        pos, value.public_key, sizeof(value.public_key)
+    );
+    pos = olm::pickle_bytes(
+        pos, value.private_key, sizeof(value.private_key)
+    );
+    return pos;
+}
+
+
+std::uint8_t const * olm::unpickle(
+    std::uint8_t const * pos, std::uint8_t const * end,
+    olm::Ed25519KeyPair & value
+) {
+    pos = olm::unpickle_bytes(
+        pos, end, value.public_key, sizeof(value.public_key)
+    );
+    pos = olm::unpickle_bytes(
+        pos, end, value.private_key, sizeof(value.private_key)
+    );
+    return pos;
+}
diff --git a/src/session.cpp b/src/session.cpp
index 08498e4..0457042 100644
--- a/src/session.cpp
+++ b/src/session.cpp
@@ -74,15 +74,16 @@ std::size_t olm::Session::new_outbound_session(
     olm::curve25519_generate_key(random + 32, ratchet_key);
 
     received_message = false;
-    alice_identity_key.id = local_account.identity_key.id;
-    alice_identity_key.key = local_account.identity_key.key;
+    alice_identity_key.id = 0;
+    alice_identity_key.key = local_account.identity_keys.curve25519_key;
     alice_base_key = base_key;
     bob_one_time_key_id = one_time_key.id;
 
     std::uint8_t shared_secret[96];
 
     olm::curve25519_shared_secret(
-        local_account.identity_key.key, one_time_key.key, shared_secret
+        local_account.identity_keys.curve25519_key,
+        one_time_key.key, shared_secret
     );
     olm::curve25519_shared_secret(
          base_key, identity_key, shared_secret + 32
@@ -148,7 +149,7 @@ std::size_t olm::Session::new_inbound_session(
     olm::Curve25519PublicKey ratchet_key;
     std::memcpy(ratchet_key.public_key, message_reader.ratchet_key, 32);
 
-    olm::LocalKey const * bob_one_time_key = local_account.lookup_key(
+    olm::OneTimeKey const * bob_one_time_key = local_account.lookup_key(
         bob_one_time_key_id
     );
 
@@ -163,7 +164,8 @@ std::size_t olm::Session::new_inbound_session(
         bob_one_time_key->key, alice_identity_key.key, shared_secret
     );
     olm::curve25519_shared_secret(
-        local_account.identity_key.key, alice_base_key, shared_secret + 32
+        local_account.identity_keys.curve25519_key,
+        alice_base_key, shared_secret + 32
     );
     olm::curve25519_shared_secret(
         bob_one_time_key->key, alice_base_key, shared_secret + 64
diff --git a/tests/test_olm.cpp b/tests/test_olm.cpp
index 74d19e7..8ea2462 100644
--- a/tests/test_olm.cpp
+++ b/tests/test_olm.cpp
@@ -75,9 +75,11 @@ std::uint8_t b_random[::olm_create_account_random_length(b_account)];
 mock_random_b(b_random, sizeof(b_random));
 ::olm_create_account(b_account, b_random, sizeof(b_random));
 
-std::uint8_t b_id_keys[::olm_account_identity_keys_length(b_account)];
+std::uint8_t b_id_keys[::olm_account_identity_keys_length(b_account, 0, 0, 0, 0)];
 std::uint8_t b_ot_keys[::olm_account_one_time_keys_length(b_account)];
-::olm_account_identity_keys(b_account, b_id_keys, sizeof(b_id_keys));
+::olm_account_identity_keys(
+    b_account, nullptr, 0, nullptr, 0, 0, 0, b_id_keys, sizeof(b_id_keys)
+);
 ::olm_account_one_time_keys(b_account, b_ot_keys, sizeof(b_ot_keys));
 
 std::uint8_t a_session_buffer[::olm_session_size()];
@@ -86,7 +88,7 @@ std::uint8_t a_rand[::olm_create_outbound_session_random_length(a_session)];
 mock_random_a(a_rand, sizeof(a_rand));
 assert_not_equals(std::size_t(-1), ::olm_create_outbound_session(
     a_session, a_account,
-    b_id_keys + 14, 43,
+    b_id_keys + 88, 43,
     ::atol((char *)(b_ot_keys + 62)), b_ot_keys + 74, 43,
     a_rand, sizeof(a_rand)
 ));
@@ -177,9 +179,11 @@ std::uint8_t b_random[::olm_create_account_random_length(b_account)];
 mock_random_b(b_random, sizeof(b_random));
 ::olm_create_account(b_account, b_random, sizeof(b_random));
 
-std::uint8_t b_id_keys[::olm_account_identity_keys_length(b_account)];
+std::uint8_t b_id_keys[::olm_account_identity_keys_length(b_account, 0, 0, 0, 0)];
 std::uint8_t b_ot_keys[::olm_account_one_time_keys_length(b_account)];
-::olm_account_identity_keys(b_account, b_id_keys, sizeof(b_id_keys));
+::olm_account_identity_keys(
+    b_account, nullptr, 0, nullptr, 0, 0, 0, b_id_keys, sizeof(b_id_keys)
+);
 ::olm_account_one_time_keys(b_account, b_ot_keys, sizeof(b_ot_keys));
 
 std::uint8_t a_session_buffer[::olm_session_size()];
@@ -188,7 +192,7 @@ std::uint8_t a_rand[::olm_create_outbound_session_random_length(a_session)];
 mock_random_a(a_rand, sizeof(a_rand));
 assert_not_equals(std::size_t(-1), ::olm_create_outbound_session(
     a_session, a_account,
-    b_id_keys + 14, 43,
+    b_id_keys + 88, 43,
     ::atol((char *)(b_ot_keys + 62)), b_ot_keys + 74, 43,
     a_rand, sizeof(a_rand)
 ));