Cache process handles for FormatPackageMessage (#360)

1 files changed, 141 insertions, 127 deletions

diff --git a/src/zenhttp/httpshared.cpp b/src/zenhttp/httpshared.cpp
index 7aade56d2..4a67b69e7 100644
--- a/src/zenhttp/httpshared.cpp
+++ b/src/zenhttp/httpshared.cpp
@@ -14,6 +14,7 @@
 #include <zencore/stream.h>
 #include <zencore/testing.h>
 #include <zencore/testutils.h>
+#include <zencore/trace.h>
 
 #include <span>
 #include <vector>
@@ -27,20 +28,20 @@ namespace zen {
 const std::string_view HandlePrefix(":?#:");
 
 std::vector<IoBuffer>
-FormatPackageMessage(const CbPackage& Data, int TargetProcessPid)
+FormatPackageMessage(const CbPackage& Data, void* TargetProcessHandle)
 {
-	return FormatPackageMessage(Data, FormatFlags::kDefault, TargetProcessPid);
+	return FormatPackageMessage(Data, FormatFlags::kDefault, TargetProcessHandle);
 }
 CompositeBuffer
-FormatPackageMessageBuffer(const CbPackage& Data, int TargetProcessPid)
+FormatPackageMessageBuffer(const CbPackage& Data, void* TargetProcessHandle)
 {
-	return FormatPackageMessageBuffer(Data, FormatFlags::kDefault, TargetProcessPid);
+	return FormatPackageMessageBuffer(Data, FormatFlags::kDefault, TargetProcessHandle);
 }
 
 CompositeBuffer
-FormatPackageMessageBuffer(const CbPackage& Data, FormatFlags Flags, int TargetProcessPid)
+FormatPackageMessageBuffer(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle)
 {
-	std::vector<IoBuffer> Message = FormatPackageMessage(Data, Flags, TargetProcessPid);
+	std::vector<IoBuffer> Message = FormatPackageMessage(Data, Flags, TargetProcessHandle);
 
 	std::vector<SharedBuffer> Buffers;
 
@@ -52,43 +53,125 @@ FormatPackageMessageBuffer(const CbPackage& Data, FormatFlags Flags, int TargetP
 	return CompositeBuffer(std::move(Buffers));
 }
 
-std::vector<IoBuffer>
-FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, int TargetProcessPid)
+static void
+MarshalLocal(CbAttachmentEntry*&		  AttachmentInfo,
+			 const std::string&			  Path8,
+			 CbAttachmentReferenceHeader& LocalRef,
+			 const IoHash&				  AttachmentHash,
+			 bool						  IsCompressed,
+			 std::vector<IoBuffer>&		  ResponseBuffers)
 {
-	void* TargetProcessHandle = nullptr;
-#if ZEN_PLATFORM_WINDOWS
-	std::vector<HANDLE> DuplicatedHandles;
-	auto				_ = MakeGuard([&DuplicatedHandles, &TargetProcessHandle]() {
-		   if (TargetProcessHandle == nullptr)
-		   {
-			   return;
-		   }
-
-		   for (HANDLE DuplicatedHandle : DuplicatedHandles)
-		   {
-			   HANDLE ClosingHandle;
-			   if (::DuplicateHandle((HANDLE)TargetProcessHandle,
-									 DuplicatedHandle,
-									 GetCurrentProcess(),
-									 &ClosingHandle,
-									 0,
-									 FALSE,
-									 DUPLICATE_CLOSE_SOURCE | DUPLICATE_SAME_ACCESS) == TRUE)
-			   {
-				   ::CloseHandle(ClosingHandle);
-			   }
-		   }
-		   ::CloseHandle((HANDLE)TargetProcessHandle);
-		   TargetProcessHandle = nullptr;
-	   });
-
-	if (EnumHasAllFlags(Flags, FormatFlags::kAllowLocalReferences) && TargetProcessPid != 0)
+	IoBuffer RefBuffer(sizeof(CbAttachmentReferenceHeader) + Path8.size());
+
+	CbAttachmentReferenceHeader* RefHdr = RefBuffer.MutableData<CbAttachmentReferenceHeader>();
+	*RefHdr++							= LocalRef;
+	memcpy(RefHdr, Path8.data(), Path8.size());
+
+	*AttachmentInfo++ = {.PayloadSize = RefBuffer.GetSize(),
+						 .Flags		  = (IsCompressed ? uint32_t(CbAttachmentEntry::kIsCompressed) : 0u) | CbAttachmentEntry::kIsLocalRef,
+						 .AttachmentHash = AttachmentHash};
+
+	ResponseBuffers.push_back(std::move(RefBuffer));
+};
+
+static bool
+IsLocalRef(tsl::robin_map<void*, std::string>& FileNameMap,
+		   std::vector<void*>&				   DuplicatedHandles,
+		   const CompositeBuffer&			   AttachmentBinary,
+		   bool								   DenyPartialLocalReferences,
+		   void*							   TargetProcessHandle,
+		   CbAttachmentReferenceHeader&		   LocalRef,
+		   std::string&						   Path8)
+{
+	const SharedBuffer&	  Segment = AttachmentBinary.GetSegments().front();
+	IoBufferFileReference Ref;
+	const IoBuffer&		  SegmentBuffer = Segment.AsIoBuffer();
+
+	if (!SegmentBuffer.GetFileReference(Ref))
+	{
+		return false;
+	}
+
+	if (DenyPartialLocalReferences && !SegmentBuffer.IsWholeFile())
+	{
+		return false;
+	}
+
+	if (auto It = FileNameMap.find(Ref.FileHandle); It != FileNameMap.end())
 	{
-		TargetProcessHandle = OpenProcess(PROCESS_DUP_HANDLE, FALSE, TargetProcessPid);
+		Path8 = It->second;
 	}
-#else
-	ZEN_UNUSED(TargetProcessPid);
-	void* DuplicatedHandles = nullptr;
+	else
+	{
+		bool UseFilePath = true;
+#if ZEN_PLATFORM_WINDOWS
+		if (TargetProcessHandle != nullptr)
+		{
+			HANDLE TargetHandle = INVALID_HANDLE_VALUE;
+			BOOL   OK			= ::DuplicateHandle(GetCurrentProcess(),
+										Ref.FileHandle,
+										(HANDLE)TargetProcessHandle,
+										&TargetHandle,
+										FILE_GENERIC_READ,
+										FALSE,
+										0);
+			if (OK)
+			{
+				DuplicatedHandles.push_back((void*)TargetHandle);
+				Path8		= fmt::format("{}{}", HandlePrefix, reinterpret_cast<uint64_t>(TargetHandle));
+				UseFilePath = false;
+			}
+		}
+#else	// ZEN_PLATFORM_WINDOWS
+		ZEN_UNUSED(TargetProcessHandle);
+		ZEN_UNUSED(DuplicatedHandles);
+		// Not supported on Linux/Mac. Could potentially use pidfd_getfd() but that requires a fairly new Linux kernel/includes and to
+		// deal with acceess rights etc.
+#endif	// ZEN_PLATFORM_WINDOWS
+		if (UseFilePath)
+		{
+			ExtendablePathBuilder<256> LocalRefFile;
+			LocalRefFile.Append(std::filesystem::absolute(PathFromHandle(Ref.FileHandle)));
+			Path8 = LocalRefFile.ToUtf8();
+		}
+		FileNameMap.insert_or_assign(Ref.FileHandle, Path8);
+	}
+
+	LocalRef.AbsolutePathLength = gsl::narrow<uint16_t>(Path8.size());
+	LocalRef.PayloadByteOffset	= Ref.FileChunkOffset;
+	LocalRef.PayloadByteSize	= Ref.FileChunkSize;
+
+	return true;
+};
+
+std::vector<IoBuffer>
+FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, void* TargetProcessHandle)
+{
+	ZEN_TRACE_CPU("FormatPackageMessage");
+
+	std::vector<void*> DuplicatedHandles;
+#if ZEN_PLATFORM_WINDOWS
+	auto _ = MakeGuard([&DuplicatedHandles, &TargetProcessHandle]() {
+		if (TargetProcessHandle == nullptr)
+		{
+			return;
+		}
+
+		for (void* DuplicatedHandle : DuplicatedHandles)
+		{
+			HANDLE ClosingHandle;
+			if (::DuplicateHandle((HANDLE)TargetProcessHandle,
+								  (HANDLE)DuplicatedHandle,
+								  GetCurrentProcess(),
+								  &ClosingHandle,
+								  0,
+								  FALSE,
+								  DUPLICATE_CLOSE_SOURCE | DUPLICATE_SAME_ACCESS) == TRUE)
+			{
+				::CloseHandle(ClosingHandle);
+			}
+		}
+	});
 #endif	// ZEN_PLATFORM_WINDOWS
 
 	const std::span<const CbAttachment>& Attachments = Data.GetAttachments();
@@ -118,91 +201,8 @@ FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, int TargetProcess
 	*AttachmentInfo++ = {.PayloadSize = RootIoBuffer.Size(), .Flags = CbAttachmentEntry::kIsObject, .AttachmentHash = Data.GetObjectHash()};
 
 	// Attachment payloads
-
-	auto MarshalLocal = [&AttachmentInfo, &ResponseBuffers](const std::string&			 Path8,
-															CbAttachmentReferenceHeader& LocalRef,
-															const IoHash&				 AttachmentHash,
-															bool						 IsCompressed) {
-		IoBuffer RefBuffer(sizeof(CbAttachmentReferenceHeader) + Path8.size());
-
-		CbAttachmentReferenceHeader* RefHdr = RefBuffer.MutableData<CbAttachmentReferenceHeader>();
-		*RefHdr++							= LocalRef;
-		memcpy(RefHdr, Path8.data(), Path8.size());
-
-		*AttachmentInfo++ = {.PayloadSize = RefBuffer.GetSize(),
-							 .Flags = (IsCompressed ? uint32_t(CbAttachmentEntry::kIsCompressed) : 0u) | CbAttachmentEntry::kIsLocalRef,
-							 .AttachmentHash = AttachmentHash};
-
-		ResponseBuffers.push_back(std::move(RefBuffer));
-	};
-
 	tsl::robin_map<void*, std::string> FileNameMap;
 
-	auto IsLocalRef = [&FileNameMap, &DuplicatedHandles](const CompositeBuffer&		  AttachmentBinary,
-														 bool						  DenyPartialLocalReferences,
-														 void*						  TargetProcessHandle,
-														 CbAttachmentReferenceHeader& LocalRef,
-														 std::string&				  Path8) -> bool {
-		const SharedBuffer&	  Segment = AttachmentBinary.GetSegments().front();
-		IoBufferFileReference Ref;
-		const IoBuffer&		  SegmentBuffer = Segment.AsIoBuffer();
-
-		if (!SegmentBuffer.GetFileReference(Ref))
-		{
-			return false;
-		}
-
-		if (DenyPartialLocalReferences && !SegmentBuffer.IsWholeFile())
-		{
-			return false;
-		}
-
-		if (auto It = FileNameMap.find(Ref.FileHandle); It != FileNameMap.end())
-		{
-			Path8 = It->second;
-		}
-		else
-		{
-			bool UseFilePath = true;
-#if ZEN_PLATFORM_WINDOWS
-			if (TargetProcessHandle != nullptr)
-			{
-				HANDLE TargetHandle = INVALID_HANDLE_VALUE;
-				BOOL   OK			= ::DuplicateHandle(GetCurrentProcess(),
-											Ref.FileHandle,
-											(HANDLE)TargetProcessHandle,
-											&TargetHandle,
-											FILE_GENERIC_READ,
-											FALSE,
-											0);
-				if (OK)
-				{
-					DuplicatedHandles.push_back(TargetHandle);
-					Path8		= fmt::format("{}{}", HandlePrefix, reinterpret_cast<uint64_t>(TargetHandle));
-					UseFilePath = false;
-				}
-			}
-#else	// ZEN_PLATFORM_WINDOWS
-			ZEN_UNUSED(TargetProcessHandle);
-			// Not supported on Linux/Mac. Could potentially use pidfd_getfd() but that requires a fairly new Linux kernel/includes and to
-			// deal with acceess rights etc.
-#endif	// ZEN_PLATFORM_WINDOWS
-			if (UseFilePath)
-			{
-				ExtendablePathBuilder<256> LocalRefFile;
-				LocalRefFile.Append(std::filesystem::absolute(PathFromHandle(Ref.FileHandle)));
-				Path8 = LocalRefFile.ToUtf8();
-			}
-			FileNameMap.insert_or_assign(Ref.FileHandle, Path8);
-		}
-
-		LocalRef.AbsolutePathLength = gsl::narrow<uint16_t>(Path8.size());
-		LocalRef.PayloadByteOffset	= Ref.FileChunkOffset;
-		LocalRef.PayloadByteSize	= Ref.FileChunkSize;
-
-		return true;
-	};
-
 	for (const CbAttachment& Attachment : Attachments)
 	{
 		if (Attachment.IsNull())
@@ -226,7 +226,13 @@ FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, int TargetProcess
 
 			if (MarshalByLocalRef)
 			{
-				MarshalByLocalRef = IsLocalRef(Compressed, DenyPartialLocalReferences, TargetProcessHandle, LocalRef, Path8);
+				MarshalByLocalRef = IsLocalRef(FileNameMap,
+											   DuplicatedHandles,
+											   Compressed,
+											   DenyPartialLocalReferences,
+											   TargetProcessHandle,
+											   LocalRef,
+											   Path8);
 			}
 
 			if (MarshalByLocalRef)
@@ -236,7 +242,7 @@ FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, int TargetProcess
 #if ZEN_PLATFORM_WINDOWS
 				IsHandle = Path8.starts_with(HandlePrefix);
 #endif
-				MarshalLocal(Path8, LocalRef, AttachmentHash, IsCompressed);
+				MarshalLocal(AttachmentInfo, Path8, LocalRef, AttachmentHash, IsCompressed, ResponseBuffers);
 				ZEN_DEBUG("Marshalled '{}' as file {} of {} bytes", Path8, IsHandle ? "handle" : "path", Compressed.GetSize());
 			}
 			else
@@ -272,7 +278,13 @@ FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, int TargetProcess
 
 			if (MarshalByLocalRef)
 			{
-				MarshalByLocalRef = IsLocalRef(AttachmentBinary, DenyPartialLocalReferences, TargetProcessHandle, LocalRef, Path8);
+				MarshalByLocalRef = IsLocalRef(FileNameMap,
+											   DuplicatedHandles,
+											   AttachmentBinary,
+											   DenyPartialLocalReferences,
+											   TargetProcessHandle,
+											   LocalRef,
+											   Path8);
 			}
 
 			if (MarshalByLocalRef)
@@ -282,7 +294,7 @@ FormatPackageMessage(const CbPackage& Data, FormatFlags Flags, int TargetProcess
 #if ZEN_PLATFORM_WINDOWS
 				IsHandle = Path8.starts_with(HandlePrefix);
 #endif
-				MarshalLocal(Path8, LocalRef, AttachmentHash, IsCompressed);
+				MarshalLocal(AttachmentInfo, Path8, LocalRef, AttachmentHash, IsCompressed, ResponseBuffers);
 				ZEN_DEBUG("Marshalled '{}' as file {} of {} bytes", Path8, IsHandle ? "handle" : "path", AttachmentBinary.GetSize());
 			}
 			else
@@ -332,6 +344,8 @@ IsPackageMessage(IoBuffer Payload)
 CbPackage
 ParsePackageMessage(IoBuffer Payload, std::function<IoBuffer(const IoHash&, uint64_t)> CreateBuffer)
 {
+	ZEN_TRACE_CPU("ParsePackageMessage");
+
 	if (!Payload)
 	{
 		return {};