1 files changed, 46 insertions, 13 deletions
diff --git a/Swiften/StringCodecs/MD5.cpp b/Swiften/StringCodecs/MD5.cpp
index 718f52e..6871f79 100644
--- a/Swiften/StringCodecs/MD5.cpp
+++ b/Swiften/StringCodecs/MD5.cpp
@@ -33,12 +33,13 @@
 
 #pragma GCC diagnostic ignored "-Wold-style-cast"
 
-#include "Swiften/StringCodecs/MD5.h"
+#include <Swiften/StringCodecs/MD5.h>
 
 #include <cassert>
+#include <string.h>
 
-#include "Swiften/Base/ByteArray.h"
-#include "Swiften/Base/Platform.h"
+#include <Swiften/Base/ByteArray.h>
+#include <Swiften/Base/Platform.h>
 
 namespace Swift {
 
@@ -124,13 +125,13 @@ static void md5_process(md5_state_t *pms, const md5_byte_t *data /*[64]*/) {
 		a = pms->abcd[0], b = pms->abcd[1],
 		c = pms->abcd[2], d = pms->abcd[3];
 		md5_word_t t;
-#ifdef SWIFTEN_BIG_ENDIAN
-		/* Define storage only for big-endian CPUs. */
-		md5_word_t X[16];
-#else
+#ifdef SWIFTEN_LITTLE_ENDIAN
 		/* Define storage for little-endian or both types of CPUs. */
 		md5_word_t xbuf[16];
 		const md5_word_t *X;
+#else
+		/* Define storage only for big-endian CPUs. */
+		md5_word_t X[16];
 #endif
 
 		{
@@ -350,16 +351,48 @@ md5_finish(md5_state_t *pms, md5_byte_t digest[16])
 				digest[i] = (md5_byte_t)(pms->abcd[i >> 2] >> ((i & 3) << 3));
 }
 
-ByteArray MD5::getHash(const ByteArray& data) {
+namespace {
+	template<typename SourceType>
+	ByteArray getMD5Hash(const SourceType& data) {
+		ByteArray digest;
+		digest.resize(16);
+
+		md5_state_t state;
+		md5_init(&state);
+		md5_append(&state, reinterpret_cast<const md5_byte_t*>(vecptr(data)), data.size());
+		md5_finish(&state, reinterpret_cast<md5_byte_t*>(vecptr(digest)));
+
+		return digest;
+	}
+}
+
+MD5::MD5() {
+	state = new md5_state_t;
+	md5_init(state);
+}
+
+MD5::~MD5() {
+	delete state;
+}
+
+MD5& MD5::update(const std::vector<unsigned char>& input) {
+	md5_append(state, reinterpret_cast<const md5_byte_t*>(vecptr(input)), input.size());
+	return *this;
+}
+
+std::vector<unsigned char> MD5::getHash() {
 	ByteArray digest;
 	digest.resize(16);
+	md5_finish(state, reinterpret_cast<md5_byte_t*>(vecptr(digest)));
+	return digest;
+}
 
-	md5_state_t state;
-	md5_init(&state);
-	md5_append(&state, reinterpret_cast<const md5_byte_t*>(data.getData()), data.getSize());
-	md5_finish(&state, reinterpret_cast<md5_byte_t*>(digest.getData()));
+ByteArray MD5::getHash(const ByteArray& data) {
+	return getMD5Hash(data);
+}
 
-	return digest;
+ByteArray MD5::getHash(const SafeByteArray& data) {
+	return getMD5Hash(data);
 }
 
 }