From 08a7e44a966047a10d7e959d4a8cdeaaf4139ce0 Mon Sep 17 00:00:00 2001
From: Mark Haines <mark.haines@matrix.org>
Date: Fri, 12 Jun 2015 09:08:15 +0100
Subject: Pass the message body to decrypt_max_plaintext_length so we can get a
 more accurate estimate, rename encrypt_max_output_length to
 encrypt_output_length and change the api to return the exact number of bytes
 needed to hold the message

---
 src/ratchet.cpp | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

(limited to 'src/ratchet.cpp')

diff --git a/src/ratchet.cpp b/src/ratchet.cpp
index 5097643..91e5ce6 100644
--- a/src/ratchet.cpp
+++ b/src/ratchet.cpp
@@ -348,7 +348,7 @@ std::size_t axolotl::Ratchet::unpickle(
 }
 
 
-std::size_t axolotl::Ratchet::encrypt_max_output_length(
+std::size_t axolotl::Ratchet::encrypt_output_length(
     std::size_t plaintext_length
 ) {
     std::size_t counter = 0;
@@ -374,7 +374,7 @@ std::size_t axolotl::Ratchet::encrypt(
     std::uint8_t const * random, std::size_t random_length,
     std::uint8_t * output, std::size_t max_output_length
 ) {
-    std::size_t output_length = encrypt_max_output_length(plaintext_length);
+    std::size_t output_length = encrypt_output_length(plaintext_length);
 
     if (random_length < encrypt_random_length()) {
         last_error = axolotl::ErrorCode::NOT_ENOUGH_RANDOM;
@@ -428,9 +428,19 @@ std::size_t axolotl::Ratchet::encrypt(
 
 
 std::size_t axolotl::Ratchet::decrypt_max_plaintext_length(
-    std::size_t input_length
+    std::uint8_t const * input, std::size_t input_length
 ) {
-    return input_length;
+    axolotl::MessageReader reader;
+    axolotl::decode_message(
+        reader, input, input_length, ratchet_cipher.mac_length()
+    );
+
+    if (!reader.ciphertext) {
+        last_error = axolotl::ErrorCode::BAD_MESSAGE_FORMAT;
+        return std::size_t(-1);
+    }
+
+    return ratchet_cipher.decrypt_max_plaintext_length(reader.ciphertext_length);
 }
 
 
@@ -438,11 +448,6 @@ std::size_t axolotl::Ratchet::decrypt(
     std::uint8_t const * input, std::size_t input_length,
     std::uint8_t * plaintext, std::size_t max_plaintext_length
 ) {
-    if (max_plaintext_length < decrypt_max_plaintext_length(input_length)) {
-        last_error = axolotl::ErrorCode::OUTPUT_BUFFER_TOO_SMALL;
-        return std::size_t(-1);
-    }
-
     axolotl::MessageReader reader;
     axolotl::decode_message(
         reader, input, input_length, ratchet_cipher.mac_length()
@@ -458,6 +463,15 @@ std::size_t axolotl::Ratchet::decrypt(
         return std::size_t(-1);
     }
 
+    std::size_t max_length = ratchet_cipher.decrypt_max_plaintext_length(
+        reader.ciphertext_length
+    );
+
+    if (max_plaintext_length < max_length) {
+        last_error = axolotl::ErrorCode::OUTPUT_BUFFER_TOO_SMALL;
+        return std::size_t(-1);
+    }
+
     if (reader.ratchet_key_length != KEY_LENGTH) {
         last_error = axolotl::ErrorCode::BAD_MESSAGE_FORMAT;
         return std::size_t(-1);
-- 
cgit v1.2.3