Implement the session key exchange

10 files changed, 368 insertions, 24 deletions

diff --git a/include/axolotl/account.hh b/include/axolotl/account.hh
index 5edb799..dd9c819 100644
--- a/include/axolotl/account.hh
+++ b/include/axolotl/account.hh
@@ -2,6 +2,8 @@
 #define AXOLOTL_ACCOUNT_HH_
 
 #include "axolotl/list.hh"
+#include "axolotl/crypto.hh"
+#include "axolotl/error.hh"
 
 #include <cstdint>
 
@@ -25,16 +27,21 @@ struct Account {
     LocalKey identity_key;
     LocalKey last_resort_one_time_key;
     List<LocalKey, MAX_ONE_TIME_KEYS> one_time_keys;
+    ErrorCode last_error;
 
     /** Number of random bytes needed to create a new account */
     std::size_t new_account_random_length();
 
     /** Create a new account. Returns NOT_ENOUGH_RANDOM if the number of random
      * bytes is too small. */
-    ErrorCode new_account(
+    std::size_t new_account(
         uint8_t const * random, std::size_t random_length
     );
 
+    LocalKey const * lookup_key(
+        std::uint32_t id
+    );
+
     /** The number of bytes needed to persist this account. */
     std::size_t pickle_length();
 
diff --git a/include/axolotl/crypto.hh b/include/axolotl/crypto.hh
index 42e4b61..7564e8f 100644
--- a/include/axolotl/crypto.hh
+++ b/include/axolotl/crypto.hh
@@ -12,6 +12,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#ifndef AXOLOTL_CRYPTO_HH_
+#define AXOLOTL_CRYPTO_HH_
+
 #include <cstdint>
 #include <cstddef>
 
@@ -141,3 +144,5 @@ void hkdf_sha256(
 );
 
 } // namespace axolotl
+
+#endif /* AXOLOTL_CRYPTO_HH_ */
diff --git a/include/axolotl/error.hh b/include/axolotl/error.hh
index 712b9eb..3bf0e63 100644
--- a/include/axolotl/error.hh
+++ b/include/axolotl/error.hh
@@ -10,6 +10,7 @@ enum struct ErrorCode {
     BAD_MESSAGE_VERSION = 3,  /*!< The message version is unsupported */
     BAD_MESSAGE_FORMAT = 4, /*!< The message couldn't be decoded */
     BAD_MESSAGE_MAC = 5, /*!< The message couldn't be decrypted */
+    BAD_MESSAGE_KEY_ID = 6, /*!< The message references an unknown key id */
 };
 
 } // namespace axolotl
diff --git a/include/axolotl/list.hh b/include/axolotl/list.hh
index ae8900c..604f00f 100644
--- a/include/axolotl/list.hh
+++ b/include/axolotl/list.hh
@@ -12,6 +12,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#ifndef AXOLOTL_LIST_HH_
+#define AXOLOTL_LIST_HH_
+
 #include <cstddef>
 
 namespace axolotl {
@@ -112,3 +115,5 @@ private:
 };
 
 } // namespace axolotl
+
+#endif /* AXOLOTL_LIST_HH_ */
diff --git a/include/axolotl/message.hh b/include/axolotl/message.hh
index 2b9bc99..5bce277 100644
--- a/include/axolotl/message.hh
+++ b/include/axolotl/message.hh
@@ -88,6 +88,7 @@ struct PreKeyMessageReader {
     std::uint8_t const * message; std::size_t message_length;
 };
 
+
 /**
  * The length of the buffer needed to hold a message.
  */
@@ -99,6 +100,7 @@ std::size_t encode_one_time_key_message_length(
     std::size_t message_length
 );
 
