path: root/src/zenserver-test/objectstore-tests.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#if ZEN_WITH_TESTS
#	include "zenserver-test.h"
#	include <zencore/memoryview.h>
#	include <zencore/testing.h>
#	include <zencore/testutils.h>
#	include <zenhttp/httpclient.h>
#	include <zenutil/cloud/s3client.h>
#	include <zenutil/zenserverprocess.h>

ZEN_THIRD_PARTY_INCLUDES_START
#	include <tsl/robin_set.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::tests {

using namespace std::literals;

TEST_SUITE_BEGIN("server.objectstore");

TEST_CASE("objectstore")
{
	ZenServerInstance Instance(TestEnv);

	const uint16_t Port = Instance.SpawnServerAndWaitUntilReady("--objectstore-enabled");
	CHECK(Port != 0);

	// --- objectstore.blobs ---
	{
		INFO("objectstore.blobs");

		std::string_view Bucket = "bkt"sv;

		std::vector<IoHash>	  CompressedBlobsHashes;
		std::vector<uint64_t> BlobsSizes;
		std::vector<uint64_t> CompressedBlobsSizes;

		HttpClient Client(Instance.GetBaseUri() + "/obj/");

		for (size_t I = 0; I < 5; I++)
		{
			IoBuffer Blob = CreateSemiRandomBlob(4711 + I * 7);
			BlobsSizes.push_back(Blob.GetSize());
			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
			CompressedBlobsHashes.push_back(CompressedBlob.DecodeRawHash());
			CompressedBlobsSizes.push_back(CompressedBlob.GetCompressedSize());
			IoBuffer Payload = std::move(CompressedBlob).GetCompressed().Flatten().AsIoBuffer();
			Payload.SetContentType(ZenContentType::kCompressedBinary);

			std::string ObjectPath = fmt::format("{}/{}.utoc",
												 CompressedBlobsHashes.back().ToHexString().substr(0, 2),
												 CompressedBlobsHashes.back().ToHexString());

			HttpClient::Response Result = Client.Put(fmt::format("bucket/{}/{}.utoc", Bucket, ObjectPath), Payload);
			CHECK(Result);
		}

		for (size_t I = 0; I < 5; I++)
		{
			std::string ObjectPath =
				fmt::format("{}/{}.utoc", CompressedBlobsHashes[I].ToHexString().substr(0, 2), CompressedBlobsHashes[I].ToHexString());
			HttpClient::Response Result = Client.Get(fmt::format("bucket/{}/{}.utoc", Bucket, ObjectPath));
			CHECK(Result);
			CHECK_EQ(Result.ResponsePayload.GetSize(), CompressedBlobsSizes[I]);
			IoHash			 RawHash;
			uint64_t		 RawSize;
			CompressedBuffer Compressed =
				CompressedBuffer::FromCompressed(SharedBuffer(std::move(Result.ResponsePayload)), RawHash, RawSize);
			CHECK(Compressed);
			CHECK_EQ(RawHash, CompressedBlobsHashes[I]);
			CHECK_EQ(RawSize, BlobsSizes[I]);
		}
	}

	// --- objectstore.s3client ---
	{
		INFO("objectstore.s3client");

		S3ClientOptions Opts;
		Opts.BucketName					 = "s3test";
		Opts.Region						 = "us-east-1";
		Opts.Endpoint					 = fmt::format("http://localhost:{}/obj/bucket", Port);
		Opts.PathStyle					 = true;
		Opts.Credentials.AccessKeyId	 = "testkey";
		Opts.Credentials.SecretAccessKey = "testsecret";

		S3Client Client(Opts);

		// -- PUT + GET roundtrip --
		std::string_view TestData = "hello from s3client via objectstore"sv;
		IoBuffer		 Content  = IoBufferBuilder::MakeFromMemory(MakeMemoryView(TestData));
		S3Result		 PutRes	  = Client.PutObject("test/hello.txt", std::move(Content));
		REQUIRE_MESSAGE(PutRes.IsSuccess(), PutRes.Error);

		S3GetObjectResult GetRes = Client.GetObject("test/hello.txt");
		REQUIRE_MESSAGE(GetRes.IsSuccess(), GetRes.Error);
		CHECK(GetRes.AsText() == TestData);

		// -- PUT overwrites --
		IoBuffer Original  = IoBufferBuilder::MakeFromMemory(MakeMemoryView("original"sv));
		IoBuffer Overwrite = IoBufferBuilder::MakeFromMemory(MakeMemoryView("overwritten"sv));
		REQUIRE(Client.PutObject("overwrite/file.txt", std::move(Original)).IsSuccess());
		REQUIRE(Client.PutObject("overwrite/file.txt", std::move(Overwrite)).IsSuccess());

		S3GetObjectResult OverwriteGet = Client.GetObject("overwrite/file.txt");
		REQUIRE(OverwriteGet.IsSuccess());
		CHECK(OverwriteGet.AsText() == "overwritten"sv);

		// -- GET not found --
		S3GetObjectResult NotFoundGet = Client.GetObject("nonexistent/file.dat");
		CHECK_FALSE(NotFoundGet.IsSuccess());

		// -- HEAD found --
		std::string_view HeadData	 = "head test data"sv;
		IoBuffer		 HeadContent = IoBufferBuilder::MakeFromMemory(MakeMemoryView(HeadData));
		REQUIRE(Client.PutObject("head/meta.txt", std::move(HeadContent)).IsSuccess());

		S3HeadObjectResult HeadRes = Client.HeadObject("head/meta.txt");
		REQUIRE_MESSAGE(HeadRes.IsSuccess(), HeadRes.Error);
		CHECK(HeadRes.Status == HeadObjectResult::Found);
		CHECK(HeadRes.Info.Size == HeadData.size());

		// -- HEAD not found --
		S3HeadObjectResult HeadNotFound = Client.HeadObject("nonexistent/file.dat");
		CHECK(HeadNotFound.IsSuccess());
		CHECK(HeadNotFound.Status == HeadObjectResult::NotFound);

		// -- LIST objects --
		for (int i = 0; i < 3; ++i)
		{
			std::string Key		= fmt::format("listing/item-{}.txt", i);
			std::string Payload = fmt::format("content-{}", i);
			IoBuffer	Buf		= IoBufferBuilder::MakeFromMemory(MakeMemoryView(Payload));
			REQUIRE(Client.PutObject(Key, std::move(Buf)).IsSuccess());
		}

		S3ListObjectsResult ListRes = Client.ListObjects("listing/");
		REQUIRE_MESSAGE(ListRes.IsSuccess(), ListRes.Error);
		REQUIRE(ListRes.Objects.size() == 3);

		std::vector<std::string> Keys;
		for (const S3ObjectInfo& Obj : ListRes.Objects)
		{
			Keys.push_back(Obj.Key);
			CHECK(Obj.Size > 0);
		}
		std::sort(Keys.begin(), Keys.end());
		CHECK(Keys[0] == "listing/item-0.txt");
		CHECK(Keys[1] == "listing/item-1.txt");
		CHECK(Keys[2] == "listing/item-2.txt");

		// -- LIST empty prefix --
		S3ListObjectsResult EmptyList = Client.ListObjects("no-such-prefix/");
		REQUIRE(EmptyList.IsSuccess());
		CHECK(EmptyList.Objects.empty());
	}

	// --- objectstore.range-requests ---
	{
		INFO("objectstore.range-requests");

		HttpClient Client(Instance.GetBaseUri() + "/obj/");

		IoBuffer	Blob	   = CreateRandomBlob(1024);
		MemoryView	BlobView   = Blob.GetView();
		std::string ObjectPath = "bucket/bkt/range-test/data.bin";

		HttpClient::Response PutResult = Client.Put(ObjectPath, IoBuffer(Blob));
		REQUIRE(PutResult);

		// Full GET without Range header
		{
			HttpClient::Response Result = Client.Get(ObjectPath);
			CHECK(Result.StatusCode == HttpResponseCode::OK);
			CHECK_EQ(Result.ResponsePayload.GetSize(), 1024u);
			CHECK(Result.ResponsePayload.GetView().EqualBytes(BlobView));
		}

		// Single range: bytes 100-199
		{
			HttpClient::Response Result = Client.Get(ObjectPath, {{"Range", "bytes=100-199"}});
			CHECK(Result.StatusCode == HttpResponseCode::PartialContent);
			CHECK_EQ(Result.ResponsePayload.GetSize(), 100u);
			CHECK(Result.ResponsePayload.GetView().EqualBytes(BlobView.Mid(100, 100)));
		}

		// Range starting at zero: bytes 0-49
		{
			HttpClient::Response Result = Client.Get(ObjectPath, {{"Range", "bytes=0-49"}});
			CHECK(Result.StatusCode == HttpResponseCode::PartialContent);
			CHECK_EQ(Result.ResponsePayload.GetSize(), 50u);
			CHECK(Result.ResponsePayload.GetView().EqualBytes(BlobView.Mid(0, 50)));
		}

		// Range at end of file: bytes 1000-1023
		{
			HttpClient::Response Result = Client.Get(ObjectPath, {{"Range", "bytes=1000-1023"}});
			CHECK(Result.StatusCode == HttpResponseCode::PartialContent);
			CHECK_EQ(Result.ResponsePayload.GetSize(), 24u);
			CHECK(Result.ResponsePayload.GetView().EqualBytes(BlobView.Mid(1000, 24)));
		}

		// Multiple ranges: bytes 0-49 and 100-149
		{
			HttpClient::Response Result = Client.Get(ObjectPath, {{"Range", "bytes=0-49,100-149"}});
			CHECK(Result.StatusCode == HttpResponseCode::PartialContent);

			std::string_view Body(reinterpret_cast<const char*>(Result.ResponsePayload.GetData()), Result.ResponsePayload.GetSize());

			// Verify multipart structure contains both range payloads
			CHECK(Body.find("Content-Range: bytes 0-49/1024") != std::string_view::npos);
			CHECK(Body.find("Content-Range: bytes 100-149/1024") != std::string_view::npos);

			// Extract and verify actual data for first range
			auto FindPartData = [&](std::string_view ContentRange) -> std::string_view {
				size_t Pos = Body.find(ContentRange);
				if (Pos == std::string_view::npos)
				{
					return {};
				}
				// Skip past the Content-Range line and the blank line separator
				Pos = Body.find("\r\n\r\n", Pos);
				if (Pos == std::string_view::npos)
				{
					return {};
				}
				Pos += 4;
				size_t End = Body.find("\r\n--", Pos);
				if (End == std::string_view::npos)
				{
					return {};
				}
				return Body.substr(Pos, End - Pos);
			};

			std::string_view Part1 = FindPartData("Content-Range: bytes 0-49/1024");
			CHECK_EQ(Part1.size(), 50u);
			CHECK(MemoryView(Part1.data(), Part1.size()).EqualBytes(BlobView.Mid(0, 50)));

			std::string_view Part2 = FindPartData("Content-Range: bytes 100-149/1024");
			CHECK_EQ(Part2.size(), 50u);
			CHECK(MemoryView(Part2.data(), Part2.size()).EqualBytes(BlobView.Mid(100, 50)));
		}

		// Out-of-bounds single range
		{
			HttpClient::Response Result = Client.Get(ObjectPath, {{"Range", "bytes=2000-2099"}});
			CHECK(Result.StatusCode == HttpResponseCode::RangeNotSatisfiable);
		}

		// Out-of-bounds multi-range
		{
			HttpClient::Response Result = Client.Get(ObjectPath, {{"Range", "bytes=0-49,2000-2099"}});
			CHECK(Result.StatusCode == HttpResponseCode::RangeNotSatisfiable);
		}
	}
}

