Add a test for pickling and unpickling sessions, fix off by one error when unpickling sessions

3 files changed, 80 insertions, 12 deletions

diff --git a/olm.py b/olm.py
index cd2daad..80ac62e 100755
--- a/olm.py
+++ b/olm.py
@@ -130,12 +130,12 @@ class Account(object):
 
     def generate_one_time_keys(self, count):
         random_length = lib.olm_account_generate_one_time_keys_random_length(
-            self.ptr
+            self.ptr, count
         )
         random = read_random(random_length)
         random_buffer = create_string_buffer(random)
         lib.olm_account_generate_one_time_keys(
-            self.ptr, random_buffer, random_length
+            self.ptr, count, random_buffer, random_length
         )
 
     def clear(self):
@@ -323,17 +323,21 @@ if __name__ == '__main__':
 
     keys = commands.add_parser("keys", help="List public keys for an account")
     keys.add_argument("account_file", help="Local account file")
+    keys.add_argument("--json", action="store_true", help="Output as JSON")
 
     def do_keys(args):
         account = Account()
         with open(args.account_file, "rb") as f:
             account.unpickle(args.key, f.read())
-        result1 = {
+        result = {
             "account_keys": account.identity_keys(),
             "one_time_keys": account.one_time_keys(),
         }
         try:
-            yaml.safe_dump(result1, sys.stdout, default_flow_style=False)
+            if args.json:
+                json.dump(result, sys.stdout, indent=4)
+            else:
+                yaml.safe_dump(result, sys.stdout, default_flow_style=False)
         except:
             pass
 
@@ -356,6 +360,22 @@ if __name__ == '__main__':
 
     sign.set_defaults(func=do_sign)
 
+
+    generate_keys = commands.add_parser("generate_keys", help="Generate one time keys")
+    generate_keys.add_argument("account_file", help="Local account file")
+    generate_keys.add_argument("count", type=int, help="Number of keys to generate")
+
+    def do_generate_keys(args):
+        account = Account()
+        with open(args.account_file, "rb") as f:
+            account.unpickle(args.key, f.read())
+        account.generate_one_time_keys(args.count)
+        with open(args.account_file, "wb") as f:
+            f.write(account.pickle(args.key))
+
+    generate_keys.set_defaults(func=do_generate_keys)
+
+
     outbound = commands.add_parser("outbound", help="Create an outbound session")
     outbound.add_argument("account_file", help="Local account file")
     outbound.add_argument("session_file", help="Local session file")
@@ -365,7 +385,7 @@ if __name__ == '__main__':
     def do_outbound(args):
         if os.path.exists(args.session_file):
             sys.stderr.write("Session %r file already exists" % (
-                args.account_file,
+                args.session_file,
             ))
             sys.exit(1)
         account = Account()
diff --git a/src/olm.cpp b/src/olm.cpp
index 6f6d8ba..53cae54 100644
--- a/src/olm.cpp
+++ b/src/olm.cpp
@@ -152,7 +152,7 @@ std::size_t b64_input(
     return raw_length;
 }
 
-const char * errors[11] {
+static const char * ERRORS[11] {
     "SUCCESS",
     "NOT_ENOUGH_RANDOM",
     "OUTPUT_BUFFER_TOO_SMALL",
@@ -181,8 +181,8 @@ const char * olm_account_last_error(
     OlmSession * account
 ) {
     unsigned error = unsigned(from_c(account)->last_error);
-    if (error < 9) {
-        return errors[error];
+    if (error < sizeof(ERRORS)) {
+        return ERRORS[error];
     } else {
         return "UNKNOWN_ERROR";
     }
@@ -193,8 +193,8 @@ const char * olm_session_last_error(
     OlmSession * session
 ) {
     unsigned error = unsigned(from_c(session)->last_error);
-    if (error < 9) {
-        return errors[error];
+    if (error < sizeof(ERRORS)) {
+        return ERRORS[error];
     } else {
         return "UNKNOWN_ERROR";
     }
@@ -337,7 +337,7 @@ size_t olm_unpickle_session(
         return std::size_t(-1);
     }
 
-    std::uint8_t * const end = pos + raw_length + 1;
+    std::uint8_t * const end = pos + raw_length;
     /* On success unpickle will return (pos + raw_length). If unpickling
      * terminates too soon then it will return a pointer before
      * (pos + raw_length). On error unpickle will return (pos + raw_length + 1).
diff --git a/tests/test_olm.cpp b/tests/test_olm.cpp
index ccd023f..551b3cd 100644
--- a/tests/test_olm.cpp
+++ b/tests/test_olm.cpp
@@ -55,7 +55,9 @@ std::memcpy(pickle2, pickle1, pickle_length);
 
 std::uint8_t account_buffer2[::olm_account_size()];
 ::OlmAccount *account2 = ::olm_account(account_buffer2);
-::olm_unpickle_account(account2, "secret_key", 10, pickle2, pickle_length);
+assert_not_equals(std::size_t(-1), ::olm_unpickle_account(
+    account2, "secret_key", 10, pickle2, pickle_length
+));
 assert_equals(pickle_length, ::olm_pickle_account_length(account2));
 ::olm_pickle_account(account2, "secret_key", 10, pickle2, pickle_length);
 
@@ -63,6 +65,52 @@ assert_equals(pickle1, pickle2, pickle_length);
 
 }
 
+{ /** Pickle session test */
+
+TestCase test_case("Pickle session test");
+MockRandom mock_random('P');
+
+std::uint8_t account_buffer[::olm_account_size()];
+::OlmAccount *account = ::olm_account(account_buffer);
+std::uint8_t random[::olm_create_account_random_length(account)];
+mock_random(random, sizeof(random));
+::olm_create_account(account, random, sizeof(random));
+
+std::uint8_t session_buffer[::olm_session_size()];
+::OlmSession *session = ::olm_session(session_buffer);
+std::uint8_t identity_key[32];
+std::uint8_t one_time_key[32];
+mock_random(identity_key, sizeof(identity_key));
+mock_random(one_time_key, sizeof(one_time_key));
+std::uint8_t random2[::olm_create_outbound_session_random_length(session)];
+mock_random(random2, sizeof(random2));
+
+::olm_create_outbound_session(
+    session, account,
+    identity_key, sizeof(identity_key),
+    one_time_key, sizeof(one_time_key),
+    random2, sizeof(random2)
+);
+
+
+std::size_t pickle_length = ::olm_pickle_session_length(session);
+std::uint8_t pickle1[pickle_length];
+::olm_pickle_session(session, "secret_key", 10, pickle1, pickle_length);
+std::uint8_t pickle2[pickle_length];
+std::memcpy(pickle2, pickle1, pickle_length);
+
+std::uint8_t session_buffer2[::olm_session_size()];
+::OlmSession *session2 = ::olm_session(session_buffer2);
+assert_not_equals(std::size_t(-1), ::olm_unpickle_session(
+    session2, "secret_key", 10, pickle2, pickle_length
+));
+assert_equals(pickle_length, ::olm_pickle_session_length(session2));
+::olm_pickle_session(session2, "secret_key", 10, pickle2, pickle_length);
+
+assert_equals(pickle1, pickle2, pickle_length);
+
+}
+
 { /** Loopback test */
 
 TestCase test_case("Loopback test");