From 76d211bbdf3a57afd72b3f3ac736c42cc17e3cac Mon Sep 17 00:00:00 2001
From: Maxim Mamontov <faust.madf@gmail.com>
Date: Sun, 14 Sep 2014 00:55:55 +0300
Subject: [PATCH] Handle callback errors in bfstreams.

---
 stglibs/crypto.lib/Makefile               |  6 ++++--
 stglibs/crypto.lib/bfstream.cpp           | 26 +++++++++++++++++++----
 stglibs/crypto.lib/include/stg/bfstream.h |  8 +++++--
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/stglibs/crypto.lib/Makefile b/stglibs/crypto.lib/Makefile
index 6e83e01e..47e91203 100644
--- a/stglibs/crypto.lib/Makefile
+++ b/stglibs/crypto.lib/Makefile
@@ -5,9 +5,11 @@
 LIB_NAME = stgcrypto
 
 SRCS = ag_md5.c \
-	blowfish.c
+	blowfish.c \
+	bfstream.cpp
 
 INCS = ag_md5.h \
-	blowfish.h
+	blowfish.h \
+	bfstream.h
 
 include ../Makefile.in
diff --git a/stglibs/crypto.lib/bfstream.cpp b/stglibs/crypto.lib/bfstream.cpp
index 24aa7e10..21a0d55d 100644
--- a/stglibs/crypto.lib/bfstream.cpp
+++ b/stglibs/crypto.lib/bfstream.cpp
@@ -16,14 +16,15 @@ const size_t BUFFER_SIZE = BFSTREAM_BUF_SIZE;
 class COMMON
 {
     public:
-        typedef void (* CALLBACK)(const void * block, size_t size, void * data);
+        typedef bool (* CALLBACK)(const void * block, size_t size, void * data);
         typedef void (* PROC)(void * dest, const void * source, size_t length, const BLOWFISH_CTX * ctx);
 
         COMMON(const std::string & key, CALLBACK callback, void * data, PROC proc)
             : m_ptr(m_buffer),
               m_callback(callback),
               m_data(data),
-              m_proc(proc)
+              m_proc(proc),
+              m_ok(true)
         {
         InitContext(key.c_str(), key.length(), &m_ctx);
         memset(m_buffer, 0, sizeof(m_buffer));
@@ -38,14 +39,18 @@ class COMMON
             size -= sizeof(m_buffer) - dataSize; // Adjust size
             data = static_cast<const char *>(data) + sizeof(m_buffer) - dataSize; // Adjust data pointer
             m_proc(m_buffer, m_buffer, sizeof(m_buffer), &m_ctx); // Process
-            m_callback(m_buffer, sizeof(m_buffer), m_data); // Consume
+            m_ok = m_ok && m_callback(m_buffer, sizeof(m_buffer), m_data); // Consume
             m_ptr = m_buffer;
             }
+        if (!m_ok)
+            return;
         memcpy(m_ptr, data, size);
         m_ptr += size;
         m_tryConsume(last);
         }
 
+        bool isOk() const { return m_ok; }
+
     private:
         char m_buffer[BUFFER_SIZE];
         char * m_ptr;
@@ -53,6 +58,7 @@ class COMMON
         void * m_data;
         BLOWFISH_CTX m_ctx;
         PROC m_proc;
+        bool m_ok;
 
         void m_tryConsume(bool last)
         {
@@ -66,7 +72,9 @@ class COMMON
         if (dataSize == 0)
             return;
         m_proc(m_buffer, m_buffer, dataSize, &m_ctx);
-        m_callback(m_buffer, dataSize, m_data);
+        m_ok = m_ok && m_callback(m_buffer, dataSize, m_data);
+        if (!m_ok)
+            return;
         if (remainder > 0)
             memmove(m_buffer, m_buffer + dataSize, remainder);
         m_ptr = m_buffer + remainder;
@@ -108,6 +116,11 @@ void ENCRYPT_STREAM::Put(const void * data, size_t size, bool last)
 m_impl->Put(data, size, last);
 }
 
+bool ENCRYPT_STREAM::isOk() const
+{
+return m_impl->isOk();
+}
+
 DECRYPT_STREAM::DECRYPT_STREAM(const std::string & key, CALLBACK callback, void * data)
     : m_impl(new IMPL(key, callback, data))
 {}
@@ -121,3 +134,8 @@ void DECRYPT_STREAM::Put(const void * data, size_t size, bool last)
 {
 m_impl->Put(data, size, last);
 }
+
+bool DECRYPT_STREAM::isOk() const
+{
+return m_impl->isOk();
+}
diff --git a/stglibs/crypto.lib/include/stg/bfstream.h b/stglibs/crypto.lib/include/stg/bfstream.h
index 41b44330..bead8d00 100644
--- a/stglibs/crypto.lib/include/stg/bfstream.h
+++ b/stglibs/crypto.lib/include/stg/bfstream.h
@@ -10,12 +10,14 @@ namespace STG
 class ENCRYPT_STREAM
 {
     public:
-        typedef void (* CALLBACK)(const void * block, size_t size, void * data);
+        typedef bool (* CALLBACK)(const void * block, size_t size, void * data);
 
         ENCRYPT_STREAM(const std::string & key, CALLBACK callback, void * data);
         ~ENCRYPT_STREAM();
         void Put(const void * data, size_t size, bool last = false);
 
+        bool isOk() const;
+
     private:
         class IMPL;
 
@@ -25,12 +27,14 @@ class ENCRYPT_STREAM
 class DECRYPT_STREAM
 {
     public:
-        typedef void (* CALLBACK)(const void * block, size_t size, void * data);
+        typedef bool (* CALLBACK)(const void * block, size_t size, void * data);
 
         DECRYPT_STREAM(const std::string & key, CALLBACK callback, void * data);
         ~DECRYPT_STREAM();
         void Put(const void * data, size_t size, bool last = false);
 
+        bool isOk() const;
+
     private:
         class IMPL;
 
-- 
2.44.2