1 files changed, 170 insertions, 69 deletions

diff --git a/src/axolotl.cpp b/src/axolotl.cpp
index a1ea03f..7384d79 100644
--- a/src/axolotl.cpp
+++ b/src/axolotl.cpp
@@ -1,43 +1,53 @@
 #include "axolotl/axolotl.hh"
+#include "axolotl/message.hh"
 
+#include <cstring>
 
 namespace {
 
 std::uint8_t PROTOCOL_VERSION = 3;
 std::size_t MAC_LENGTH = 8;
-std::size_t KEY_LENGTH = Curve25519PublicKey::Length;
+std::size_t KEY_LENGTH = axolotl::Curve25519PublicKey::LENGTH;
 std::uint8_t MESSAGE_KEY_SEED[1] = {0x01};
 std::uint8_t CHAIN_KEY_SEED[1] = {0x02};
 std::size_t MAX_MESSAGE_GAP = 2000;
 
+template<typename T>
+void unset(
+    T & value
+) {
+    std::memset(&value, 0, sizeof(T));
+}
+
+
 void create_chain_key(
     axolotl::SharedKey const & root_key,
-    Curve25519KeyPair const & our_key,
-    Curve25519PublicKey const & their_key,
-    std::uint8_t const * info, std::size_t info_length,
-    SharedSecret & new_root_key,
-    ChainKey & new_chain_key
+    axolotl::Curve25519KeyPair const & our_key,
+    axolotl::Curve25519PublicKey const & their_key,
+    axolotl::KdfInfo const & info,
+    axolotl::SharedKey & new_root_key,
+    axolotl::ChainKey & new_chain_key
 ) {
-    axolotl::SharedSecret secret;
+    axolotl::SharedKey secret;
     axolotl::curve25519_shared_secret(our_key, their_key, secret);
     std::uint8_t derived_secrets[64];
     axolotl::hkdf_sha256(
         secret, sizeof(secret),
         root_key, sizeof(root_key),
-        info, info_length,
+        info.ratchet_info, info.ratchet_info_length,
         derived_secrets, sizeof(derived_secrets)
     );
     std::memcpy(new_root_key, derived_secrets, 32);
     std::memcpy(new_chain_key.key, derived_secrets + 32, 32);
     new_chain_key.index = 0;
-    std::memset(derived_secrets, 0, sizeof(derived_secrets);
-    std::memset(secret, 0, sizeof(secret));
+    unset(derived_secrets);
+    unset(secret);
 }
 
 
 void advance_chain_key(
-    ChainKey const & chain_key,
-    ChainKey & new_chain_key,
+    axolotl::ChainKey const & chain_key,
+    axolotl::ChainKey & new_chain_key
 ) {
     axolotl::hmac_sha256(
         chain_key.key, sizeof(chain_key.key),
@@ -49,11 +59,11 @@ void advance_chain_key(
 
 
 void create_message_keys(
-    ChainKey const & chain_key,
-    std::uint8_t const * info, std::size_t info_length,
-    MessageKey & message_key
+    axolotl::ChainKey const & chain_key,
+    axolotl::KdfInfo const & info,
+    axolotl::MessageKey & message_key
 ) {
-    axolotl::SharedSecret secret;
+    axolotl::SharedKey secret;
     axolotl::hmac_sha256(
         chain_key.key, sizeof(chain_key.key),
         MESSAGE_KEY_SEED, sizeof(MESSAGE_KEY_SEED),
@@ -62,45 +72,43 @@ void create_message_keys(
     std::uint8_t derived_secrets[80];
     axolotl::hkdf_sha256(
         secret, sizeof(secret),
-        root_key, sizeof(root_key),
-        info, info_length,
+        NULL, 0,
+        info.message_info, info.message_info_length,
         derived_secrets, sizeof(derived_secrets)
     );
-    std::memcpy(message_key.cipher_key, derived_secrets, 32);
+    std::memcpy(message_key.cipher_key.key, derived_secrets, 32);
     std::memcpy(message_key.mac_key, derived_secrets + 32, 32);
-    std::memcpy(message_key.iv, derived_secrets + 64, 16);
+    std::memcpy(message_key.iv.iv, derived_secrets + 64, 16);
     message_key.index = chain_key.index;
-    std::memset(derived_secrets, 0, sizeof(derived_secrets);
-    std::memset(secret, 0, sizeof(secret));
+    unset(derived_secrets);
+    unset(secret);
 }
 
 
 bool verify_mac(
-    MessageKey const & message_key,
+    axolotl::MessageKey const & message_key,
     std::uint8_t const * input,
     axolotl::MessageReader const & reader
 ) {
-    std::uint8_t mac[HMAC_SHA256_OUTPUT_LENGTH];
+    std::uint8_t mac[axolotl::HMAC_SHA256_OUTPUT_LENGTH];
     axolotl::hmac_sha256(
-        keys.mac_key, sizeof(keys.mac_key),
-        ciphertext, reader.body_length,
+        message_key.mac_key, sizeof(message_key.mac_key),
+        input, reader.body_length,
         mac
     );
 
     bool result = std::memcmp(mac, reader.mac, MAC_LENGTH) == 0;
-    std::memset(&mac, 0, HMAC_SHA256_OUTPUT_LENGTH);
+    unset(mac);
     return result;
 }
 
 
 bool verify_mac_for_existing_chain(
     axolotl::Session const & session,
-    axolotl::ReceiverChain const & chain,
+    axolotl::ChainKey const & chain,
     std::uint8_t const * input,
     axolotl::MessageReader const & reader
 ) {
-    ReceiverChain new_chain = chain;
-
     if (reader.counter < chain.index) {
         return false;
     }
@@ -110,18 +118,17 @@ bool verify_mac_for_existing_chain(
         return false;
     }
 
+    axolotl::ChainKey new_chain = chain;
+
     while (new_chain.index < reader.counter) {
         advance_chain_key(new_chain, new_chain);
     }
 
-    MessageKey message_key;
-    create_message_keys(
-        new_chain_key, sender.message_info, sender.message_info_length,
-        message_key
-    );
+    axolotl::MessageKey message_key;
+    create_message_keys(new_chain, session.kdf_info, message_key);
 
     bool result = verify_mac(message_key, input, reader);
-    std::memset(&new_chain, 0, sizeof(new_chain.ratchet_key);
+    unset(new_chain);
     return result;
 }
 
@@ -131,8 +138,8 @@ bool verify_mac_for_new_chain(
     std::uint8_t const * input,
     axolotl::MessageReader const & reader
 ) {
-    SharedSecret new_root_key;
-    ReceiverChain new_chain;
+    axolotl::SharedKey new_root_key;
+    axolotl::ReceiverChain new_chain;
 
     /* They shouldn't move to a new chain until we've sent them a message
      * acknowledging the last one */
@@ -144,30 +151,78 @@ bool verify_mac_for_new_chain(
     if (reader.counter > MAX_MESSAGE_GAP) {
         return false;
     }
-    std::memcpy(new_chain.ratchet_key, reader.ratchet_key, KEY_LENGTH);
+    std::memcpy(
+        new_chain.ratchet_key.public_key, reader.ratchet_key, KEY_LENGTH
+    );
 
     create_chain_key(
-        root_key, sender_chain[0].ratchet_key, new_chain.ratchet_key,
-        session.kdf_info.ratchet_info, session.kdf_info.ratchet_info_length,
-        new_root_key, new_chain
+        session.root_key, session.sender_chain[0].ratchet_key,
+        new_chain.ratchet_key, session.kdf_info,
+        new_root_key, new_chain.chain_key
     );
 
     bool result = verify_mac_for_existing_chain(
-        session, new_chain, input, reader
+        session, new_chain.chain_key, input, reader
     );
-    std::memset(&new_root_key, 0, sizeof(new_root_key));
-    std::memset(&new_chain, 0, sizeof(new_chain.ratchet_key);
+    unset(new_root_key);
+    unset(new_chain);
     return result;
 }
 
 } // namespace
 
 
+axolotl::Session::Session(
+    axolotl::KdfInfo const & kdf_info
+) : kdf_info(kdf_info), last_error(axolotl::ErrorCode::SUCCESS) {
+}
+
+
+void axolotl::Session::initialise_as_bob(
+    std::uint8_t const * shared_secret, std::size_t shared_secret_length,
+    axolotl::Curve25519PublicKey const & their_ratchet_key
+) {
+    std::uint8_t derived_secrets[64];
+    axolotl::hkdf_sha256(
+        shared_secret, shared_secret_length,
+        NULL, 0,
+        kdf_info.root_info, kdf_info.root_info_length,
+        derived_secrets, sizeof(derived_secrets)
+    );
+    receiver_chains.insert();
+    std::memcpy(root_key, derived_secrets, 32);
+    std::memcpy(receiver_chains[0].chain_key.key, derived_secrets + 32, 32);
+    receiver_chains[0].ratchet_key = their_ratchet_key;
+    unset(derived_secrets);
+}
+
+
+void axolotl::Session::initialise_as_alice(
+    std::uint8_t const * shared_secret, std::size_t shared_secret_length,
+    axolotl::Curve25519KeyPair const & our_ratchet_key
+) {
+    std::uint8_t derived_secrets[64];
+    axolotl::hkdf_sha256(
+        shared_secret, shared_secret_length,
+        NULL, 0,
+        kdf_info.root_info, kdf_info.root_info_length,
+        derived_secrets, sizeof(derived_secrets)
+    );
+    sender_chain.insert();
+    std::memcpy(root_key, derived_secrets, 32);
+    std::memcpy(sender_chain[0].chain_key.key, derived_secrets + 32, 32);
+    sender_chain[0].ratchet_key = our_ratchet_key;
+    unset(derived_secrets);
+}
+
+
 std::size_t axolotl::Session::encrypt_max_output_length(
     std::size_t plaintext_length
 ) {
-    std::size_t key_length = 1 + varstring_length(Curve25519PublicKey::Length);
-    std::size_t counter = sender_chain.empty() ? 0 : sender_chain[0].index;
+    std::size_t counter = 0;
+    if (!sender_chain.empty()) {
+        counter = sender_chain[0].chain_key.index;
+    }
     std::size_t padded = axolotl::aes_encrypt_cbc_length(plaintext_length);
     return axolotl::encode_message_length(
         counter, KEY_LENGTH, padded, MAC_LENGTH
@@ -176,7 +231,7 @@ std::size_t axolotl::Session::encrypt_max_output_length(
 
 
 std::size_t axolotl::Session::encrypt_random_length() {
-    return sender_chain.size() ? Curve25519PublicKey::Length : 0;
+    return sender_chain.empty() ? KEY_LENGTH : 0;
 }
 
 
@@ -189,29 +244,36 @@ std::size_t axolotl::Session::encrypt(
         last_error = axolotl::ErrorCode::NOT_ENOUGH_RANDOM;
         return std::size_t(-1);
     }
-    if (max_output_length < encrypt_max_output_length()) {
+    if (max_output_length < encrypt_max_output_length(plaintext_length)) {
         last_error = axolotl::ErrorCode::OUTPUT_BUFFER_TOO_SMALL;
         return std::size_t(-1);
     }
 
     if (sender_chain.empty()) {
-        /** create sender chain */
+        sender_chain.insert();
+        axolotl::generate_key(random, sender_chain[0].ratchet_key);
+        create_chain_key(
+            root_key,
+            sender_chain[0].ratchet_key,
+            receiver_chains[0].ratchet_key,
+            kdf_info,
+            root_key, sender_chain[0].chain_key
+        );
     }
 
     MessageKey keys;
-
-    /** create message keys and advance chain */
+    create_message_keys(sender_chain[0].chain_key, kdf_info, keys);
+    advance_chain_key(sender_chain[0].chain_key, sender_chain[0].chain_key);
 
     std::size_t padded = axolotl::aes_encrypt_cbc_length(plaintext_length);
-    std::size_t key_length = Curve25519PublicKey::Length;
     std::uint32_t counter = keys.index;
     const Curve25519PublicKey &ratchet_key = sender_chain[0].ratchet_key;
 
     axolotl::MessageWriter writer(axolotl::encode_message(
-        PROTOCOL_VERSION, counter, key_length, padded, cipher_text
+        PROTOCOL_VERSION, counter, KEY_LENGTH, padded, output
     ));
 
-    std::memcpy(writer.ratchet_key, ratchet_key.public_key, key_length);
+    std::memcpy(writer.ratchet_key, ratchet_key.public_key, KEY_LENGTH);
 
     axolotl::aes_encrypt_cbc(
         keys.cipher_key, keys.iv,
@@ -219,19 +281,20 @@ std::size_t axolotl::Session::encrypt(
         writer.ciphertext
     );
 
-    std::uint8_t mac[HMAC_SHA256_OUTPUT_LENGTH];
+    std::uint8_t mac[axolotl::HMAC_SHA256_OUTPUT_LENGTH];
     axolotl::hmac_sha256(
         keys.mac_key, sizeof(keys.mac_key),
-        ciphertext, writer.body_length,
+        output, writer.body_length,
         mac
     );
     std::memcpy(writer.mac, mac, MAC_LENGTH);
 
+    unset(keys);
     return writer.body_length + MAC_LENGTH;
 }
 
 
-std::size_t decrypt_max_plaintext_length(
+std::size_t axolotl::Session::decrypt_max_plaintext_length(
     std::size_t input_length
 ) {
     return input_length;
@@ -256,8 +319,7 @@ std::size_t axolotl::Session::decrypt(
         return std::size_t(-1);
     }
 
-    if (reader.body_length == 0
-            || reader.ratchet_key_length != Curve25519PublicKey::Length) {
+    if (reader.body_length == 0 || reader.ratchet_key_length != KEY_LENGTH) {
         last_error = axolotl::ErrorCode::BAD_MESSAGE_FORMAT;
         return std::size_t(-1);
     }
@@ -265,7 +327,8 @@ std::size_t axolotl::Session::decrypt(
     ReceiverChain * chain = NULL;
     for (axolotl::ReceiverChain & receiver_chain : receiver_chains) {
         if (0 == std::memcmp(
-                receiver_chain.ratchet_key, reader.ratchet_key, KEY_LENGTH
+                receiver_chain.ratchet_key.public_key, reader.ratchet_key,
+                KEY_LENGTH
         )) {
             chain = &receiver_chain;
             break;
@@ -278,15 +341,16 @@ std::size_t axolotl::Session::decrypt(
             return std::size_t(-1);
         }
     } else {
-        if (chain->index > reader.counter) {
+        if (chain->chain_key.index > reader.counter) {
             /* Chain already advanced beyond the key for this message
              * Check if the message keys are in the skipped key list. */
-            for (const axolotl::SkippedMessageKey & skipped
-                    : skipped_message_keys) {
+            for (axolotl::SkippedMessageKey & skipped : skipped_message_keys) {
                 if (reader.counter == skipped.message_key.index
                         && 0 == std::memcmp(
-                            skipped.ratchet_key, reader.ratchet_key, KEY_LENGTH
-                        )) {
+                            skipped.ratchet_key.public_key, reader.ratchet_key,
+                            KEY_LENGTH
+                        )
+                ) {
                     /* Found the key for this message. Check the MAC. */
                     if (!verify_mac(skipped.message_key, input, reader)) {
                         last_error = axolotl::ErrorCode::BAD_MESSAGE_MAC;
@@ -307,6 +371,7 @@ std::size_t axolotl::Session::decrypt(
 
                     /* Remove the key from the skipped keys now that we've
                      * decoded the message it corresponds to. */
+                    unset(skipped);
                     skipped_message_keys.erase(&skipped);
                     return result;
                 }
@@ -314,18 +379,54 @@ std::size_t axolotl::Session::decrypt(
             /* No matching keys for the message, fail with bad mac */
             last_error = axolotl::ErrorCode::BAD_MESSAGE_MAC;
             return std::size_t(-1);
-        } else if (!verify_mac_for_existing_chain(*chain, input, reader)) {
+        } else if (!verify_mac_for_existing_chain(
+               *this, chain->chain_key, input, reader
+        )) {
             last_error = axolotl::ErrorCode::BAD_MESSAGE_MAC;
             return std::size_t(-1);
         }
     }
 
     if (!chain) {
-        
+        /* They have started using a new empheral ratchet key.
+         * We need to derive a new set of chain keys.
+         * We can discard our previous empheral ratchet key.
+         * We will generate a new key when we send the next message. */
+        chain = receiver_chains.insert();
+        std::memcpy(
+            chain->ratchet_key.public_key, reader.ratchet_key, KEY_LENGTH
+        );
+        create_chain_key(
+            root_key, sender_chain[0].ratchet_key, chain->ratchet_key,
+            kdf_info, root_key, chain->chain_key
+        );
+        unset(sender_chain[0]);
+        sender_chain.erase(sender_chain.begin());
     }
 
+    while (chain->chain_key.index < reader.counter) {
+        axolotl::SkippedMessageKey & key = *skipped_message_keys.insert();
+        create_message_keys(chain->chain_key, kdf_info, key.message_key);
+        key.ratchet_key = chain->ratchet_key;
+        advance_chain_key(chain->chain_key, chain->chain_key);
+    }
 
+    axolotl::MessageKey message_key;
+    create_message_keys(chain->chain_key, kdf_info, message_key);
+    std::size_t result = axolotl::aes_decrypt_cbc(
+        message_key.cipher_key,
+        message_key.iv,
+        reader.ciphertext, reader.ciphertext_length,
+        plaintext
+    );
+    unset(message_key);
 
+    advance_chain_key(chain->chain_key, chain->chain_key);
 
-
+    if (result == std::size_t(-1)) {
+        last_error = axolotl::ErrorCode::BAD_MESSAGE_MAC;
+        return std::size_t(-1);
+    } else {
+        return result;
+    }
 }