5 files changed, 24 insertions, 15 deletions

diff --git a/protocols/WhatsAppWeb/src/iq.cpp b/protocols/WhatsAppWeb/src/iq.cpp
index 756f64c622..d47384067e 100644
--- a/protocols/WhatsAppWeb/src/iq.cpp
+++ b/protocols/WhatsAppWeb/src/iq.cpp
@@ -194,7 +194,7 @@ void WhatsAppProto::OnReceiveMessage(const WANode &node)
 		if (it->title != "enc" || it->content.length() == 0)
 			continue;
 
-		MBinBuffer msgBody;
+		SignalBuffer msgBody;
 		auto *pszType = it->getAttr("type");
 		try {
 			if (!mir_strcmp(pszType, "pkmsg") || !mir_strcmp(pszType, "msg")) {
@@ -206,10 +206,13 @@ void WhatsAppProto::OnReceiveMessage(const WANode &node)
 			}
 			else throw "Invalid e2e type";
 
+			if (!msgBody)
+				throw "Invalid e2e message";
+
 			iDecryptable++;
 
 			proto::Message encMsg;
-			encMsg << msgBody;
+			encMsg.ParseFromArray(msgBody.data(), msgBody.len());
 			if (encMsg.devicesentmessage().has_message())
 				encMsg = encMsg.devicesentmessage().message();
 
diff --git a/protocols/WhatsAppWeb/src/proto.h b/protocols/WhatsAppWeb/src/proto.h
index d6ffee6e20..7d0e7bdce0 100644
--- a/protocols/WhatsAppWeb/src/proto.h
+++ b/protocols/WhatsAppWeb/src/proto.h
@@ -198,8 +198,8 @@ public:
 
 	MSignalSession *createSession(const CMStringA &szName, int deviceId);
 
-	MBinBuffer decryptSignalProto(const CMStringA &from, const char *pszType, const MBinBuffer &encrypted);
-	MBinBuffer decryptGroupSignalProto(const CMStringA &from, const CMStringA &author, const MBinBuffer &encrypted);
+	signal_buffer* decryptSignalProto(const CMStringA &from, const char *pszType, const MBinBuffer &encrypted);
+	signal_buffer* decryptGroupSignalProto(const CMStringA &from, const CMStringA &author, const MBinBuffer &encrypted);
 
 	void generatePrekeys(int count);
 
diff --git a/protocols/WhatsAppWeb/src/signal.cpp b/protocols/WhatsAppWeb/src/signal.cpp
index 1e089ea6c0..ddb7da6dd0 100644
--- a/protocols/WhatsAppWeb/src/signal.cpp
+++ b/protocols/WhatsAppWeb/src/signal.cpp
@@ -444,7 +444,7 @@ MSignalSession* MSignalStore::createSession(const CMStringA &szName, int deviceI
 
 /////////////////////////////////////////////////////////////////////////////////////////
 
-MBinBuffer MSignalStore::decryptSignalProto(const CMStringA &from, const char *pszType, const MBinBuffer &encrypted)
+signal_buffer* MSignalStore::decryptSignalProto(const CMStringA &from, const char *pszType, const MBinBuffer &encrypted)
 {
 	WAJid jid(from);
 	auto *pSession = createSession(jid.user, 0);
@@ -471,18 +471,12 @@ MBinBuffer MSignalStore::decryptSignalProto(const CMStringA &from, const char *p
 		signal_message_destroy((signal_type_base *)pMsg);
 	}
 
-	MBinBuffer ret;
-	if (result != nullptr) {
-		ret.append(result->data, result->len);
-		signal_buffer_free(result);
-	}
-	return ret;
+	return result;
 }
 
-MBinBuffer MSignalStore::decryptGroupSignalProto(const CMStringA &group, const CMStringA &sender, const MBinBuffer &encrypted)
+signal_buffer* MSignalStore::decryptGroupSignalProto(const CMStringA &group, const CMStringA &sender, const MBinBuffer &encrypted)
 {
-	MBinBuffer ret;
-	return ret;
+	return nullptr;
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////
diff --git a/utils/mir_signal.cpp b/utils/mir_signal.cpp
index afca4f4188..d4e03e6803 100644
--- a/utils/mir_signal.cpp
+++ b/utils/mir_signal.cpp
@@ -37,4 +37,12 @@ SignalBuffer::~SignalBuffer()
 CMStringA SignalBuffer::toBase64() const
 {
 	return ptrA(mir_base64_encode(m_buf->data, m_buf->len)).get();
-}
\ No newline at end of file
+}
+
+void SignalBuffer::operator=(signal_buffer *buf)
+{
+	if (m_buf)
+		signal_buffer_free(m_buf);
+
+	m_buf = buf;
+}
diff --git a/utils/mir_signal.h b/utils/mir_signal.h
index ef443374ff..313714edac 100644
--- a/utils/mir_signal.h
+++ b/utils/mir_signal.h
@@ -40,6 +40,7 @@ class SignalBuffer
 	signal_buffer *m_buf;
 
 public:
+	SignalBuffer() : m_buf(nullptr) {}
 	SignalBuffer(const class MBinBuffer &buf);
 	SignalBuffer(const ec_public_key *key);
 	SignalBuffer(const ec_private_key *key);
@@ -47,6 +48,9 @@ public:
 	SignalBuffer(const session_signed_pre_key *pre_key);
 	~SignalBuffer();
 
+	void operator=(signal_buffer *buf);
+
+	__forceinline operator bool() const { return m_buf != nullptr; }
 	__forceinline uint8_t* data() const { return m_buf->data; }
 	__forceinline unsigned len() const { return unsigned(m_buf->len); }