From 2ef1f6f4fc5bdc069483a527ab3a1b060c71fcad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sa=C3=BAl=20Ibarra=20Corretg=C3=A9?= <s@saghul.net>
Date: Wed, 23 Sep 2020 10:13:46 +0200
Subject: [PATCH] SAS: add olm_sas_is_their_key_set

Also make olm_sas_generate_bytes fail if their key wasn't set.
---
 include/olm/error.h         |  8 ++++----
 include/olm/sas.h           | 12 ++++++++++++
 javascript/olm_sas.js       |  6 ++++++
 javascript/test/sas.spec.js |  9 +++++++++
 src/error.c                 |  1 +
 src/sas.c                   | 12 ++++++++++++
 6 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/include/olm/error.h b/include/olm/error.h
index ee2187c..17e819d 100644
--- a/include/olm/error.h
+++ b/include/olm/error.h
@@ -53,10 +53,10 @@ enum OlmErrorCode {
 
     OLM_INPUT_BUFFER_TOO_SMALL = 15,
 
-    // Not an error code, just here to pad out the enum past 16 because
-    // otherwise the compiler warns about a redunant check. If you're
-    // adding an error code, replace this one!
-    OLM_ERROR_NOT_INVENTED_YET = 16,
+    /**
+     * SAS doesn't have their key set.
+     */
+    OLM_SAS_THEIR_KEY_NOT_SET = 16,
 
     /* remember to update the list of string constants in error.c when updating
      * this list. */
diff --git a/include/olm/sas.h b/include/olm/sas.h
index ee44eff..4a5e751 100644
--- a/include/olm/sas.h
+++ b/include/olm/sas.h
@@ -105,6 +105,15 @@ size_t olm_sas_set_their_key(
     void * their_key, size_t their_key_length
 );
 
+/** Checks if their key was set.
+ *
+ * @param[in] sas the SAS object.
+ *
+ */
+int olm_sas_is_their_key_set(
+    OlmSAS *sas
+);
+
 /** Generate bytes to use for the short authentication string.
  *
  * @param[in] sas the SAS object.
@@ -114,6 +123,9 @@ size_t olm_sas_set_their_key(
  * @param[out] output the output buffer.
  * @param[in] output_length the size of the output buffer.  For hex-based SAS
  *     as in the Matrix spec, this will be 5.
+ *
+ * @return `olm_error()` on failure. If their key wasn't set then
+ * `olm_sas_last_error()` will be `SAS_THEIR_KEY_NOT_SET`.
  */
 size_t olm_sas_generate_bytes(
     OlmSAS * sas,
diff --git a/javascript/olm_sas.js b/javascript/olm_sas.js
index 5cf3fb0..38535d5 100644
--- a/javascript/olm_sas.js
+++ b/javascript/olm_sas.js
@@ -42,6 +42,12 @@ SAS.prototype['set_their_key'] = restore_stack(function(their_key) {
     );
 });
 
+SAS.prototype['is_their_key_set'] = restore_stack(function() {
+    return sas_method(Module['_olm_sas_is_their_key_set'])(
+        this.ptr
+    ) ? true : false;
+});
+
 SAS.prototype['generate_bytes'] = restore_stack(function(info, length) {
     var info_array = array_from_string(info);
     var info_buffer = stack(info_array);
diff --git a/javascript/test/sas.spec.js b/javascript/test/sas.spec.js
index af7ea65..4ab4120 100644
--- a/javascript/test/sas.spec.js
+++ b/javascript/test/sas.spec.js
@@ -50,4 +50,13 @@ describe("sas", function() {
         bob.set_their_key(alice.get_pubkey());
         expect(alice.calculate_mac("test", "MAC").toString()).toEqual(bob.calculate_mac("test", "MAC").toString());
     });
+
+    it('should fail to generate bytes if their key is not set', function () {
+        expect(alice.is_their_key_set()).toBeFalsy();
+        expect(() => {
+            alice.generate_bytes("SAS", 5);
+        }).toThrow();
+        alice.set_their_key(bob.get_pubkey());
+        expect(alice.is_their_key_set()).toBeTruthy();
+    });
 });
diff --git a/src/error.c b/src/error.c
index 5147b5c..7d39637 100644
--- a/src/error.c
+++ b/src/error.c
@@ -32,6 +32,7 @@ static const char * ERRORS[] = {
     "BAD_LEGACY_ACCOUNT_PICKLE",
     "BAD_SIGNATURE",
     "OLM_INPUT_BUFFER_TOO_SMALL",
+    "OLM_SAS_THEIR_KEY_NOT_SET"
 };
 
 const char * _olm_error_to_string(enum OlmErrorCode error)
diff --git a/src/sas.c b/src/sas.c
index 76e4576..8d239b1 100644
--- a/src/sas.c
+++ b/src/sas.c
@@ -23,6 +23,7 @@ struct OlmSAS {
     enum OlmErrorCode last_error;
     struct _olm_curve25519_key_pair curve25519_key;
     uint8_t secret[CURVE25519_SHARED_SECRET_LENGTH];
+    int their_key_set;
 };
 
 const char * olm_sas_last_error(
@@ -95,14 +96,25 @@ size_t olm_sas_set_their_key(
     }
     _olm_decode_base64(their_key, their_key_length, their_key);
     _olm_crypto_curve25519_shared_secret(&sas->curve25519_key, their_key, sas->secret);
+    sas->their_key_set = 1;
     return 0;
 }
 
+int olm_sas_is_their_key_set(
+    OlmSAS *sas
+) {
+    return sas->their_key_set;
+}
+
 size_t olm_sas_generate_bytes(
     OlmSAS * sas,
     const void * info, size_t info_length,
     void * output, size_t output_length
 ) {
+    if (!sas->their_key_set) {
+        sas->last_error = OLM_SAS_THEIR_KEY_NOT_SET;
+        return (size_t)-1;
+    }
     _olm_crypto_hkdf_sha256(
         sas->secret, sizeof(sas->secret),
         NULL, 0,