WhatsApp: user's key retrieving

6 files changed, 159 insertions, 71 deletions

diff --git a/protocols/WhatsApp/src/iq.cpp b/protocols/WhatsApp/src/iq.cpp
index 2f8b7fb3d6..5e31f71372 100644
--- a/protocols/WhatsApp/src/iq.cpp
+++ b/protocols/WhatsApp/src/iq.cpp
@@ -82,12 +82,15 @@ void WhatsAppProto::OnIqDoNothing(const WANode&)
 
 /////////////////////////////////////////////////////////////////////////////////////////
 
-void WhatsAppProto::OnIqGetKeys(const WANode &node)
+void WhatsAppProto::OnIqGetKeys(const WANode &node, void *pUserInfo)
 {
-	if (auto *pList = node.getChild("list"))
-		for (auto &it : pList->getChildren())
-			if (it->title == "user")
-				m_signalStore.injectSession(it);
+	for (auto &it : node.getChild("list")->getChildren())
+		if (it->title == "user")
+			m_signalStore.injectSession(it);
+
+	// don't forget to send delayed message when all keys are retrieved
+	if (auto *pTask = (WASendTask *)pUserInfo)
+		SendTask(pTask);
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////
@@ -111,7 +114,7 @@ void WhatsAppProto::OnIqGetUsync(const WANode &node)
 				pKey->addChild("user")->addAttr("jid", it->toString());
 		}
 		if (pKey->getChildren().getCount() > 0)
-			WSSendNode(iq, &WhatsAppProto::OnIqGetKeys);
+			WSSendNode(iq, &WhatsAppProto::OnIqGetKeys, nullptr);
 	}
 }
 
@@ -206,7 +209,6 @@ void WhatsAppProto::OnIqPairSuccess(const WANode &node)
 				signal_buffer_free(result);
 			}
 
