From 7f138c4cfc6d6d6fc3e9b29f5cbb750d2462ca96 Mon Sep 17 00:00:00 2001
From: Stefan Boberg <stefan@boberg.org>
Date: Thu, 16 Sep 2021 15:41:11 +0200
Subject: Improved package serialization to allow round tripping

---
 zenhttp/httpshared.cpp | 93 +++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 77 insertions(+), 16 deletions(-)

(limited to 'zenhttp/httpshared.cpp')

diff --git a/zenhttp/httpshared.cpp b/zenhttp/httpshared.cpp
index 68252a763..2dbf95959 100644
--- a/zenhttp/httpshared.cpp
+++ b/zenhttp/httpshared.cpp
@@ -1,6 +1,6 @@
 // Copyright Epic Games, Inc. All Rights Reserved.
 
-#include "httpshared.h"
+#include <zenhttp/httpshared.h>
 
 #include <zencore/compactbinarypackage.h>
 #include <zencore/compositebuffer.h>
@@ -58,22 +58,54 @@ FormatPackageMessage(const CbPackage& Data)
 	IoBuffer RootIoBuffer = Data.GetObject().GetBuffer().AsIoBuffer();
 	ResponseBuffers.push_back(RootIoBuffer);  // Root object
 
-	*AttachmentInfo++ = {.AttachmentSize = RootIoBuffer.Size(), .AttachmentHash = Data.GetObjectHash()};
+	*AttachmentInfo++ = {.AttachmentSize = RootIoBuffer.Size(),
+						 .Flags			 = CbAttachmentEntry::kIsObject,
+						 .AttachmentHash = Data.GetObjectHash()};
 
 	// Attachment payloads
 
 	for (const CbAttachment& Attachment : Attachments)
 	{
-		CompressedBuffer AttachmentBuffer = Attachment.AsCompressedBinary();
-		CompositeBuffer	 Compressed		  = AttachmentBuffer.GetCompressed();
+		if (Attachment.IsNull())
+		{
+			ZEN_NOT_IMPLEMENTED("Null attachments are not supported");
+		}
+		else if (CompressedBuffer AttachmentBuffer = Attachment.AsCompressedBinary())
+		{
+			CompositeBuffer Compressed = AttachmentBuffer.GetCompressed();
 
-		*AttachmentInfo++ = {.AttachmentSize = AttachmentBuffer.GetCompressedSize(),
-							 .AttachmentHash = IoHash::FromBLAKE3(AttachmentBuffer.GetRawHash())};
+			*AttachmentInfo++ = {.AttachmentSize = AttachmentBuffer.GetCompressedSize(),
+								 .Flags			 = CbAttachmentEntry::kIsCompressed,
+								 .AttachmentHash = IoHash::FromBLAKE3(AttachmentBuffer.GetRawHash())};
 
-		for (const SharedBuffer& Segment : Compressed.GetSegments())
+			for (const SharedBuffer& Segment : Compressed.GetSegments())
+			{
+				ResponseBuffers.push_back(Segment.AsIoBuffer());
+				TotalAttachmentsSize += Segment.GetSize();
+			}
+		}
+		else if (CbObject AttachmentObject = Attachment.AsObject())
 		{
-			ResponseBuffers.push_back(Segment.AsIoBuffer());
-			TotalAttachmentsSize += Segment.GetSize();
+			IoBuffer ObjIoBuffer = AttachmentObject.GetBuffer().AsIoBuffer();
+			ResponseBuffers.push_back(ObjIoBuffer);
+
+			*AttachmentInfo++ = {.AttachmentSize = ObjIoBuffer.Size(),
+								 .Flags			 = CbAttachmentEntry::kIsObject,
+								 .AttachmentHash = Attachment.GetHash()};
+		}
+		else if (CompositeBuffer AttachmentBinary = Attachment.AsCompositeBinary())
+		{
+			*AttachmentInfo++ = {.AttachmentSize = AttachmentBinary.GetSize(), .Flags = 0, .AttachmentHash = Attachment.GetHash()};
+
+			for (const SharedBuffer& Segment : AttachmentBinary.GetSegments())
+			{
+				ResponseBuffers.push_back(Segment.AsIoBuffer());
+				TotalAttachmentsSize += Segment.GetSize();
+			}
+		}
+		else
+		{
+			ZEN_NOT_IMPLEMENTED("Unknown attachment kind");
 		}
 	}
 
@@ -119,16 +151,45 @@ ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint
 
 		Reader.Read(AttachmentBuffer.MutableData(), AttachmentSize);
 
-		CompressedBuffer CompBuf(CompressedBuffer::FromCompressed(SharedBuffer(AttachmentBuffer)));
-
-		if (i == 0)
+		if (Entry.Flags & CbAttachmentEntry::kIsCompressed)
 		{
-			Package.SetObject(LoadCompactBinaryObject(std::move(CompBuf)));
+			CompressedBuffer CompBuf(CompressedBuffer::FromCompressed(SharedBuffer(AttachmentBuffer)));
+
+			if (Entry.Flags & CbAttachmentEntry::kIsObject)
+			{
+				if (i == 0)
+				{
+					Package.SetObject(LoadCompactBinaryObject(std::move(CompBuf)));
+				}
+				else
+				{
+					ZEN_NOT_IMPLEMENTED("Object attachments are not currently supported");
+				}
+			}
+			else
+			{
+				CbAttachment Attachment(std::move(CompBuf));
+				Package.AddAttachment(Attachment);
+			}
 		}
-		else
+		else /* not compressed */
 		{
-			CbAttachment Attachment(std::move(CompBuf));
-			Package.AddAttachment(Attachment);
+			if (Entry.Flags & CbAttachmentEntry::kIsObject)
+			{
+				if (i == 0)
+				{
+					Package.SetObject(LoadCompactBinaryObject(AttachmentBuffer));
+				}
+				else
+				{
+					ZEN_NOT_IMPLEMENTED("Object attachments are not currently supported");
+				}
+			}
+			else
+			{
+				CbAttachment Attachment(SharedBuffer{AttachmentBuffer});
+				Package.AddAttachment(Attachment);
+			}
 		}
 	}
 
-- 
cgit v1.2.3