1 files changed, 44 insertions, 30 deletions
diff --git a/Swiften/StringCodecs/Base64.cpp b/Swiften/StringCodecs/Base64.cpp
index 795d982..d8511b4 100644
--- a/Swiften/StringCodecs/Base64.cpp
+++ b/Swiften/StringCodecs/Base64.cpp
@@ -5,50 +5,64 @@
  */
 
 #include <boost/numeric/conversion/cast.hpp>
+
 #include <algorithm>
 
-#include "Swiften/StringCodecs/Base64.h"
+#include <Swiften/StringCodecs/Base64.h>
+#include <Swiften/Base/Algorithm.h>
+#include <Swiften/Base/SafeString.h>
 
 namespace Swift {
 
 #pragma GCC diagnostic ignored "-Wold-style-cast"
 
-std::string Base64::encode(const ByteArray &s) {
-	int i;
-	int len = s.getSize();
-	char tbl[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=";
-	int a, b, c;
-
-	std::string p;
-	p.resize((len+2)/3*4);
-	int at = 0;
-	for( i = 0; i < len; i += 3 ) {
-		a = ((unsigned char) (s[i]) & 3) << 4;
-		if(i + 1 < len) {
-			a += (unsigned char) (s[i + 1]) >> 4;
-			b = ((unsigned char) (s[i + 1]) & 0xF) << 2;
-			if(i + 2 < len) {
-				b += (unsigned char) (s[i + 2]) >> 6;
-				c = (unsigned char) (s[i + 2]) & 0x3F;
+namespace {
+	template<typename TargetType, typename SourceType>
+	TargetType base64Encode(const SourceType& s) {
+		int i;
+		int len = s.size();
+		char tbl[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=";
+		int a, b, c;
+
+		TargetType p;
+		p.resize((len+2)/3*4);
+		int at = 0;
+		for( i = 0; i < len; i += 3 ) {
+			a = ((unsigned char) (s[i]) & 3) << 4;
+			if(i + 1 < len) {
+				a += (unsigned char) (s[i + 1]) >> 4;
+				b = ((unsigned char) (s[i + 1]) & 0xF) << 2;
+				if(i + 2 < len) {
+					b += (unsigned char) (s[i + 2]) >> 6;
+					c = (unsigned char) (s[i + 2]) & 0x3F;
+				}
+				else
+					c = 64;
+			}
+			else {
+				b = c = 64;
 			}
-			else
-				c = 64;
-		}
-		else {
-			b = c = 64;
-		}
 
-		p[at++] = tbl[(unsigned char) (s[i]) >> 2];
-		p[at++] = tbl[a];
-		p[at++] = tbl[b];
-		p[at++] = tbl[c];
+			p[at++] = tbl[(unsigned char) (s[i]) >> 2];
+			p[at++] = tbl[a];
+			p[at++] = tbl[b];
+			p[at++] = tbl[c];
+		}
+		return p;
 	}
-	return p;
+}
+
+std::string Base64::encode(const ByteArray &s) {
+	return base64Encode<std::string, ByteArray>(s);
+}
+
+SafeString Base64::encode(const SafeByteArray &s) {
+	return base64Encode<SafeString, SafeByteArray>(s);
 }
 
 ByteArray Base64::decode(const std::string& input) {
 	std::string inputWithoutNewlines(input);
-	inputWithoutNewlines.erase(std::remove(inputWithoutNewlines.begin(), inputWithoutNewlines.end(), '\n'), inputWithoutNewlines.end());
+	erase(inputWithoutNewlines, '\n');
 
 	const std::string& s = inputWithoutNewlines;
 	ByteArray p;