4 files changed, 59 insertions, 5 deletions

diff --git a/autotests/testolmsession.cpp b/autotests/testolmsession.cpp
index 1b7fbb9b..6fa2a380 100644
--- a/autotests/testolmsession.cpp
+++ b/autotests/testolmsession.cpp
@@ -1,7 +1,13 @@
+// SPDX-FileCopyrightText: 2021 Carl Schwan <carlschwan@kde.org>
+//
+// SPDX-License-Identifier: LGPL-2.1-or-later
+
 #include "olm/session.h"
+#include "testolmsession.h"
 
 using namespace Quotient;
 
+#ifdef Quotient_E2EE_ENABLED
 std::pair<std::unique_ptr<QOlmSession>, std::unique_ptr<QOlmSession>> createSessionPair()
 {
     QByteArray pickledAccountA("eOBXIKivUT6YYowRH031BNv7zNmzqM5B7CpXdyeaPvala5mt7/OeqrG1qVA7vA1SYloFyvJPIy0QNkD3j1HiPl5vtZHN53rtfZ9exXDok03zjmssqn4IJsqcA7Fbo1FZeKafG0NFcWwCPTdmcV7REqxjqGm3I4K8MQFa45AdTGSUu2C12cWeOcbSMlcINiMral+Uyah1sgPmLJ18h1qcnskXUXQvpffZ5DiUw1Iz5zxnwOQF1GVyowPJD7Zdugvj75RQnDxAn6CzyvrY2k2CuedwqDC3fIXM2xdUNWttW4nC2g4InpBhCVvNwhZYxlUb5BUEjmPI2AB3dAL5ry6o9MFncmbN6x5x");
@@ -18,7 +24,7 @@ std::pair<std::unique_ptr<QOlmSession>, std::unique_ptr<QOlmSession>> createSess
     auto outbound = std::get<std::unique_ptr<QOlmSession>>(accountA
         .createOutboundSession(identityKeyB, oneTimeKeyB));
 
-    const auto preKey = std::get<Message>(outbound->encrypt("")); // Payload does not matter for PreKey
+    const auto preKey = outbound->encrypt(""); // Payload does not matter for PreKey
 
     if (preKey.type() != Message::General) {
         throw "Wrong first message type received, can't create session";
@@ -26,3 +32,14 @@ std::pair<std::unique_ptr<QOlmSession>, std::unique_ptr<QOlmSession>> createSess
     auto inbound = std::get<std::unique_ptr<QOlmSession>>(accountB.createInboundSession(preKey));
     return std::make_pair<std::unique_ptr<QOlmSession>, std::unique_ptr<QOlmSession>>(std::move(inbound), std::move(outbound));
 }
+#endif
+
+void TestOlmSession::olmOutboundSessionCreation()
+{
+#ifdef Quotient_E2EE_ENABLED
+    const auto [_, outboundSession] = createSessionPair();
+    QCOMPARE(0, outboundSession->hasReceivedMessage());
+#endif
+}
+
+QTEST_MAIN(TestOlmSession)
diff --git a/autotests/testolmsession.h b/autotests/testolmsession.h
new file mode 100644
index 00000000..7e3fc6e4
--- /dev/null
+++ b/autotests/testolmsession.h
@@ -0,0 +1,14 @@
+// SPDX-FileCopyrightText: 2021 Carl Schwan <carlschwan@kde.org>
+//
+// SPDX-License-Identifier: LGPL-2.1-or-later
+
+#ifdef Quotient_E2EE_ENABLED
+#include <QtTest/QtTest>
+
+class TestOlmSession : public QObject
+{
+    Q_OBJECT
+private Q_SLOTS:
+    void olmOutboundSessionCreation();
+};
+#endif
diff --git a/lib/olm/session.cpp b/lib/olm/session.cpp
index f6cab650..0beb136e 100644
--- a/lib/olm/session.cpp
+++ b/lib/olm/session.cpp
@@ -121,11 +121,12 @@ std::variant<std::unique_ptr<QOlmSession>, OlmError> QOlmSession::unpickle(QByte
     return std::make_unique<QOlmSession>(olmSession);
 }
 
-std::variant<Message, OlmError> QOlmSession::encrypt(const QString &plaintext)
+Message QOlmSession::encrypt(const QString &plaintext)
 {
     QByteArray plaintextBuf = plaintext.toUtf8();
     const auto messageMaxLen = olm_encrypt_message_length(m_session, plaintextBuf.length());
     QByteArray messageBuf(messageMaxLen, '0');
+    const auto messageType = encryptMessageType();
     const auto randomLen = olm_encrypt_random_length(m_session);
     QByteArray randomBuf = getRandom(randomLen);
     const auto error = olm_encrypt(m_session,
@@ -134,10 +135,22 @@ std::variant<Message, OlmError> QOlmSession::encrypt(const QString &plaintext)
                                    reinterpret_cast<uint8_t *>(messageBuf.data()), messageBuf.length());
 
     if (error == olm_error()) {
-        return lastError(m_session);
+        throw lastError(m_session);
     }
 
-    return Message::fromCiphertext(messageBuf);
+    return Message(messageBuf, messageType);
+}
+
+Message::Type QOlmSession::encryptMessageType()
+{
+    const auto messageTypeResult = olm_encrypt_message_type(m_session);
+    if (messageTypeResult == olm_error()) {
+        throw lastError(m_session);
+    }
+    if (messageTypeResult == OLM_MESSAGE_TYPE_PRE_KEY) {
+        return Message::PreKey;
+    }
+    return Message::General;
 }
 
 QByteArray QOlmSession::sessionId() const
@@ -152,6 +165,11 @@ QByteArray QOlmSession::sessionId() const
     return idBuffer;
 }
 
+bool QOlmSession::hasReceivedMessage() const
+{
+    return olm_session_has_received_message(m_session);
+}
+
 QOlmSession::QOlmSession(OlmSession *session)
     : m_session(session)
 {
diff --git a/lib/olm/session.h b/lib/olm/session.h
index 89f5d822..f9221dec 100644
--- a/lib/olm/session.h
+++ b/lib/olm/session.h
@@ -31,12 +31,17 @@ public:
     //! Deserialises from encrypted Base64 that was previously obtained by pickling a `QOlmSession`.
     static std::variant<std::unique_ptr<QOlmSession>, OlmError> unpickle(QByteArray &pickled, const PicklingMode &mode);
     //! Encrypts a plaintext message using the session.
-    std::variant<Message, OlmError> encrypt(const QString &plaintext);
+    Message encrypt(const QString &plaintext);
     // TODO: WiP
 
     //! Get a base64-encoded identifier for this session.
     QByteArray sessionId() const;
 
+    //! The type of the next message that will be returned from encryption.
+    Message::Type encryptMessageType();
+
+    bool hasReceivedMessage() const;
+
     QOlmSession(OlmSession* session);
 private:
     //! Helper function for creating new sessions and handling errors.