3 files changed, 41 insertions, 17 deletions
diff --git a/Swiften/IDN/IDNA.cpp b/Swiften/IDN/IDNA.cpp
index 6b6c7a4..16b4183 100644
--- a/Swiften/IDN/IDNA.cpp
+++ b/Swiften/IDN/IDNA.cpp
@@ -4,7 +4,7 @@
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
 
-#include "Swiften/IDN/IDNA.h"
+#include <Swiften/IDN/IDNA.h>
 
 #include <stringprep.h>
 #include <vector>
diff --git a/Swiften/IDN/StringPrep.cpp b/Swiften/IDN/StringPrep.cpp
index d54fb0b..db09523 100644
--- a/Swiften/IDN/StringPrep.cpp
+++ b/Swiften/IDN/StringPrep.cpp
@@ -4,36 +4,58 @@
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
 
-#include "Swiften/IDN/StringPrep.h"
+#include <Swiften/IDN/StringPrep.h>
 
 #include <stringprep.h>
 #include <vector>
 #include <cassert>
+#include <Swiften/Base/SafeAllocator.h>
 
-namespace Swift {
+using namespace Swift;
+
+	namespace {
+	static const int MAX_STRINGPREP_SIZE = 1024;
 
-static const int MAX_STRINGPREP_SIZE = 1024;
+	const Stringprep_profile* getLibIDNProfile(StringPrep::Profile profile) {
+		switch(profile) {
+			case StringPrep::NamePrep: return stringprep_nameprep; break;
+			case StringPrep::XMPPNodePrep: return stringprep_xmpp_nodeprep; break;
+			case StringPrep::XMPPResourcePrep: return stringprep_xmpp_resourceprep; break;
+			case StringPrep::SASLPrep: return stringprep_saslprep; break;
+		}
+		assert(false);
+		return 0;
+	}
 
-const Stringprep_profile* getLibIDNProfile(StringPrep::Profile profile) {
-	switch(profile) {
-		case StringPrep::NamePrep: return stringprep_nameprep; break;
-		case StringPrep::XMPPNodePrep: return stringprep_xmpp_nodeprep; break;
-		case StringPrep::XMPPResourcePrep: return stringprep_xmpp_resourceprep; break;
-		case StringPrep::SASLPrep: return stringprep_saslprep; break;
+	template<typename StringType, typename ContainerType>
+	ContainerType getStringPrepared(const StringType& s, StringPrep::Profile profile) {
+		ContainerType input(s.begin(), s.end());
+		input.resize(MAX_STRINGPREP_SIZE);
+		if (stringprep(&input[0], MAX_STRINGPREP_SIZE, static_cast<Stringprep_profile_flags>(0), getLibIDNProfile(profile)) == 0) {
+			return input;
+		}
+		else {
+			return ContainerType();
+		}
 	}
-	assert(false);
-	return 0;
 }
 
+namespace Swift {
+
 std::string StringPrep::getPrepared(const std::string& s, Profile profile) {
-	std::vector<char> input(s.begin(), s.end());
-	input.resize(MAX_STRINGPREP_SIZE);
-	if (stringprep(&input[0], MAX_STRINGPREP_SIZE, static_cast<Stringprep_profile_flags>(0), getLibIDNProfile(profile)) == 0) {
-		return std::string(&input[0]);
+	std::vector<char> preparedData = getStringPrepared< std::string, std::vector<char> >(s, profile);
+	if (preparedData.empty()) {
+		throw std::exception();
 	}
-	else {
+	return std::string(vecptr(preparedData));
+}
+
+SafeByteArray StringPrep::getPrepared(const SafeByteArray& s, Profile profile) {
+	std::vector<char, SafeAllocator<char> > preparedData = getStringPrepared<SafeByteArray, std::vector<char, SafeAllocator<char> > >(s, profile);
+	if (preparedData.empty()) {
 		throw std::exception();
 	}
+	return createSafeByteArray(reinterpret_cast<const char*>(vecptr(preparedData)));
 }
 
 }
diff --git a/Swiften/IDN/StringPrep.h b/Swiften/IDN/StringPrep.h
index f40553b..dc8fadc 100644
--- a/Swiften/IDN/StringPrep.h
+++ b/Swiften/IDN/StringPrep.h
@@ -7,6 +7,7 @@
 #pragma once
 
 #include <string>
+#include <Swiften/Base/SafeByteArray.h>
 
 namespace Swift {
 	class StringPrep {
@@ -19,5 +20,6 @@ namespace Swift {
 			};
 
 			static std::string getPrepared(const std::string& s, Profile profile);
+			static SafeByteArray getPrepared(const SafeByteArray& s, Profile profile);
 	};
 }