From 711f5bcc6210bf2806ae477739d771303dcd6e65 Mon Sep 17 00:00:00 2001
From: orignal <i2porignal@yandex.ru>
Date: Mon, 7 Apr 2025 20:33:19 -0400
Subject: [PATCH] store ML-KEM section for possible next NSR

---
 libi2pd/ECIESX25519AEADRatchetSession.cpp | 21 +++++++++++++++++++--
 libi2pd/ECIESX25519AEADRatchetSession.h   |  2 ++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/libi2pd/ECIESX25519AEADRatchetSession.cpp b/libi2pd/ECIESX25519AEADRatchetSession.cpp
index 467a3752..cfcbf5bb 100644
--- a/libi2pd/ECIESX25519AEADRatchetSession.cpp
+++ b/libi2pd/ECIESX25519AEADRatchetSession.cpp
@@ -665,6 +665,8 @@ namespace garlic
 				LogPrint (eLogWarning, "Garlic: NSR ML-KEM ciphertext section AEAD encryption failed");
 				return false;
 			}
+			m_NSREncodedPQKey = std::make_unique<std::array<uint8_t, i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16> >();
+			memcpy (m_NSREncodedPQKey->data (), out + offset, i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16);
 			MixHash (out + offset, i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16);
 			MixKey (sharedSecret);
 			offset += i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16;
@@ -718,8 +720,19 @@ namespace garlic
 		size_t offset = 40;
 #if OPENSSL_PQ		
 		if (m_PQKeys)
-			// TODO: encrypted ML-KEM section
-			offset += i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16;
+		{	
+			if (m_NSREncodedPQKey)
+			{	
+				memcpy (out + offset, m_NSREncodedPQKey->data (), i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16);
+				MixHash (out + offset, i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16);
+				offset += i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16;
+			}
+			else
+			{
+				LogPrint (eLogWarning, "Garlic: No stored ML-KEM keys");
+				return false;
+			}
+		}
 #endif
 		if (!Encrypt (m_NSRH /* can be anything */, out + offset, 0)) // encrypt, ciphertext = ENCRYPT(k, n, ZEROLEN, ad)
 		{
@@ -914,6 +927,10 @@ namespace garlic
 				m_State = eSessionStateEstablished;
 				m_NSRSendTagset = nullptr;
 				m_EphemeralKeys = nullptr;
+#if OPENSSL_PQ
+				m_PQKeys = nullptr;
+				m_NSREncodedPQKey = nullptr;
+#endif				
 				[[fallthrough]];
 			case eSessionStateEstablished:
 				if (m_SendReverseKey && receiveTagset->GetTagSetID () == m_NextReceiveRatchet->GetReceiveTagSetID ())
diff --git a/libi2pd/ECIESX25519AEADRatchetSession.h b/libi2pd/ECIESX25519AEADRatchetSession.h
index cd32cb98..b988263a 100644
--- a/libi2pd/ECIESX25519AEADRatchetSession.h
+++ b/libi2pd/ECIESX25519AEADRatchetSession.h
@@ -14,6 +14,7 @@
 #include <functional>
 #include <memory>
 #include <vector>
+#include <array>
 #include <list>
 #include <unordered_map>
 #include "Identity.h"
@@ -228,6 +229,7 @@ namespace garlic
 			std::shared_ptr<i2p::crypto::X25519Keys> m_EphemeralKeys;
 #if OPENSSL_PQ	
 			std::unique_ptr<i2p::crypto::MLKEM512Keys> m_PQKeys;
+			std::unique_ptr<std::array<uint8_t, i2p::crypto::MLKEM512_CIPHER_TEXT_LENGTH + 16> > m_NSREncodedPQKey;
 #endif			
 			SessionState m_State = eSessionStateNew;
 			uint64_t m_SessionCreatedTimestamp = 0, m_LastActivityTimestamp = 0, // incoming (in seconds)