X-Git-Url: https://git.stg.codes/stg.git/blobdiff_plain/27fd26be047af5ccd9d023cf0134fcbf272f18cd..c8543b19b95c1be44e9cf72f666f214de9da67f2:/stglibs/crypto.lib/bfstream.cpp?ds=sidebyside

diff --git a/stglibs/crypto.lib/bfstream.cpp b/stglibs/crypto.lib/bfstream.cpp
index 4f376c5d..3325bd1a 100644
--- a/stglibs/crypto.lib/bfstream.cpp
+++ b/stglibs/crypto.lib/bfstream.cpp
@@ -16,35 +16,42 @@ 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));
         }
 
         void Put(const void * data, size_t size, bool last)
         {
         size_t dataSize = m_ptr - m_buffer;
-        if (dataSize + size > sizeof(m_buffer))
+        while (dataSize + size > sizeof(m_buffer))
             {
             memcpy(m_ptr, data, sizeof(m_buffer) - dataSize); // Fill buffer
             size -= sizeof(m_buffer) - dataSize; // Adjust size
-            data += sizeof(m_buffer) - dataSize; // Adjust data pointer
+            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;
+            dataSize = 0;
             }
+        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;
@@ -52,6 +59,7 @@ class COMMON
         void * m_data;
         BLOWFISH_CTX m_ctx;
         PROC m_proc;
+        bool m_ok;
 
         void m_tryConsume(bool last)
         {
@@ -62,10 +70,12 @@ class COMMON
             dataSize += 8;
             remainder = 0;
             }
-        if (dataSize == 0)
+        if (!last && dataSize == 0) // Allow to call callback with 0 data on last call.
             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;
@@ -107,6 +117,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))
 {}
@@ -120,3 +135,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();
+}