From 686bbdeb9c1962207f1faeae6646af94fd1007b8 Mon Sep 17 00:00:00 2001
From: Grzegorz Rychlik <Grzegorz.Rychlik@f-secure.com>
Date: Thu, 14 Nov 2019 17:49:30 +0100
Subject: [PATCH] Send the snapshot over ApiBridge only if it changed

---
 Src/Core/GateRelay.cpp | 19 ++++++++++++-------
 Src/Core/Profiler.cpp  | 12 +++++-------
 Src/Core/Profiler.h    | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 14 deletions(-)

diff --git a/Src/Core/GateRelay.cpp b/Src/Core/GateRelay.cpp
index a9826b2..274df43 100644
--- a/Src/Core/GateRelay.cpp
+++ b/Src/Core/GateRelay.cpp
@@ -168,18 +168,23 @@ void MWR::C3::Core::GateRelay::RunApiBrige(std::string_view apiBrigdeIp, std::ui
 			Log({ "API bridge connection established on " + std::string{apiBrigdeIp} +':' + std::to_string(apiBrigdePort), MWR::C3::LogMessage::Severity::Information });
 			reconnectWait = 0s;
 			// Enter main loop.
+			auto sp = m_Profiler->GetSnapshotProxy();
 			while (m_IsAlive && connection.IsSending())
 			{
 				// Read socket.
 				std::this_thread::sleep_for(300ms);
-				try
+				auto newSnapshot = sp.GetSnapshotIfChanged();
+				if (newSnapshot)
 				{
-					connection.Send(Crypto::Encrypt(ByteView{ json{ { "messageType", "GetProfile" }, { "messageData", m_Profiler->Get().m_Gateway.CreateProfileSnapshot() }}.dump() }, m_SessionKeys.second));
-				}
-				catch (std::exception& exception)
-				{
-					Log({ "Caught an exception while sending Profile. "s + exception.what(), MWR::C3::LogMessage::Severity::Error });
-					break;
+					try
+					{
+						connection.Send(Crypto::Encrypt(ByteView{ json{ { "messageType", "GetProfile" }, { "messageData", *newSnapshot }}.dump() }, m_SessionKeys.second));
+					}
+					catch (std::exception& exception)
+					{
+						Log({ "Caught an exception while sending Profile. "s + exception.what(), MWR::C3::LogMessage::Severity::Error });
+						break;
+					}
 				}
 			}
 		}
diff --git a/Src/Core/Profiler.cpp b/Src/Core/Profiler.cpp
index 82ba20e..55fc09f 100644
--- a/Src/Core/Profiler.cpp
+++ b/Src/Core/Profiler.cpp
@@ -172,18 +172,16 @@ MWR::ByteVector MWR::C3::Core::Profiler::Translate(std::string const& type, json
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 void MWR::C3::Core::Profiler::DumpSnapshots()
 {
-	std::optional<std::size_t> oldHash;
+	auto sp = GetSnapshotProxy();
 	while (true)
 	{
-		const auto snapshotTmpPath = std::filesystem::path(m_SnapshotPath).replace_extension(".tmp");
-		const auto snapshot = Get().m_Gateway.CreateProfileSnapshot().dump(4);
-		const auto snapshotHash = std::hash<std::string>{}(snapshot);
-		if (!oldHash || *oldHash != snapshotHash)
+		const auto newSnapshot = sp.GetSnapshotIfChanged();
+		if (newSnapshot)
 		{
-			oldHash = snapshotHash;
+			const auto snapshotTmpPath = std::filesystem::path(m_SnapshotPath).replace_extension(".tmp");
 			{
 				std::ofstream snapshotTmp{ snapshotTmpPath };
-				snapshotTmp << snapshot << std::endl;
+				snapshotTmp << newSnapshot->dump(4) << std::endl;
 			}
 			std::filesystem::rename(snapshotTmpPath, m_SnapshotPath);
 		}
diff --git a/Src/Core/Profiler.h b/Src/Core/Profiler.h
index 31ab15a..6c2649c 100644
--- a/Src/Core/Profiler.h
+++ b/Src/Core/Profiler.h
@@ -536,6 +536,42 @@ namespace MWR::C3::Core
 		/// @returns binder id
 		uint32_t GetBinderTo(uint32_t);
 
+		/// Helper to wrap calls to CreateProfileShnapshot
+		class SnapshotProxy
+		{
+		public:
+			/// Create a snapshot proxy
+			/// @param profiler to wrap CreateProfileShnapshot calls
+			SnapshotProxy(Profiler& profiler) :
+				m_Profiler(profiler)
+			{
+			}
+
+			/// Create a snapshot
+			/// @return std::nullopt if snaphot hasn't change since the last call
+			std::optional<json> GetSnapshotIfChanged()
+			{
+				auto currentSnapshot = m_Profiler.Get().m_Gateway.CreateProfileSnapshot();
+				auto currentHash = std::hash<json>{}(currentSnapshot);
+				if (previousHash && currentHash == *previousHash)
+					return {};
+
+				previousHash = currentHash;
+				return currentSnapshot;
+			}
+
+		private:
+			/// Proxied profiler
+			Profiler& m_Profiler;
+
+			/// helper state variable
+			std::optional<size_t> previousHash;
+		};
+
+		/// Create Snapshot proxy for this profiler
+		/// @returns snapshot proxy for this profiler
+		SnapshotProxy GetSnapshotProxy() { return SnapshotProxy(*this); }
+
 	protected:
 		std::optional<Gateway> m_Gateway;																				///< The "virtual gateway object".