Fixed bug when cloning CbObject.

4 files changed, 198 insertions, 29 deletions

diff --git a/zenserver-test/zenserver-test.cpp b/zenserver-test/zenserver-test.cpp
index a5cfabe08..e3f859d31 100644
--- a/zenserver-test/zenserver-test.cpp
+++ b/zenserver-test/zenserver-test.cpp
@@ -1111,6 +1111,45 @@ TEST_CASE("project.pipe")
 }
 #	endif
 
+namespace utils {
+
+	struct ZenConfig
+	{
+		std::filesystem::path DataDir;
+		uint16_t			  Port;
+		std::string			  BaseUri;
+		std::string			  Args;
+
+		static ZenConfig New(uint16_t Port = 13337, std::string Args = "")
+		{
+			return ZenConfig{.DataDir = TestEnv.CreateNewTestDir(),
+							 .Port	  = Port,
+							 .BaseUri = "http://localhost:{}/z$"_format(Port),
+							 .Args	  = std::move(Args)};
+		}
+
+		static ZenConfig NewWithUpstream(uint16_t UpstreamPort)
+		{
+			return New(13337, "--debug --upstream-thread-count=0 --upstream-zen-url=http://localhost:{}"_format(UpstreamPort));
+		}
+
+		void Spawn(ZenServerInstance& Inst)
+		{
+			Inst.SetTestDir(DataDir);
+			Inst.SpawnServer(Port, Args);
+			Inst.WaitUntilReady();
+		}
+	};
+
+	void SpawnServer(ZenServerInstance& Server, ZenConfig& Cfg)
+	{
+		Server.SetTestDir(Cfg.DataDir);
+		Server.SpawnServer(Cfg.Port, Cfg.Args);
+		Server.WaitUntilReady();
+	}
+
+}  // namespace utils
+
 TEST_CASE("zcache.basic")
 {
 	using namespace std::literals;
@@ -1417,34 +1456,7 @@ TEST_CASE("zcache.cbpackage")
 TEST_CASE("zcache.policy")
 {
 	using namespace std::literals;
-
-	struct ZenConfig
-	{
-		std::filesystem::path DataDir;
-		uint16_t			  Port;
-		std::string			  BaseUri;
-		std::string			  Args;
-
-		static ZenConfig New(uint16_t Port = 13337, std::string Args = "")
-		{
-			return ZenConfig{.DataDir = TestEnv.CreateNewTestDir(),
-							 .Port	  = Port,
-							 .BaseUri = "http://localhost:{}/z$"_format(Port),
-							 .Args	  = std::move(Args)};
-		}
-
-		static ZenConfig NewWithUpstream(uint16_t UpstreamPort)
-		{
-			return New(13337, "--upstream-thread-count=0 --upstream-zen-url=http://localhost:{}"_format(UpstreamPort));
-		}
-
-		void Spawn(ZenServerInstance& Inst)
-		{
-			Inst.SetTestDir(DataDir);
-			Inst.SpawnServer(Port, Args);
-			Inst.WaitUntilReady();
-		}
-	};
+	using namespace utils;
 
 	auto GenerateData = [](uint64_t Size, zen::IoHash& OutHash) -> zen::UniqueBuffer {
 		auto	 Buf  = zen::UniqueBuffer::Alloc(Size);
@@ -1921,6 +1933,102 @@ TEST_CASE("zcache.batch")
 		return zen::IoBuffer(zen::IoBuffer::Clone, MemStream.Data(), MemStream.Size());
 	};
 
+	auto PutCacheRecords = [&CreateCacheRecord, &ToIoBuffer](std::string_view BaseUri,
+															 std::string_view Query,
+															 std::string_view Bucket,
+															 size_t			  Num,
+															 size_t			  PayloadSize = 1024) -> std::vector<CacheKey> {
+		std::vector<zen::CacheKey> OutKeys;
+
+		for (uint32_t Key = 1; Key <= Num; ++Key)
+		{
+			const zen::CacheKey CacheKey	= zen::CacheKey::Create(Bucket, zen::IoHash::HashBuffer(&Key, sizeof uint32_t));
+			CbPackage			CacheRecord = CreateCacheRecord(CacheKey, PayloadSize);
+
+			OutKeys.push_back(CacheKey);
+
+			IoBuffer Payload = ToIoBuffer(CacheRecord);
+
+			cpr::Response Result = cpr::Put(cpr::Url{"{}/{}/{}{}"_format(BaseUri, CacheKey.Bucket, CacheKey.Hash, Query)},
+											cpr::Body{(const char*)Payload.Data(), Payload.Size()},
+											cpr::Header{{"Content-Type", "application/x-ue-cbpkg"}});
+
+			CHECK(Result.status_code == 201);
+		}
+
+		return OutKeys;
+	};
+
+	struct GetCacheRecordResult
+	{
+		zen::CbPackage				  Response;
+		std::vector<zen::CbFieldView> Records;
+		bool						  Success;
+	};
+
+	auto GetCacheRecords =
+		[](std::string_view BaseUri, std::span<zen::CacheKey> Keys, const zen::CacheRecordPolicy& Policy) -> GetCacheRecordResult {
+		using namespace zen;
+
+		CbObjectWriter Request;
+		Request << "Method"sv
+				<< "GetCacheRecords"sv;
+		Request.BeginObject("Params"sv);
+
+		Request.BeginArray("CacheKeys"sv);
+		for (const CacheKey& Key : Keys)
+		{
+			Request.BeginObject();
+			Request << "Bucket"sv << Key.Bucket << "Hash"sv << Key.Hash;
+			Request.EndObject();
+		}
+		Request.EndArray();
+
+		Request.BeginObject("Policy");
+		CacheRecordPolicy::Save(Policy, Request);
+		Request.EndObject();
+
+		Request.EndObject();
+
+		BinaryWriter Body;
+		Request.Save(Body);
+
+		cpr::Response Result = cpr::Post(cpr::Url{"{}/$rpc"_format(BaseUri)},
+										 cpr::Header{{"Content-Type", "application/x-ue-cb"}, {"Accept", "application/x-ue-cbpkg"}},
+										 cpr::Body{(const char*)Body.GetData(), Body.GetSize()});
+
+		GetCacheRecordResult OutResult;
+
+		if (Result.status_code == 200)
+		{
+			CbPackage Response;
+			if (Response.TryLoad(zen::IoBuffer(zen::IoBuffer::Wrap, Result.text.data(), Result.text.size())))
+			{
+				OutResult.Response			= std::move(Response);
+				CbObjectView ResponseObject = OutResult.Response.GetObject();
+
+				for (CbFieldView RecordView : ResponseObject["Result"])
+				{
+					ExtendableStringBuilder<256> Tmp;
+					auto JSON = RecordView.AsObjectView().ToJson(Tmp).ToView();
+					OutResult.Records.push_back(RecordView);
+				}
+
+				OutResult.Success = true;
+			}
+		}
+
+		return OutResult;
+	};
+
+	auto LoadKey = [](zen::CbFieldView KeyView) -> zen::CacheKey {
+		if (zen::CbObjectView KeyObj = KeyView.AsObjectView())
+		{
+			return CacheKey::Create(KeyObj["Bucket"sv].AsString(), KeyObj["Hash"].AsHash());
+		}
+		return CacheKey::Empty;
+	};
+
 	SUBCASE("get cache records")
 	{
 		std::filesystem::path TestDir	 = TestEnv.CreateNewTestDir();
@@ -2114,6 +2222,59 @@ TEST_CASE("zcache.batch")
 			CHECK(ReturnCount == (CacheKeys.size() * 2));
 		}
 	}
+
+	SUBCASE("policy - 'QueryLocal' does not query upstream")
+	{
+		using namespace utils;
+
+		ZenConfig		  UpstreamCfg = ZenConfig::New(13338);
+		ZenServerInstance UpstreamServer(TestEnv);
+		ZenConfig		  LocalCfg = ZenConfig::NewWithUpstream(13338);
+		ZenServerInstance LocalServer(TestEnv);
+
+		SpawnServer(UpstreamServer, UpstreamCfg);
+		SpawnServer(LocalServer, LocalCfg);
+
+		std::vector<zen::CacheKey> Keys = PutCacheRecords(UpstreamCfg.BaseUri, ""sv, "mastodon"sv, 4);
+
+		CacheRecordPolicy	 Policy(CachePolicy::QueryLocal);
+		GetCacheRecordResult Result = GetCacheRecords(LocalCfg.BaseUri, Keys, Policy);
+
+		CHECK(Result.Records.size() == Keys.size());
+
+		for (CbFieldView RecordView : Result.Records)
+		{
+			CHECK(RecordView.IsNull());
+		}
+	}
+
+	SUBCASE("policy - 'QueryRemote' does query upstream")
+	{
+		using namespace utils;
+
+		ZenConfig		  UpstreamCfg = ZenConfig::New(13338);
+		ZenServerInstance UpstreamServer(TestEnv);
+		ZenConfig		  LocalCfg = ZenConfig::NewWithUpstream(13338);
+		ZenServerInstance LocalServer(TestEnv);
+
+		SpawnServer(UpstreamServer, UpstreamCfg);
+		SpawnServer(LocalServer, LocalCfg);
+
+		std::vector<zen::CacheKey> Keys = PutCacheRecords(UpstreamCfg.BaseUri, ""sv, "mastodon"sv, 4);
+
+		CacheRecordPolicy	 Policy(CachePolicy::QueryLocal | CachePolicy::QueryRemote);
+		GetCacheRecordResult Result = GetCacheRecords(LocalCfg.BaseUri, Keys, Policy);
+
+		CHECK(Result.Records.size() == Keys.size());
+
+		for (size_t Index = 0; CbFieldView RecordView : Result.Response.GetObject()["Result"sv])
+		{
+			CbObjectView RecordObj = RecordView.AsObjectView();
+			const zen::CacheKey& ExpectedKey = Keys[Index++];
+			zen::CacheKey Key = LoadKey(RecordObj["CacheKey"sv]);
+			CHECK(Key == ExpectedKey);
+		}
+	}
 }
 
 struct RemoteExecutionRequest
