1 files changed, 40 insertions, 28 deletions
diff --git a/Swiften/StringCodecs/Base64.cpp b/Swiften/StringCodecs/Base64.cpp
index 4ec2e16..d8511b4 100644
--- a/Swiften/StringCodecs/Base64.cpp
+++ b/Swiften/StringCodecs/Base64.cpp
@@ -10,42 +10,54 @@
 
 #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.size();
-	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) {