From c2bdfab158da4be7d422a768a3d481cbbd814370 Mon Sep 17 00:00:00 2001
From: Felix Geyer <debfx@fobos.de>
Date: Thu, 16 Sep 2010 18:20:46 +0200
Subject: [PATCH] Implement writing to the SymmetricCipherStream.

---
 src/streams/SymmetricCipherStream.cpp | 87 +++++++++++++++++++++++----
 src/streams/SymmetricCipherStream.h   |  4 +-
 tests/TestSymmetricCipher.cpp         | 24 ++++++++
 3 files changed, 101 insertions(+), 14 deletions(-)

diff --git a/src/streams/SymmetricCipherStream.cpp b/src/streams/SymmetricCipherStream.cpp
index 483abefc3..30963bd7a 100644
--- a/src/streams/SymmetricCipherStream.cpp
+++ b/src/streams/SymmetricCipherStream.cpp
@@ -21,7 +21,8 @@ SymmetricCipherStream::SymmetricCipherStream(QIODevice* baseDevice, SymmetricCip
                       SymmetricCipher::Direction direction, const QByteArray& key, const QByteArray& iv)
                           : LayeredStream(baseDevice)
                           , m_bufferPos(0)
-                          , m_eof(false)
+                          , m_bufferFilling(false)
+                          , m_error(false)
 {
     m_cipher = new SymmetricCipher(algo, mode, direction, key, iv);
 }
@@ -30,22 +31,30 @@ bool SymmetricCipherStream::reset()
 {
     m_buffer.clear();
     m_bufferPos = 0;
+    m_bufferFilling = false;
+    m_error = false;
     m_cipher->reset();
     return true;
 }
 
+void SymmetricCipherStream::close()
+{
+    if (isWritable()) {
+        writeBlock();
+    }
+
+    LayeredStream::close();
+}
+
 qint64 SymmetricCipherStream::readData(char* data, qint64 maxSize)
 {
-    // TODO m_eof is probably wrong and should be removed
-    if (m_eof) {
-        return 0;
-    }
+    Q_ASSERT(maxSize >= 0);
 
     qint64 bytesRemaining = maxSize;
     qint64 offset = 0;
 
     while (bytesRemaining > 0) {
-        if (m_bufferPos == m_buffer.size()) {
+        if ((m_bufferPos == m_buffer.size()) || m_bufferFilling) {
             if (!readBlock()) {
                 return maxSize - bytesRemaining;
             }
@@ -65,28 +74,80 @@ qint64 SymmetricCipherStream::readData(char* data, qint64 maxSize)
 
 bool SymmetricCipherStream::readBlock()
 {
-    m_buffer = m_baseDevice->read(m_cipher->blockSize());
+    if (m_bufferFilling) {
+        m_buffer.append(m_baseDevice->read(m_cipher->blockSize() - m_buffer.size()));
+    }
+    else {
+        m_buffer = m_baseDevice->read(m_cipher->blockSize());
+    }
 
     if (m_buffer.size() != m_cipher->blockSize()) {
-        m_eof = true;
-        // TODO check if m_buffer.size()!=0
+        m_bufferFilling = true;
         return false;
     }
     else {
         m_cipher->processInPlace(m_buffer);
         m_bufferPos = 0;
+        m_bufferFilling = false;
         return true;
     }
 }
 
 qint64 SymmetricCipherStream::writeData(const char* data, qint64 maxSize)
 {
-    // TODO implement
-    return 0;
+    Q_ASSERT(maxSize >= 0);
+
+    if (m_error) {
+        return 0;
+    }
+
+    qint64 bytesRemaining = maxSize;
+    qint64 offset = 0;
+
+    while (bytesRemaining > 0) {
+        int bytesToCopy = qMin(bytesRemaining, static_cast<qint64>(m_cipher->blockSize() - m_buffer.size()));
+
+        m_buffer.append(data + offset, bytesToCopy);
+
+        offset += bytesToCopy;
+        bytesRemaining -= bytesToCopy;
+
+        if (m_buffer.size() == m_cipher->blockSize()) {
+            if (!writeBlock()) {
+                if (m_error) {
+                    return -1;
+                }
+                else {
+                    return maxSize - bytesRemaining;
+                }
+            }
+        }
+    }
+
+    return maxSize;
 }
 
 bool SymmetricCipherStream::writeBlock()
 {
-    // TODO implement
-    return false;
+    if (m_buffer.isEmpty()) {
+        return true;
+    }
+    else if (m_buffer.size() != m_cipher->blockSize()) {
+        int padLen = m_cipher->blockSize() - m_buffer.size();
+        for (int i=m_buffer.size(); i<m_cipher->blockSize(); i++) {
+            m_buffer.append(static_cast<char>(padLen));
+        }
+    }
+
+    m_cipher->processInPlace(m_buffer);
+
+    if (m_baseDevice->write(m_buffer) != m_buffer.size()) {
+        m_error = true;
+        // TODO copy error string
+        return false;
+    }
+    else {
+        m_buffer.clear();
+        return true;
+    }
 }
diff --git a/src/streams/SymmetricCipherStream.h b/src/streams/SymmetricCipherStream.h
index 1891f91c8..90bfc7257 100644
--- a/src/streams/SymmetricCipherStream.h
+++ b/src/streams/SymmetricCipherStream.h
@@ -31,6 +31,7 @@ public:
     SymmetricCipherStream(QIODevice* baseDevice, SymmetricCipher::Algorithm algo, SymmetricCipher::Mode mode,
                           SymmetricCipher::Direction direction, const QByteArray& key, const QByteArray& iv);
     bool reset();
+    void close();
 
 protected:
     qint64 readData(char* data, qint64 maxSize);
@@ -43,7 +44,8 @@ private:
     SymmetricCipher* m_cipher;
     QByteArray m_buffer;
     int m_bufferPos;
-    bool m_eof;
+    bool m_bufferFilling;
+    bool m_error;
 };
 
 #endif // KEEPASSX_SYMMETRICCIPHERSTREAM_H
diff --git a/tests/TestSymmetricCipher.cpp b/tests/TestSymmetricCipher.cpp
index 17feef1ba..70c6b6e42 100644
--- a/tests/TestSymmetricCipher.cpp
+++ b/tests/TestSymmetricCipher.cpp
@@ -53,6 +53,30 @@ void TestSymmetricCipher::testAes256CbcEncryption()
 
     QCOMPARE(QString(cipher.process(plainText).toHex()),
              QString(cipherText.toHex()));
+
+    QBuffer buffer;
+    SymmetricCipherStream stream(&buffer, SymmetricCipher::Aes256, SymmetricCipher::Cbc, SymmetricCipher::Encrypt, key, iv);
+    buffer.open(QIODevice::WriteOnly);
+    stream.open(QIODevice::WriteOnly);
+
+    buffer.reset();
+    buffer.buffer().clear();
+    stream.reset();
+    stream.write(plainText.left(16));
+    QCOMPARE(QString(buffer.data().toHex()), QString(cipherText.left(16).toHex()));
+
+    buffer.reset();
+    buffer.buffer().clear();
+    stream.reset();
+    stream.write(plainText.left(10));
+    QCOMPARE(QString(buffer.data().toHex()), QString());
+
+    buffer.reset();
+    buffer.buffer().clear();
+    stream.reset();
+    stream.write(plainText.left(10));
+    stream.close();
+    QCOMPARE(buffer.data().size(), 16);
 }
 
 void TestSymmetricCipher::testAes256CbcDecryption()