diff --git a/zenserver/cache/structuredcache.cpp b/zenserver/cache/structuredcache.cpp
index 1b6406562..d3da98620 100644
--- a/zenserver/cache/structuredcache.cpp
+++ b/zenserver/cache/structuredcache.cpp
@@ -919,7 +919,8 @@ HttpStructuredCacheService::HandleRpcGetCacheRecords(zen::HttpServerRequest& Req
 							  ToString(HttpContentType::kCbObject));
 
 					ZEN_ASSERT(Params.KeyIndex < CacheValues.size());
-					CacheValues[Params.KeyIndex] = IoBufferBuilder::MakeCloneFromMemory(Params.Record.GetView());
+
+					CacheValues[Params.KeyIndex] = CbObject::Clone(Params.Record).GetBuffer().AsIoBuffer();
 					m_CacheStats.HitCount++;
 					m_CacheStats.UpstreamHitCount++;
 				}
diff --git a/zenutil/cache/cachepolicy.cpp b/zenutil/cache/cachepolicy.cpp
index 968f6aa53..f718bf841 100644
--- a/zenutil/cache/cachepolicy.cpp
+++ b/zenutil/cache/cachepolicy.cpp
@@ -102,6 +102,12 @@ ParseSkipCachePolicy(std::string_view SkipPolicy, CachePolicy Default)
 	return Result;
 }
 
+CacheRecordPolicy::CacheRecordPolicy(const CachePolicy RecordPolicy, const CachePolicy PayloadPolicy)
+: m_RecordPolicy(RecordPolicy)
+, m_DefaultPayloadPolicy(PayloadPolicy)
+{
+}
+
 CachePolicy
 CacheRecordPolicy::GetPayloadPolicy(const Oid& PayloadId) const
 {
diff --git a/zenutil/include/zenutil/cache/cachepolicy.h b/zenutil/include/zenutil/cache/cachepolicy.h
index da0db435c..0b5a70ea8 100644
--- a/zenutil/include/zenutil/cache/cachepolicy.h
+++ b/zenutil/include/zenutil/cache/cachepolicy.h
@@ -45,6 +45,7 @@ class CacheRecordPolicy
 {
 public:
 	CacheRecordPolicy() = default;
+	CacheRecordPolicy(const CachePolicy RecordPolicy, const CachePolicy DefaultPayloadPolicy = CachePolicy::Default);
 
 	CachePolicy GetRecordPolicy() const { return m_RecordPolicy; }
 	CachePolicy GetPayloadPolicy(const Oid& PayloadId) const;