TEST_CASE("objectstore.range-requests-download")
{
	ZenServerInstance Instance(TestEnv);
	const uint16_t	  Port = Instance.SpawnServerAndWaitUntilReady("--objectstore-enabled");
	REQUIRE(Port != 0);

	HttpClient Client(Instance.GetBaseUri() + "/obj/");

	IoBuffer	Blob	   = CreateRandomBlob(1024);
	MemoryView	BlobView   = Blob.GetView();
	std::string ObjectPath = "bucket/bkt/range-download-test/data.bin";

	HttpClient::Response PutResult = Client.Put(ObjectPath, IoBuffer(Blob));
	REQUIRE(PutResult);

	ScopedTemporaryDirectory DownloadDir;

	// Single range via Download: verify Ranges is populated and GetRanges maps correctly
	{
		HttpClient::Response Result = Client.Download(ObjectPath, DownloadDir.Path(), {{"Range", "bytes=100-199"}});
		CHECK(Result.StatusCode == HttpResponseCode::PartialContent);
		REQUIRE_EQ(Result.Ranges.size(), 1u);
		CHECK_EQ(Result.Ranges[0].RangeOffset, 100u);
		CHECK_EQ(Result.Ranges[0].RangeLength, 100u);

		std::vector<std::pair<uint64_t, uint64_t>> RequestedRanges = {{100, 100}};
		std::vector<std::pair<uint64_t, uint64_t>> PayloadRanges   = Result.GetRanges(RequestedRanges);
		REQUIRE_EQ(PayloadRanges.size(), 1u);
		CHECK(Result.ResponsePayload.GetView().Mid(PayloadRanges[0].first, PayloadRanges[0].second).EqualBytes(BlobView.Mid(100, 100)));
	}

	// Multi-range via Download: verify Ranges is populated for both parts and GetRanges maps correctly
	{
		HttpClient::Response Result = Client.Download(ObjectPath, DownloadDir.Path(), {{"Range", "bytes=0-49,100-149"}});
		CHECK(Result.StatusCode == HttpResponseCode::PartialContent);
		REQUIRE_EQ(Result.Ranges.size(), 2u);
		CHECK_EQ(Result.Ranges[0].RangeOffset, 0u);
		CHECK_EQ(Result.Ranges[0].RangeLength, 50u);
		CHECK_EQ(Result.Ranges[1].RangeOffset, 100u);
		CHECK_EQ(Result.Ranges[1].RangeLength, 50u);

		std::vector<std::pair<uint64_t, uint64_t>> RequestedRanges = {{0, 50}, {100, 50}};
		std::vector<std::pair<uint64_t, uint64_t>> PayloadRanges   = Result.GetRanges(RequestedRanges);
		REQUIRE_EQ(PayloadRanges.size(), 2u);
		CHECK(Result.ResponsePayload.GetView().Mid(PayloadRanges[0].first, PayloadRanges[0].second).EqualBytes(BlobView.Mid(0, 50)));
		CHECK(Result.ResponsePayload.GetView().Mid(PayloadRanges[1].first, PayloadRanges[1].second).EqualBytes(BlobView.Mid(100, 50)));
	}
}

TEST_SUITE_END();

}  // namespace zen::tests
#endif