Introduce a SignatureBuilder to build signatures of arbitrary size

tracer/src/Datadog.Tracer.Native/CMakeLists.txt

@@ -177,6 +177,7 @@ add_library("Datadog.Tracer.Native.static" STATIC
         environment_variables_util.cpp
         method_rewriter.cpp
         tracer_tokens.cpp
+        signature_builder.cpp
         debugger_environment_variables_util.cpp
         debugger_probes_tracker.cpp
         iast/aspect.cpp

tracer/src/Datadog.Tracer.Native/Datadog.Tracer.Native.vcxproj

@@ -336,6 +336,7 @@
     <ClInclude Include="rejit_handler.h" />
     <ClInclude Include="rejit_preprocessor.h" />
     <ClInclude Include="rejit_work_offloader.h" />
+    <ClInclude Include="signature_builder.h" />
     <ClInclude Include="stats.h" />
     <ClInclude Include="Synchronized.hpp" />
     <ClInclude Include="threadUtils.h" />
@@ -394,6 +395,7 @@
     <ClCompile Include="rejit_handler.cpp" />
     <ClCompile Include="rejit_preprocessor.cpp" />
     <ClCompile Include="rejit_work_offloader.cpp" />
+    <ClCompile Include="signature_builder.cpp" />
     <ClCompile Include="threadUtils.cpp" />
     <ClCompile Include="tracer_handler_module_method.cpp" />
     <ClCompile Include="tracer_integration_definition.cpp" />

tracer/src/Datadog.Tracer.Native/Datadog.Tracer.Native.vcxproj.filters

@@ -117,6 +117,7 @@
       <Filter>Iast</Filter>
     </ClCompile>
     <ClCompile Include="Generated\generated_definitions.cpp" />
+    <ClCompile Include="signature_builder.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="calltarget_tokens.h" />
@@ -270,6 +271,7 @@
     <ClInclude Include="Generated\generated_definitions.h" />
     <ClInclude Include="Generated\generated_callsites.g.h" />
     <ClInclude Include="Generated\generated_calltargets.g.h" />
+    <ClInclude Include="signature_builder.h" />
   </ItemGroup>
   <ItemGroup>
     <Filter Include="Debugger">

tracer/src/Datadog.Tracer.Native/calltarget_tokens.cpp

@@ -4,6 +4,7 @@
 #include "il_rewriter_wrapper.h"
 #include "logger.h"
 #include "module_metadata.h"
+#include "signature_builder.h"
 
 namespace trace
 {
@@ -434,11 +435,6 @@ HRESULT CallTargetTokens::ModifyLocalSig(ILRewriter* reWriter, TypeSignature* me
         callTargetReturnSizeForNewSignature = 1 + callTargetReturnSize;
     }
 
-    // New signature size
-    ULONG newSignatureSize = originalSignatureSize + returnSignatureTypeSize + (1 + exTypeRefSize) +
-                             callTargetReturnSizeForNewSignature + (1 + callTargetStateTypeRefSize);
-    ULONG newSignatureOffset = 0;
-
     ULONG oldLocalsBuffer;
     ULONG oldLocalsLen = 0;
     unsigned newLocalsBuffer;
@@ -447,69 +443,60 @@ HRESULT CallTargetTokens::ModifyLocalSig(ILRewriter* reWriter, TypeSignature* me
     // Calculate the new locals count
     if (originalSignatureSize == 0)
     {
-        newSignatureSize += 2;
         newLocalsLen = CorSigCompressData(newLocalsCount, &newLocalsBuffer);
     }
     else
     {
         oldLocalsLen = CorSigUncompressData(originalSignature + 1, &oldLocalsBuffer);
         newLocalsCount += oldLocalsBuffer;
         newLocalsLen = CorSigCompressData(newLocalsCount, &newLocalsBuffer);
-        newSignatureSize += newLocalsLen - oldLocalsLen;
     }
 
     // New signature declaration
-    COR_SIGNATURE newSignatureBuffer[BUFFER_SIZE];
-    newSignatureBuffer[newSignatureOffset++] = IMAGE_CEE_CS_CALLCONV_LOCAL_SIG;
+    SignatureBuilder newSignature;
+    newSignature.Append(IMAGE_CEE_CS_CALLCONV_LOCAL_SIG);
 
     // Set the locals count
-    memcpy(&newSignatureBuffer[newSignatureOffset], &newLocalsBuffer, newLocalsLen);
-    newSignatureOffset += newLocalsLen;
+    newSignature.Append(&newLocalsBuffer, newLocalsLen);
 
     // Copy previous locals to the signature
     if (originalSignatureSize > 0)
     {
         const auto copyLength = originalSignatureSize - 1 - oldLocalsLen;
-        memcpy(&newSignatureBuffer[newSignatureOffset], originalSignature + 1 + oldLocalsLen, copyLength);
-        newSignatureOffset += copyLength;
+        newSignature.Append(originalSignature + 1 + oldLocalsLen, copyLength);
     }
 
     // Return value local
     if (returnSignatureType != nullptr)
     {
-        memcpy(&newSignatureBuffer[newSignatureOffset], returnSignatureType, returnSignatureTypeSize);
-        newSignatureOffset += returnSignatureTypeSize;
+        newSignature.Append(returnSignatureType, returnSignatureTypeSize);
     }
 
     // Exception value
-    newSignatureBuffer[newSignatureOffset++] = ELEMENT_TYPE_CLASS;
-    memcpy(&newSignatureBuffer[newSignatureOffset], &exTypeRefBuffer, exTypeRefSize);
-    newSignatureOffset += exTypeRefSize;
+    newSignature.Append(ELEMENT_TYPE_CLASS);
+    newSignature.Append(&exTypeRefBuffer, exTypeRefSize);
 
     // CallTarget Return value
     if (callTargetReturnSignature != nullptr)
     {
-        memcpy(&newSignatureBuffer[newSignatureOffset], callTargetReturnSignature, callTargetReturnSignatureSize);
-        newSignatureOffset += callTargetReturnSignatureSize;
+        newSignature.Append(callTargetReturnSignature, callTargetReturnSignatureSize);
     }
     else
     {
-        newSignatureBuffer[newSignatureOffset++] = ELEMENT_TYPE_VALUETYPE;
-        memcpy(&newSignatureBuffer[newSignatureOffset], &callTargetReturnBuffer, callTargetReturnSize);
-        newSignatureOffset += callTargetReturnSize;
+        newSignature.Append(ELEMENT_TYPE_VALUETYPE);
+        newSignature.Append(&callTargetReturnBuffer, callTargetReturnSize);
     }
 
     // Add custom locals
-    AddAdditionalLocals(methodReturnValue, methodTypeArguments, newSignatureBuffer, newSignatureOffset, newSignatureSize, isAsyncMethod);
+    AddAdditionalLocals(methodReturnValue, methodTypeArguments, newSignature, isAsyncMethod);
 
     // CallTarget state value
-    newSignatureBuffer[newSignatureOffset++] = ELEMENT_TYPE_VALUETYPE;
-    memcpy(&newSignatureBuffer[newSignatureOffset], &callTargetStateTypeRefBuffer, callTargetStateTypeRefSize);
-    newSignatureOffset += callTargetStateTypeRefSize;
+    newSignature.Append(ELEMENT_TYPE_VALUETYPE);
+    newSignature.Append(&callTargetStateTypeRefBuffer, callTargetStateTypeRefSize);
 
     // Get new locals token
     mdToken newLocalVarSig;
-    hr = module_metadata->metadata_emit->GetTokenFromSig(newSignatureBuffer, newSignatureSize, &newLocalVarSig);
+    hr = module_metadata->metadata_emit->GetTokenFromSig(newSignature.GetSignature(), newSignature.Size(), &newLocalVarSig);
     if (FAILED(hr))
     {
         Logger::Warn("Error creating new locals var signature.");
@@ -809,8 +796,7 @@ int CallTargetTokens::GetAdditionalLocalsCount(const std::vector<TypeSignature>&
 }
 
 void CallTargetTokens::AddAdditionalLocals(TypeSignature* methodReturnValue, std::vector<TypeSignature>* methodTypeArguments,
-                                           COR_SIGNATURE (&signatureBuffer)[BUFFER_SIZE], ULONG& signatureOffset,
-                                           ULONG& signatureSize, bool isAsyncMethod)
+                                           SignatureBuilder& signatureBuffer, bool isAsyncMethod)
 {
 }
 

tracer/src/Datadog.Tracer.Native/calltarget_tokens.h

@@ -12,6 +12,7 @@
 #include "clr_helpers.h"
 #include "il_rewriter.h"
 #include "integration.h"
+#include "signature_builder.h"
 #include "../../../shared/src/native-src/string.h" // NOLINT
 #include "../../../shared/src/native-src/com_ptr.h"
 
@@ -77,8 +78,7 @@ class CallTargetTokens
     virtual const shared::WSTRING& GetCallTargetRefStructType() = 0;
 
     virtual void AddAdditionalLocals(TypeSignature* methodReturnValue, std::vector<TypeSignature>* methodTypeArguments,
-                                     COR_SIGNATURE (&signatureBuffer)[BUFFER_SIZE], ULONG& signatureOffset,
-                                     ULONG& signatureSize, bool isAsyncMethod);
+        SignatureBuilder& signatureBuffer, bool isAsyncMethod);
 
     CallTargetTokens(ModuleMetadata* moduleMetadataPtr, bool enableByRefInstrumentation,
                      bool enableCallTargetStateByRef);

tracer/src/Datadog.Tracer.Native/debugger_tokens.cpp

@@ -315,46 +315,33 @@ int DebuggerTokens::GetAdditionalLocalsCount(const std::vector<TypeSignature>& m
 }
 
 void DebuggerTokens::AddAdditionalLocals(TypeSignature* methodReturnValue, std::vector<TypeSignature>* methodTypeArguments,
-                                         COR_SIGNATURE (&signatureBuffer)[BUFFER_SIZE], ULONG& signatureOffset,
-                                         ULONG& signatureSize, bool isAsyncMethod)
+    SignatureBuilder& signature, bool isAsyncMethod)
 {
     // Gets the calltarget state of line probe type buffer and size
     unsigned callTargetStateTypeRefBuffer;
     const auto callTargetStateTypeRefSize = CorSigCompressToken(GetDebuggerState(isAsyncMethod ? AsyncLineProbe : NonAsyncLineProbe), &callTargetStateTypeRefBuffer);
 
-    // Enlarge the *new* signature size
-    signatureSize += (1 + callTargetStateTypeRefSize);
-
     // CallTarget state of line probe
-    signatureBuffer[signatureOffset++] = ELEMENT_TYPE_VALUETYPE;
-    memcpy(&signatureBuffer[signatureOffset], &callTargetStateTypeRefBuffer, callTargetStateTypeRefSize);
-    signatureOffset += callTargetStateTypeRefSize;
+    signature.Append(ELEMENT_TYPE_VALUETYPE);
+    signature.Append(&callTargetStateTypeRefBuffer, callTargetStateTypeRefSize);
 
     // Gets the calltarget state of span method probe type buffer and size
     unsigned spanStateTypeRefBuffer;
     const auto spanStateTypeRefSize = CorSigCompressToken(methodSpanProbeDebuggerStateTypeRef, &spanStateTypeRefBuffer);
 
-    // Enlarge the *new* signature size
-    signatureSize += (1 + spanStateTypeRefSize);
-
     // CallTarget state of async method probe
-    signatureBuffer[signatureOffset++] = ELEMENT_TYPE_VALUETYPE;
-    memcpy(&signatureBuffer[signatureOffset], &spanStateTypeRefBuffer, spanStateTypeRefSize);
-    signatureOffset += spanStateTypeRefSize;
+    signature.Append(ELEMENT_TYPE_VALUETYPE);
+    signature.Append(&spanStateTypeRefBuffer, spanStateTypeRefSize);
 
     // CallTarget states of multi-probe scenario
 
     // Gets the calltarget state of line probe type buffer and size
     unsigned methodDebuggerStatesTypeRefBuffer;
     const auto methodDebuggerStatesTypeRefSize = CorSigCompressToken(GetDebuggerState(NonAsyncMethodMultiProbe), &methodDebuggerStatesTypeRefBuffer);
 
-    // Enlarge the *new* signature size
-    signatureSize += (2 + methodDebuggerStatesTypeRefSize);
-
-    signatureBuffer[signatureOffset++] = ELEMENT_TYPE_SZARRAY;
-    signatureBuffer[signatureOffset++] = ELEMENT_TYPE_VALUETYPE;
-    memcpy(&signatureBuffer[signatureOffset], &methodDebuggerStatesTypeRefBuffer, methodDebuggerStatesTypeRefSize);
-    signatureOffset += methodDebuggerStatesTypeRefSize;
+    signature.Append(ELEMENT_TYPE_SZARRAY);
+    signature.Append(ELEMENT_TYPE_VALUETYPE);
+    signature.Append(&methodDebuggerStatesTypeRefBuffer, methodDebuggerStatesTypeRefSize);
 }
 
 /**

tracer/src/Datadog.Tracer.Native/debugger_tokens.h

@@ -550,8 +550,7 @@ class DebuggerTokens : public CallTargetTokens
     const WSTRING& GetCallTargetRefStructType() override;
 
     void AddAdditionalLocals(TypeSignature* methodReturnValue, std::vector<TypeSignature>* methodTypeArguments,
-                             COR_SIGNATURE (&signatureBuffer)[BUFFER_SIZE], ULONG& signatureOffset,
-                             ULONG& signatureSize, bool isAsyncMethod) override;
+        SignatureBuilder& signatureBuffer, bool isAsyncMethod) override;
 
 public:
     DebuggerTokens(ModuleMetadata* module_metadata_ptr);

tracer/src/Datadog.Tracer.Native/signature_builder.cpp

@@ -0,0 +1,31 @@
+#include "signature_builder.h"
+
+SignatureBuilder::SignatureBuilder()
+{
+    _buffer = _stackSignatureBuffer;
+    _length = STACK_BUFFER_SIZE;
+}
+
+void SignatureBuilder::Append(const COR_SIGNATURE element)
+{
+    EnsureBufferSpace(1);
+    _buffer[_offset++] = element;
+}
+
+void SignatureBuilder::Append(const void* elements, const size_t length)
+{
+    EnsureBufferSpace(length);
+    memcpy(_buffer + _offset, elements, length);
+    _offset += length;
+}
+
+void SignatureBuilder::EnsureBufferSpace(int size)
+{
+    if (_offset + size >= _length)
+    {
+        _heapSignatureBuffer = std::make_unique<COR_SIGNATURE[]>(_length * 2);
+        memcpy(_heapSignatureBuffer.get(), _buffer, _length);
+        _buffer = _heapSignatureBuffer.get();
+        _length *= 2;
+    }
+}

tracer/src/Datadog.Tracer.Native/signature_builder.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <corhlpr.h>
+#include <memory>
+
+class SignatureBuilder
+{
+public:
+    SignatureBuilder();
+
+    void Append(const COR_SIGNATURE elementType);
+    void Append(const void* elements, const size_t length);
+
+    const size_t Size() const { return _offset; }
+
+    const COR_SIGNATURE* GetSignature() const { return _buffer; }
+
+private:
+    constexpr static int STACK_BUFFER_SIZE = 1000;
+
+    COR_SIGNATURE _stackSignatureBuffer[STACK_BUFFER_SIZE];
+    std::unique_ptr<COR_SIGNATURE[]> _heapSignatureBuffer;
+
+    COR_SIGNATURE* _buffer;
+
+    size_t _length;
+    size_t _offset;
+
+    void EnsureBufferSpace(int size);
+};

tracer/src/Datadog.Tracer.Native/tracer_tokens.cpp

@@ -259,24 +259,18 @@ int TracerTokens::GetAdditionalLocalsCount(const std::vector<TypeSignature>& met
 }
 
 void TracerTokens::AddAdditionalLocals(TypeSignature* methodReturnValue, std::vector<TypeSignature>* methodTypeArguments,
-                                       COR_SIGNATURE (&signatureBuffer)[BUFFER_SIZE], ULONG& signatureOffset,
-                                       ULONG& signatureSize, bool isAsyncMethod)
+                                       SignatureBuilder& signature, bool isAsyncMethod)
 {
     // Gets the exception type buffer and size
     unsigned exTypeRefBuffer;
     auto exTypeRefSize = CorSigCompressToken(exTypeRef, &exTypeRefBuffer);
 
     // Exception value for calltarget exception filters
-    signatureBuffer[signatureOffset++] = ELEMENT_TYPE_CLASS;
-    memcpy(&signatureBuffer[signatureOffset], &exTypeRefBuffer, exTypeRefSize);
-    signatureOffset += exTypeRefSize;
-    signatureSize += 1 + exTypeRefSize;
+    signature.Append(ELEMENT_TYPE_CLASS);
+    signature.Append(&exTypeRefBuffer, exTypeRefSize);
 
-
-    signatureBuffer[signatureOffset++] = ELEMENT_TYPE_CLASS;
-    memcpy(&signatureBuffer[signatureOffset], &exTypeRefBuffer, exTypeRefSize);
-    signatureOffset += exTypeRefSize;
-    signatureSize += 1 + exTypeRefSize;
+    signature.Append(ELEMENT_TYPE_CLASS);
+    signature.Append(&exTypeRefBuffer, exTypeRefSize);
 
     if (enable_by_ref_instrumentation)
     {
@@ -291,10 +285,8 @@ void TracerTokens::AddAdditionalLocals(TypeSignature* methodReturnValue, std::ve
             if (SUCCEEDED(IsTypeByRefLike(_profiler_info, *module_metadata, typeArgument, GetCorLibAssemblyRef(), isByRefLike)) &&
                 isByRefLike)
             {
-                signatureBuffer[signatureOffset++] = ELEMENT_TYPE_VALUETYPE;
-                memcpy(&signatureBuffer[signatureOffset], &callTargetRefStructTypeRefBuffer, callTargetRefStructTypeRefSize);
-                signatureOffset += callTargetRefStructTypeRefSize;
-                signatureSize += 1 + callTargetRefStructTypeRefSize;
+                signature.Append(ELEMENT_TYPE_VALUETYPE);
+                signature.Append(&callTargetRefStructTypeRefBuffer, callTargetRefStructTypeRefSize);
             }
         }
     }

tracer/src/Datadog.Tracer.Native/tracer_tokens.h

@@ -34,8 +34,7 @@ class TracerTokens : public CallTargetTokens
 
     HRESULT EnsureBaseCalltargetTokens() override;
     void AddAdditionalLocals(TypeSignature* methodReturnValue, std::vector<TypeSignature>* methodTypeArguments,
-                             COR_SIGNATURE (&signatureBuffer)[BUFFER_SIZE], ULONG& signatureOffset,
-                             ULONG& signatureSize, bool isAsyncMethod) override;
+                             SignatureBuilder& signature, bool isAsyncMethod) override;
 
 public:
     TracerTokens(ModuleMetadata* module_metadata_ptr, bool enableByRefInstrumentation,