-			setDword("SignalDeviceId", 0);
 			{
 				MBinBuffer key;
 				if (accountSignatureKey.len == 32)
@@ -217,6 +219,7 @@ void WhatsAppProto::OnIqPairSuccess(const WANode &node)
 
 			proto::CleanBinary(account->accountsignaturekey); account->has_accountsignaturekey = false;
 			MBinBuffer accountEnc(proto::Serialize(account));
+			db_set_blob(0, m_szModuleName, "WAAccount", accountEnc.data(), (int)accountEnc.length());
 
 			proto::ADVDeviceIdentity deviceIdentity(deviceDetails);
 
@@ -244,7 +247,10 @@ void WhatsAppProto::OnIqResult(const WANode &node)
 	if (auto *pszId = node.getAttr("id"))
 		for (auto &it : m_arPacketQueue)
 			if (it->szPacketId == pszId)
-				(this->*it->pHandler)(node);
+				if (it->pUserInfo)
+					(this->*it->pHandlerFull)(node, it->pUserInfo);
+				else
+					(this->*it->pHandler)(node);
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////
@@ -444,18 +450,8 @@ void WhatsAppProto::OnReceiveInfo(const WANode &node)
 			if (!m_bUpdatedPrekeys)
 				WSSendNode(WANodeIq(IQ::GET, "encrypt") << XCHILD("count"), &WhatsAppProto::OnIqCountPrekeys);
 
-			if (m_arDevices.getCount() == 0) {
-				WANodeIq iq(IQ::GET, "usync");
-
-				auto *pNode1 = iq.addChild("usync");
-				*pNode1 << CHAR_PARAM("sid", GenerateMessageId()) << CHAR_PARAM("mode", "query") << CHAR_PARAM("last", "true")
-					<< CHAR_PARAM("index", "0") << CHAR_PARAM("context", "message");
-
-				pNode1->addChild("query")->addChild("devices")->addAttr("version", "2");
-				pNode1->addChild("list")->addChild("user")->addAttr("jid", m_szJid);
-
-				WSSendNode(iq, &WhatsAppProto::OnIqGetUsync);
-			}
+			if (m_arDevices.getCount() == 0)
+				SendUsync(m_szJid);
 
 			for (auto &it : m_arCollections) {
 				if (it->version == 0) {
diff --git a/protocols/WhatsApp/src/message.cpp b/protocols/WhatsApp/src/message.cpp
index 3a768961ee..eb221e9700 100644
--- a/protocols/WhatsApp/src/message.cpp
+++ b/protocols/WhatsApp/src/message.cpp
@@ -341,9 +341,6 @@ void WhatsAppProto::OnReceiveAck(const WANode &node)
 
 bool WhatsAppProto::CreateMsgParticipant(WANode *pParticipants, const WAJid &jid, const MBinBuffer &orig)
 {
-	if (jid.device == (int)getDword(DBKEY_DEVICE_ID))
-		return false;
-
 	int type = 0;
 
 	try {
@@ -366,25 +363,13 @@ int WhatsAppProto::SendTextMessage(const char *jid, const char *pszMsg)
 {
 	WAJid toJid(jid);
 
-	char szMsgId[40];
-	__int64 msgId;
-	Utils_GetRandom(&msgId, sizeof(msgId));
-	if (msgId < 0)
-		msgId = -msgId;
-	_i64toa(msgId, szMsgId, 10);
-
-	// mother node for all participants
-	WANode payLoad("message");
-	payLoad << CHAR_PARAM("id", szMsgId) << CHAR_PARAM("type", "text") << CHAR_PARAM("to", jid);
-
-	auto *pParticipants = payLoad.addChild("participants");
+	// send task creation
+	auto *pTask = new WASendTask(jid);
 
 	// basic message 
 	Wa__Message body;
 	body.conversation = (char*)pszMsg;
 
-	bool shouldIncludeIdentity = false;
-
 	if (toJid.isGroup()) {
 		MBinBuffer encodedMsg(proto::Serialize(&body));
 		padBuffer16(encodedMsg);
@@ -392,7 +377,7 @@ int WhatsAppProto::SendTextMessage(const char *jid, const char *pszMsg)
 		MBinBuffer skmsgKey;
 		MBinBuffer cipherText(m_signalStore.encryptSenderKey(toJid, m_szJid, encodedMsg, skmsgKey));
 
-		auto *pEnc = payLoad.addChild("enc");
+		auto *pEnc = pTask->payLoad.addChild("enc");
 		*pEnc << CHAR_PARAM("v", "2") << CHAR_PARAM("type", "skmsg");
 		pEnc->content.append(cipherText.data(), cipherText.length());
 
@@ -405,16 +390,12 @@ int WhatsAppProto::SendTextMessage(const char *jid, const char *pszMsg)
 		Wa__Message msg;
 		msg.senderkeydistributionmessage = &sentBody;
 
-		MBinBuffer encodedMeMsg(proto::Serialize(&msg));
-		padBuffer16(encodedMeMsg);
+		pTask->content.append(proto::Serialize(&msg));
 
 		if (auto *pUser = FindUser(jid))
 			if (pUser->si)
 				for (auto &it : pUser->si->arUsers)
-					shouldIncludeIdentity = CreateMsgParticipant(pParticipants, WAJid(T2Utf(it->pszUID)), encodedMeMsg);
-
-		for (auto &it : m_arDevices)
-			shouldIncludeIdentity |= CreateMsgParticipant(pParticipants, *it, encodedMeMsg);
+					pTask->arDest.insert(new WAJid(T2Utf(it->pszUID)));
 	}
 	else {
 		Wa__Message__DeviceSentMessage sentBody;
@@ -424,24 +405,63 @@ int WhatsAppProto::SendTextMessage(const char *jid, const char *pszMsg)
 		Wa__Message msg;
 		msg.devicesentmessage = &sentBody;
 
-		MBinBuffer encodedMeMsg(proto::Serialize(&msg));
-		padBuffer16(encodedMeMsg);
+		pTask->content.append(proto::Serialize(&msg));
+
+		pTask->arDest.insert(new WAJid(toJid));
+	}
+
+	padBuffer16(pTask->content);
+
+	for (auto &it : m_arDevices)
+		if (it->device != (int)getDword(DBKEY_DEVICE_ID))
+			pTask->arDest.insert(new WAJid(*it));
+
+	// check session presence first
+	LIST<WAJid> missingKeys(1);
+	for (auto &it : pTask->arDest) {
+		MSignalSession tmp(it->user, it->device);
+		auto blob = getBlob(tmp.getSetting());
+		if (blob.isEmpty())
+			missingKeys.insert(it);
+	}
 
-		shouldIncludeIdentity = CreateMsgParticipant(pParticipants, toJid, encodedMeMsg);
-		for (auto &it : m_arDevices)
-			shouldIncludeIdentity |= CreateMsgParticipant(pParticipants, *it, encodedMeMsg);
+	// generate & reserve packet id
+	int pktId;
+	{
+		mir_cslock lck(m_csOwnMessages);
+		pktId = m_iPacketId++;
+		m_arOwnMsgs.insert(new WAOwnMessage(pktId, jid, pTask->szMsgId));
 	}
 
+	// if some keys are missing, schedule task for execution & retrieve keys
+	if (missingKeys.getCount()) {
+		WANodeIq iq(IQ::GET, "encrypt");
+		auto *pKey = iq.addChild("key");
+		for (auto &it: missingKeys)
+			pKey->addChild("user")->addAttr("jid", it->toString());
+		WSSendNode(iq, &WhatsAppProto::OnIqGetKeys, pTask);
+	}
+	// otherwise simply execute the task
+	else SendTask(pTask);
+
+	return pktId;
+}
+
+void WhatsAppProto::SendTask(WASendTask *pTask)
+{
+	// pack all data and send the whole payload
+	bool shouldIncludeIdentity = false;
+	auto *pParticipants = pTask->payLoad.addChild("participants");
+
+	for (auto &it : pTask->arDest)
+		shouldIncludeIdentity |= CreateMsgParticipant(pParticipants, *it, pTask->content);
+
 	if (shouldIncludeIdentity) {
 		MBinBuffer encIdentity(m_signalStore.encodeSignedIdentity(true));
-		auto *pNode = payLoad.addChild("device-identity");
+		auto *pNode = pTask->payLoad.addChild("device-identity");
 		pNode->content.assign(encIdentity.data(), encIdentity.length());
 	}
 
-	WSSendNode(payLoad);
-	
-	mir_cslock lck(m_csOwnMessages);
-	int pktId = m_iPacketId++;
-	m_arOwnMsgs.insert(new WAOwnMessage(pktId, jid, szMsgId));
-	return pktId;
+	WSSendNode(pTask->payLoad);
+	delete pTask;
 }
diff --git a/protocols/WhatsApp/src/proto.h b/protocols/WhatsApp/src/proto.h
index 36ecbfcbf9..290aa6c392 100644
--- a/protocols/WhatsApp/src/proto.h
+++ b/protocols/WhatsApp/src/proto.h
@@ -14,6 +14,7 @@ Copyright © 2019-22 George Hazan
 
 class WhatsAppProto;
 typedef void (WhatsAppProto:: *WA_PKT_HANDLER)(const WANode &node);
+typedef void (WhatsAppProto:: *WA_PKT_HANDLER_FULL)(const WANode &node, void *pUserInfo);
 
 struct WAMSG
 {
@@ -42,14 +43,23 @@ struct WAMediaKeys
 
 struct WARequest
 {
-	WARequest(const CMStringA &_1, WA_PKT_HANDLER _2, void *_3 = nullptr) :
+	WARequest(const CMStringA &_1, WA_PKT_HANDLER _2) :
 		szPacketId(_1),
 		pHandler(_2),
+		pUserInfo(nullptr)
+	{}
+
+	WARequest(const CMStringA &_1, WA_PKT_HANDLER_FULL _2, void *_3) :
+		szPacketId(_1),
+		pHandlerFull(_2),
 		pUserInfo(_3)
 	{}
 
 	CMStringA szPacketId;
-	WA_PKT_HANDLER pHandler;
+	union {
+		WA_PKT_HANDLER pHandler;
+		WA_PKT_HANDLER_FULL pHandlerFull;
+	};
 	void *pUserInfo;
 };
 
@@ -317,7 +327,9 @@ class WhatsAppProto : public PROTO<WhatsAppProto>
 
 	bool WSReadPacket(const WSHeader &hdr, MBinBuffer &buf);
 	int  WSSend(const ProtobufCMessage &msg);
-	int  WSSendNode(WANode &node, WA_PKT_HANDLER = nullptr);
+	int  WSSendNode(WANode &node);
+	int  WSSendNode(WANode &node, WA_PKT_HANDLER);
+	int  WSSendNode(WANode &node, WA_PKT_HANDLER_FULL, void *pUserInfo);
 
 	MBinBuffer DownloadEncryptedFile(const char *url, const ProtobufCBinaryData &mediaKeys, const char *pszType);
 	CMStringW  GetTmpFileName(const char *pszClass, const char *addition);
@@ -331,7 +343,9 @@ class WhatsAppProto : public PROTO<WhatsAppProto>
 	void SendAck(const WANode &node);
 	void SendReceipt(const char *pszTo, const char *pszParticipant, const char *pszId, const char *pszType);
 	void SendKeepAlive();
+	void SendTask(WASendTask *pTask);
 	int  SendTextMessage(const char *jid, const char *pszMsg);
+	void SendUsync(const char *jid);
 	void SetServerStatus(int iStatus);
 
 	/// Popups /////////////////////////////////////////////////////////////////////////////
@@ -353,7 +367,7 @@ class WhatsAppProto : public PROTO<WhatsAppProto>
 	void OnIqDoNothing(const WANode &node);
 	void OnIqGcGetAllMetadata(const WANode &node);
 	void OnIqGetAvatar(const WANode &node);
-	void OnIqGetKeys(const WANode &node);
+	void OnIqGetKeys(const WANode &node, void *pUserInfo);
 	void OnIqGetUsync(const WANode &node);
 	void OnIqPairDevice(const WANode &node);
 	void OnIqPairSuccess(const WANode &node);
diff --git a/protocols/WhatsApp/src/server.cpp b/protocols/WhatsApp/src/server.cpp
index 5b14f120ab..1fce9ed3a7 100644
--- a/protocols/WhatsApp/src/server.cpp
+++ b/protocols/WhatsApp/src/server.cpp
@@ -399,6 +399,20 @@ void WhatsAppProto::SetServerStatus(int iStatus)
 			&WhatsAppProto::OnIqDoNothing);
 }
 
+void WhatsAppProto::SendUsync(const char *jid)
+{
+	WANodeIq iq(IQ::GET, "usync");
+
+	auto *pNode1 = iq.addChild("usync");
+	*pNode1 << CHAR_PARAM("sid", GenerateMessageId()) << CHAR_PARAM("mode", "query") << CHAR_PARAM("last", "true")
+		<< CHAR_PARAM("index", "0") << CHAR_PARAM("context", "message");
+
+	pNode1->addChild("query")->addChild("devices")->addAttr("version", "2");
+	pNode1->addChild("list")->addChild("user")->addAttr("jid", jid);
+
+	WSSendNode(iq, &WhatsAppProto::OnIqGetUsync);
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////
 
 void WhatsAppProto::ShutdownSession()
diff --git a/protocols/WhatsApp/src/utils.cpp b/protocols/WhatsApp/src/utils.cpp
index de72fd5b92..2993af9579 100644
--- a/protocols/WhatsApp/src/utils.cpp
+++ b/protocols/WhatsApp/src/utils.cpp
@@ -185,21 +185,11 @@ int WhatsAppProto::WSSend(const ProtobufCMessage &msg)
 
 /////////////////////////////////////////////////////////////////////////////////////////
 
-static char zeroData[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
-
-int WhatsAppProto::WSSendNode(WANode &node, WA_PKT_HANDLER pHandler)
+int WhatsAppProto::WSSendNode(WANode &node)
 {
 	if (m_hServerConn == nullptr)
 		return 0;
 
-	if (pHandler != nullptr) {
-		CMStringA id(GenerateMessageId());
-		node.addAttr("id", id);
-
-		mir_cslock lck(m_csPacketQueue);
-		m_arPacketQueue.insert(new WARequest(id, pHandler));
-	}
-
 	CMStringA szText;
 	node.print(szText);
 	debugLogA("Sending binary node:\n%s", szText.c_str());
@@ -213,6 +203,36 @@ int WhatsAppProto::WSSendNode(WANode &node, WA_PKT_HANDLER pHandler)
 	return 1;
 }
 
+int WhatsAppProto::WSSendNode(WANode &node, WA_PKT_HANDLER pHandler)
+{
+	if (m_hServerConn == nullptr)
+		return 0;
+
+	CMStringA id(GenerateMessageId());
+	node.addAttr("id", id);
+	{
+		mir_cslock lck(m_csPacketQueue);
+		m_arPacketQueue.insert(new WARequest(id, pHandler));
+	}
+
+	return WSSendNode(node);
+}
+
+int WhatsAppProto::WSSendNode(WANode &node, WA_PKT_HANDLER_FULL pHandler, void *pUserInfo)
+{
+	if (m_hServerConn == nullptr)
+		return 0;
+
+	CMStringA id(GenerateMessageId());
+	node.addAttr("id", id);
+	{
+		mir_cslock lck(m_csPacketQueue);
+		m_arPacketQueue.insert(new WARequest(id, pHandler, pUserInfo));
+	}
+
+	return WSSendNode(node);
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////
 
 std::string decodeBinStr(const std::string &buf)
diff --git a/protocols/WhatsApp/src/utils.h b/protocols/WhatsApp/src/utils.h
index 37e48793b2..d55e816d8c 100644
--- a/protocols/WhatsApp/src/utils.h
+++ b/protocols/WhatsApp/src/utils.h
@@ -172,6 +172,30 @@ struct WAJid
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////
+// WASendTask
+
+struct WASendTask
+{
+	WASendTask(const char *jid) :
+		payLoad("message"),
+		arDest(1)
+	{
+		__int64 msgId;
+		Utils_GetRandom(&msgId, sizeof(msgId));
+		if (msgId < 0)
+			msgId = -msgId;
+		_i64toa(msgId, szMsgId, 10);
+
+		payLoad << CHAR_PARAM("id", szMsgId) << CHAR_PARAM("type", "text") << CHAR_PARAM("to", jid);
+	}
+
+	char szMsgId[40];
+	WANode payLoad;
+	OBJLIST<WAJid> arDest;
+	MBinBuffer content;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////
 // LT_HASH
 
 struct LT_HASH