1 files changed, 17 insertions, 8 deletions
diff --git a/Swiften/StringCodecs/HMAC.h b/Swiften/StringCodecs/HMAC.h
index 438a3a7..cf0abfe 100644
--- a/Swiften/StringCodecs/HMAC.h
+++ b/Swiften/StringCodecs/HMAC.h
@@ -6,28 +6,37 @@
 
 #pragma once
 
 #include <Swiften/Base/ByteArray.h>
 #include <Swiften/Base/Algorithm.h>
 #include <cassert>
 
 namespace Swift {
 	namespace HMAC_Detail {
-		static const unsigned int B = 64;
+		template<typename KeyType> struct KeyWrapper;
+		template<> struct KeyWrapper<ByteArray> {
+			ByteArray wrap(const ByteArray& hash) const {
+				return hash;
+			}
+		};
+		template<> struct KeyWrapper<SafeByteArray> {
+			SafeByteArray wrap(const ByteArray& hash) const {
+				return createSafeByteArray(hash);
+			}
+		};
 
-		template<typename Hash, typename KeyType>
+		template<typename Hash, typename KeyType, int BlockSize>
 		static ByteArray getHMAC(const KeyType& key, const ByteArray& data) {
-			assert(key.size() <= B);
 			Hash hash;
 
 			// Create the padded key
-			KeyType paddedKey(key);
-			paddedKey.resize(B, 0x0);
+			KeyType paddedKey(key.size() <= BlockSize ? key : KeyWrapper<KeyType>().wrap(hash(key)));
+			paddedKey.resize(BlockSize, 0x0);
 
 			// Create the first value
 			KeyType x(paddedKey);
 			for (unsigned int i = 0; i < x.size(); ++i) {
 				x[i] ^= 0x36;
 			}
 			append(x, data);
 
 			// Create the second value
@@ -35,23 +44,23 @@ namespace Swift {
 			for (unsigned int i = 0; i < y.size(); ++i) {
 				y[i] ^= 0x5c;
 			}
 			append(y, hash(x));
 
 			return hash(y);
 		}
 	};
 
-	template<typename Hash>
+	template<typename Hash, int BlockSize>
 	class HMAC {
 		private:
 
 		public:
 			ByteArray operator()(const ByteArray& key, const ByteArray& data) const {
-				return HMAC_Detail::getHMAC<Hash,ByteArray>(key, data);
+				return HMAC_Detail::getHMAC<Hash,ByteArray,BlockSize>(key, data);
 			}
 
 			ByteArray operator()(const SafeByteArray& key, const ByteArray& data) const {
-				return HMAC_Detail::getHMAC<Hash,SafeByteArray>(key, data);
+				return HMAC_Detail::getHMAC<Hash,SafeByteArray,BlockSize>(key, data);
 			}
 	};
 }