+
 /**
  * Writes the message headers into the output buffer.
  * Populates the writer struct with pointers into the output buffer.
diff --git a/include/axolotl/session.hh b/include/axolotl/session.hh
index c69699d..1c3395a 100644
--- a/include/axolotl/session.hh
+++ b/include/axolotl/session.hh
@@ -5,14 +5,13 @@
 
 namespace axolotl {
 
+class Account;
+
 struct RemoteKey {
     std::uint32_t id;
     Curve25519PublicKey key;
 };
 
-struct RemoteKeys {
-};
-
 
 enum struct MessageType {
     PRE_KEY_MESSAGE = 0,
@@ -21,28 +20,34 @@ enum struct MessageType {
 
 
 struct Session {
+
+    Session();
+
+    Ratchet ratchet;
+    ErrorCode last_error;
+
     bool received_message;
+
     RemoteKey alice_identity_key;
-    RemoteKey alice_base_key;
-    RemoteKey bob_identity_key;
-    RemoteKey bob_one_time_key;
-    Ratchet ratchet;
+    Curve25519PublicKey alice_base_key;
+    std::uint32_t bob_one_time_key_id;
+
 
-    void initialise_outbound_session_random_length();
+    std::size_t new_outbound_session_random_length();
 
-    void initialise_outbound_session(
+    std::size_t new_outbound_session(
         Account const & local_account,
-        RemoteKey const & identity_key,
+        Curve25519PublicKey const & identity_key,
         RemoteKey const & one_time_key,
         std::uint8_t const * random, std::size_t random_length
     );
 
-    void initialise_inbound_session(
+    std::size_t new_inbound_session(
         Account & local_account,
         std::uint8_t const * one_time_key_message, std::size_t message_length
     );
 
-    void matches_inbound_session(
+    bool matches_inbound_session(
         std::uint8_t const * one_time_key_message, std::size_t message_length
     );
 
diff --git a/src/cipher.cpp b/src/cipher.cpp
index 86cde88..8b496df 100644
--- a/src/cipher.cpp
+++ b/src/cipher.cpp
@@ -26,7 +26,7 @@ static void derive_keys(
     std::uint8_t derived_secrets[80];
     axolotl::hkdf_sha256(
         key, key_length,
-        NULL, 0,
+        nullptr, 0,
         kdf_info, kdf_info_length,
         derived_secrets, sizeof(derived_secrets)
     );
diff --git a/src/message.cpp b/src/message.cpp
index fcedd07..d9978cb 100644
--- a/src/message.cpp
+++ b/src/message.cpp
@@ -198,15 +198,15 @@ void axolotl::decode_message(
 ) {
     std::uint8_t const * pos = input;
     std::uint8_t const * end = input + input_length - mac_length;
-    std::uint8_t const * unknown = NULL;
+    std::uint8_t const * unknown = nullptr;
 
     if (pos == end) return;
     reader.version = *(pos++);
     reader.input = input;
     reader.input_length = input_length;
     reader.has_counter = false;
-    reader.ratchet_key = NULL;
-    reader.ciphertext = NULL;
+    reader.ratchet_key = nullptr;
+    reader.ciphertext = nullptr;
 
     while (pos != end) {
         pos = decode(
@@ -283,15 +283,15 @@ void axolotl::decode_one_time_key_message(
 ) {
     std::uint8_t const * pos = input;
     std::uint8_t const * end = input + input_length;
-    std::uint8_t const * unknown = NULL;
+    std::uint8_t const * unknown = nullptr;
 
     if (pos == end) return;
     reader.version = *(pos++);
     reader.has_registration_id = false;
     reader.has_one_time_key_id = false;
-    reader.identity_key = NULL;
-    reader.base_key = NULL;
-    reader.message = NULL;
+    reader.identity_key = nullptr;
+    reader.base_key = nullptr;
+    reader.message = nullptr;
 
     while (pos != end) {
         pos = decode(
diff --git a/src/ratchet.cpp b/src/ratchet.cpp
index 91e5ce6..87d79b7 100644
--- a/src/ratchet.cpp
+++ b/src/ratchet.cpp
@@ -184,7 +184,7 @@ void axolotl::Ratchet::initialise_as_bob(
     std::uint8_t derived_secrets[64];
     axolotl::hkdf_sha256(
         shared_secret, shared_secret_length,
-        NULL, 0,
+        nullptr, 0,
         kdf_info.root_info, kdf_info.root_info_length,
         derived_secrets, sizeof(derived_secrets)
     );
@@ -203,7 +203,7 @@ void axolotl::Ratchet::initialise_as_alice(
     std::uint8_t derived_secrets[64];
     axolotl::hkdf_sha256(
         shared_secret, shared_secret_length,
-        NULL, 0,
+        nullptr, 0,
         kdf_info.root_info, kdf_info.root_info_length,
         derived_secrets, sizeof(derived_secrets)
     );
@@ -477,7 +477,7 @@ std::size_t axolotl::Ratchet::decrypt(
         return std::size_t(-1);
     }
 
-    ReceiverChain * chain = NULL;
+    ReceiverChain * chain = nullptr;
     for (axolotl::ReceiverChain & receiver_chain : receiver_chains) {
         if (0 == std::memcmp(
                 receiver_chain.ratchet_key.public_key, reader.ratchet_key,
diff --git a/src/session.cpp b/src/session.cpp
new file mode 100644
index 0000000..be2bce7
--- /dev/null
+++ b/src/session.cpp
@@ -0,0 +1,319 @@
+#include "axolotl/session.hh"
+#include "axolotl/cipher.hh"
+#include "axolotl/crypto.hh"
+#include "axolotl/account.hh"
+#include "axolotl/memory.hh"
+#include "axolotl/message.hh"
+
+#include <cstring>
+
+namespace {
+
+static const std::size_t KEY_LENGTH = 32;
+static const std::uint8_t PROTOCOL_VERSION = 0x3;
+
+static const std::uint8_t ROOT_KDF_INFO[] = "AXOLOTL_ROOT";
+static const std::uint8_t RATCHET_KDF_INFO[] = "AXOLOTL_RATCHET";
+static const std::uint8_t CIPHER_KDF_INFO[] = "AXOLOTL_KEYS";
+
+static const axolotl::CipherAesSha256 AXOLOTL_CIPHER(
+    CIPHER_KDF_INFO, sizeof(CIPHER_KDF_INFO) -1
+);
+
+static const axolotl::KdfInfo AXOLOTL_KDF_INFO = {
+    ROOT_KDF_INFO, sizeof(ROOT_KDF_INFO) - 1,
+    RATCHET_KDF_INFO, sizeof(RATCHET_KDF_INFO) - 1
+};
+
+} // namespace
+
+axolotl::Session::Session(
+) : ratchet(AXOLOTL_KDF_INFO, AXOLOTL_CIPHER),
+    last_error(axolotl::ErrorCode::SUCCESS),
+    received_message(false),
+    bob_one_time_key_id(0) {
+
+}
+
+
+std::size_t axolotl::Session::new_outbound_session_random_length() {
+    return KEY_LENGTH;
+}
+
+
+std::size_t axolotl::Session::new_outbound_session(
+    axolotl::Account const & local_account,
+    axolotl::Curve25519PublicKey const & identity_key,
+    axolotl::RemoteKey const & one_time_key,
+    std::uint8_t const * random, std::size_t random_length
+) {
+    if (random_length < new_outbound_session_random_length()) {
+        last_error = axolotl::ErrorCode::NOT_ENOUGH_RANDOM;
+        return std::size_t(-1);
+    }
+
+    Curve25519KeyPair base_key;
+    axolotl::generate_key(random, base_key);
+
+    received_message = false;
+    alice_identity_key.id = local_account.identity_key.id;
+    alice_identity_key.key = local_account.identity_key.key;
+    alice_base_key = base_key;
+    bob_one_time_key_id = one_time_key.id;
+
+    std::uint8_t shared_secret[160];
+    std::memset(shared_secret, 0xFF, 32);
+
+    axolotl::curve25519_shared_secret(
+    );
+    axolotl::curve25519_shared_secret(
+         base_key, identity_key, shared_secret + 64
+    );
+    axolotl::curve25519_shared_secret(
+    );
+    axolotl::curve25519_shared_secret(
+         base_key, one_time_key.key, shared_secret + 128
+    );
+
+    axolotl::unset(base_key);
+    axolotl::unset(shared_secret);
+
+    return std::size_t(0);
+}
+
+namespace {
+
+bool check_message_fields(
+    axolotl::PreKeyMessageReader & reader
+) {
+    bool ok = true;
+    ok = ok && reader.identity_key;
+    ok = ok && reader.identity_key_length == KEY_LENGTH;
+    ok = ok && reader.message;
+    ok = ok && reader.base_key;
+    ok = ok && reader.base_key_length == KEY_LENGTH;
+    ok = ok && reader.has_one_time_key_id;
+    ok = ok && reader.has_registration_id;
+    return ok;
+}
+
+} // namespace
+
+
+std::size_t axolotl::Session::new_inbound_session(
+    axolotl::Account & local_account,
+    std::uint8_t const * one_time_key_message, std::size_t message_length
+) {
+    axolotl::PreKeyMessageReader reader;
+    decode_one_time_key_message(reader, one_time_key_message, message_length);
+
+    if (!check_message_fields(reader)) {
+        last_error = axolotl::ErrorCode::BAD_MESSAGE_FORMAT;
+        return std::size_t(-1);
+    }
+
+    alice_identity_key.id = reader.registration_id;
+    std::memcpy(alice_identity_key.key.public_key, reader.identity_key, 32);
+    std::memcpy(alice_base_key.public_key, reader.base_key, 32);
+    bob_one_time_key_id = reader.one_time_key_id;
+
+
+    axolotl::LocalKey const * bob_one_time_key = local_account.lookup_key(
+        bob_one_time_key_id
+    );
+
+        last_error = axolotl::ErrorCode::BAD_MESSAGE_KEY_ID;
+        return std::size_t(-1);
+    }
+
+    std::uint8_t shared_secret[160];
+    std::memset(shared_secret, 0xFF, 32);
+
+    axolotl::curve25519_shared_secret(
+    );
+    axolotl::curve25519_shared_secret(
+        local_account.identity_key.key, alice_base_key, shared_secret + 64
+    );
+    axolotl::curve25519_shared_secret(
+    );
+    axolotl::curve25519_shared_secret(
+        bob_one_time_key->key, alice_base_key, shared_secret + 128
+    );
+
+
+    return std::size_t(0);
+}
+
+
+bool axolotl::Session::matches_inbound_session(
+    std::uint8_t const * one_time_key_message, std::size_t message_length
+) {
+    axolotl::PreKeyMessageReader reader;
+    decode_one_time_key_message(reader, one_time_key_message, message_length);
+
+    if (!check_message_fields(reader)) {
+        return false;
+    }
+
+    bool same = true;
+    same = same && 0 == std::memcmp(
+        reader.identity_key, alice_identity_key.key.public_key, KEY_LENGTH
+    );
+    same = same && 0 == std::memcmp(
+        reader.base_key, alice_base_key.public_key, KEY_LENGTH
+    );
+    same = same && reader.one_time_key_id == bob_one_time_key_id;
+    same = same && reader.registration_id == alice_identity_key.id;
+    return same;
+}
+
+
+axolotl::MessageType axolotl::Session::encrypt_message_type() {
+    if (received_message) {
+        return axolotl::MessageType::MESSAGE;
+    } else {
+        return axolotl::MessageType::PRE_KEY_MESSAGE;
+    }
+}
+
+
+std::size_t axolotl::Session::encrypt_message_length(
+    std::size_t plaintext_length
+) {
+    std::size_t message_length = ratchet.encrypt_output_length(
+        plaintext_length
+    );
+
+    if (received_message) {
+        return message_length;
+    }
+
+    return encode_one_time_key_message_length(
+        alice_identity_key.id,
+        bob_one_time_key_id,
+        KEY_LENGTH,
+        KEY_LENGTH,
+        message_length
+    );
+}
+
+
+std::size_t axolotl::Session::encrypt_random_length() {
+    return ratchet.encrypt_random_length();
+}
+
+
+std::size_t axolotl::Session::encrypt(
+    std::uint8_t const * plaintext, std::size_t plaintext_length,
+    std::uint8_t const * random, std::size_t random_length,
+    std::uint8_t * message, std::size_t message_length
+) {
+    if (message_length < encrypt_message_length(plaintext_length)) {
+        last_error = axolotl::ErrorCode::OUTPUT_BUFFER_TOO_SMALL;
+        return std::size_t(-1);
+    }
+    std::uint8_t * message_body;
+    std::size_t message_body_length = ratchet.encrypt_output_length(
+        plaintext_length
+    );
+
+    if (received_message) {
+        message_body = message;
+    } else {
+        axolotl::PreKeyMessageWriter writer;
+        encode_one_time_key_message(
+            writer,
+            PROTOCOL_VERSION,
+            alice_identity_key.id,
+            bob_one_time_key_id,
+            KEY_LENGTH,
+            KEY_LENGTH,
+            message_body_length,
+            message
+        );
+        std::memcpy(
+            writer.identity_key, alice_identity_key.key.public_key, KEY_LENGTH
+        );
+        std::memcpy(
+            writer.base_key, alice_base_key.public_key, KEY_LENGTH
+        );
+        message_body = writer.message;
+    }
+
+    std::size_t result = ratchet.encrypt(
+        plaintext, plaintext_length,
+        random, random_length,
+        message_body, message_body_length
+    );
+
+    if (result == std::size_t(-1)) {
+        last_error = ratchet.last_error;
+        ratchet.last_error = axolotl::ErrorCode::SUCCESS;
+    }
+    return result;
+}
+
+
+std::size_t axolotl::Session::decrypt_max_plaintext_length(
+    MessageType message_type,
+    std::uint8_t const * message, std::size_t message_length
+) {
+    std::uint8_t const * message_body;
+    std::size_t message_body_length;
+    if (message_type == axolotl::MessageType::MESSAGE) {
+        message_body = message;
+        message_body_length = message_length;
+    } else {
+        axolotl::PreKeyMessageReader reader;
+        decode_one_time_key_message(reader, message, message_length);
+        if (!reader.message) {
+            last_error = axolotl::ErrorCode::BAD_MESSAGE_FORMAT;
+            return std::size_t(-1);
+        }
+        message_body = reader.message;
+        message_body_length = reader.message_length;
+    }
+
+    std::size_t result = ratchet.decrypt_max_plaintext_length(
+        message_body, message_body_length
+    );
+
+    if (result == std::size_t(-1)) {
+        last_error = ratchet.last_error;
+        ratchet.last_error = axolotl::ErrorCode::SUCCESS;
+    }
+    return result;
+}
+
+
+std::size_t axolotl::Session::decrypt(
+    axolotl::MessageType message_type,
+    std::uint8_t const * message, std::size_t message_length,
+    std::uint8_t * plaintext, std::size_t max_plaintext_length
+) {
+    std::uint8_t const * message_body;
+    std::size_t message_body_length;
+    if (message_type == axolotl::MessageType::MESSAGE) {
+        message_body = message;
+        message_body_length = message_length;
+    } else {
+        axolotl::PreKeyMessageReader reader;
+        decode_one_time_key_message(reader, message, message_length);
+        if (!reader.message) {
+            last_error = axolotl::ErrorCode::BAD_MESSAGE_FORMAT;
+            return std::size_t(-1);
+        }
+        message_body = reader.message;
+        message_body_length = reader.message_length;
+    }
+
+    std::size_t result = ratchet.decrypt(
+        message_body, message_body_length, plaintext, max_plaintext_length
+    );
+
+    if (result == std::size_t(-1)) {
+        last_error = ratchet.last_error;
+        ratchet.last_error = axolotl::ErrorCode::SUCCESS;
+    }
+    return result;
+}