Add session-aware XEvent logging for C# extensions

language-extensions/dotnet-core-CSharp/src/managed/CSharpExtension.cs

@@ -697,6 +697,10 @@ public struct SqlExtensionHostCallbacks
         /// 
         private const ushort MinSupportedHostCallbacksVersion = 1;
 
+        private static readonly uint MinimumHostCallbacksSize = checked(
+            (uint)(Marshal.OffsetOf<SqlExtensionHostCallbacks>(
+                nameof(SqlExtensionHostCallbacks.LogXEvent)).ToInt64() + IntPtr.Size));
+
         /// <summary>
         /// This delegate declares the delegate type of SetHostCallbacks.
         /// </summary>
@@ -739,6 +743,14 @@ public static short SetHostCallbacks(
                         hostCallbacks->Version);
                 }
 
+                if (hostCallbacks->SizeInBytes < MinimumHostCallbacksSize)
+                {
+                    Logging.Error("CSharpExtension::SetHostCallbacks: incomplete host callbacks structure");
+                    throw new ArgumentException(
+                        "Incomplete SQLEXTENSION_HOST_CALLBACKS structure.",
+                        nameof(hostCallbacks));
+                }
+
                 if (hostCallbacks->LogXEvent != IntPtr.Zero)
                 {
                     var logXEvent = Marshal.GetDelegateForFunctionPointer<LogXEventCallbackDelegate>(

language-extensions/dotnet-core-CSharp/src/managed/sdk/ExtensionEventLogger.cs

@@ -98,9 +98,23 @@ public static void LogVerbose(string message) =>
         /// <param name="level">Event severity.</param>
         /// <param name="message">Message to log. A null message is treated as empty.</param>
         /// <param name="errorCode">Optional error code for non-informational events.</param>
-        public static void Log(ExtensionTraceLevel level, string message, int errorCode = 0)
+        public static void Log(ExtensionTraceLevel level, string message, int errorCode = 0) =>
+            Log(level, message, errorCode, extensionName: null);
+
+        /// <summary>
+        /// Logs an event at the specified severity, attributed to a named extension.
+        /// </summary>
+        /// <param name="level">Event severity.</param>
+        /// <param name="message">Message to log. A null message is treated as empty.</param>
+        /// <param name="errorCode">Error code for non-informational events; 0 when unused.</param>
+        /// <param name="extensionName">
+        /// Name recorded as the event's originating extension. When null or empty the
+        /// Extension's default name is used. Lets an in-process library (e.g. one loaded
+        /// by a user script) attribute its events to itself rather than the host Extension.
+        /// </param>
+        public static void Log(ExtensionTraceLevel level, string message, int errorCode, string extensionName)
         {
-            Sink?.Log(level, errorCode, message);
+            Sink?.Log(level, errorCode, message, extensionName);
         }
     }
 }

language-extensions/dotnet-core-CSharp/src/managed/sdk/IExtensionLogSink.cs

@@ -34,6 +34,10 @@ internal interface IExtensionLogSink
         /// <param name="level">Event severity.</param>
         /// <param name="errorCode">Error code for non-informational events; 0 when unused.</param>
         /// <param name="message">Message to log. A null message is treated as empty.</param>
-        void Log(ExtensionTraceLevel level, int errorCode, string message);
+        /// <param name="extensionName">
+        /// Name recorded as the event's originating extension; null or empty selects the
+        /// Extension's default name.
+        /// </param>
+        void Log(ExtensionTraceLevel level, int errorCode, string message, string extensionName);
     }
 }

language-extensions/dotnet-core-CSharp/src/managed/utils/Logging.cs

@@ -138,14 +138,18 @@ internal static void ClearCurrentSession()
         /// <param name="traceLevel">Trace severity.</param>
         /// <param name="errorCode">Error code for non-informational logs.</param>
         /// <param name="message">The message to log.</param>
+        /// <param name="extensionName">
+        /// Originating extension name; null or empty selects the default name.
+        /// </param>
         internal static void LogXEventFromCurrentSession(
             ExtensionTraceLevel traceLevel,
             int                 errorCode,
-            string              message)
+            string              message,
+            string              extensionName = null)
         {
             (Guid sessionId, ushort taskId) = _currentSession.Value;
             LogXEvent(
-                extensionName: null,
+                extensionName: extensionName,
                 sessionId:     sessionId,
                 taskId:        taskId,
                 traceLevel:    traceLevel,

language-extensions/dotnet-core-CSharp/src/managed/utils/XEventLogSink.cs

@@ -27,7 +27,7 @@ internal sealed class XEventLogSink : IExtensionLogSink
         public bool IsEnabled => Logging.HasLogXEventCallback;
 
         /// <inheritdoc/>
-        public void Log(ExtensionTraceLevel level, int errorCode, string message) =>
-            Logging.LogXEventFromCurrentSession(level, errorCode, message);
+        public void Log(ExtensionTraceLevel level, int errorCode, string message, string extensionName) =>
+            Logging.LogXEventFromCurrentSession(level, errorCode, message, extensionName);
     }
 }

language-extensions/dotnet-core-CSharp/src/native/nativecsharpextension.cpp

@@ -359,8 +359,9 @@ SQLRETURN Cleanup()
     if (g_dotnet_runtime != nullptr)
     {
         SQLEXTENSION_HOST_CALLBACKS nullCallbacks = {};
-        nullCallbacks.Version   = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
-        nullCallbacks.LogXEvent = nullptr;
+        nullCallbacks.Version     = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
+        nullCallbacks.SizeInBytes = sizeof(nullCallbacks);
+        nullCallbacks.LogXEvent   = nullptr;
 
         g_dotnet_runtime->call_managed_method<decltype(&SetHostCallbacks)>(
             nameof(SetHostCallbacks),
@@ -563,13 +564,43 @@ SQLRETURN SetHostCallbacks(
         return SQL_ERROR;
     }
 
-    // Take a shallow copy of the caller's struct so g_hostCallbacks cannot
-    // dangle if the host passed a stack-allocated SQLEXTENSION_HOST_CALLBACKS.
+    // Smallest structure this Extension can consume: everything through LogXEvent,
+    // the last field it reads.
     //
-    g_hostCallbacksCopy = *hostCallbacks;
+    const SQLUINTEGER minHostCallbacksSize =
+        static_cast<SQLUINTEGER>(offsetof(SQLEXTENSION_HOST_CALLBACKS, LogXEvent) + sizeof(hostCallbacks->LogXEvent));
+    if (hostCallbacks->SizeInBytes < minHostCallbacksSize)
+    {
+        LOG_ERROR("SetHostCallbacks called with incomplete host callbacks structure");
+        return SQL_ERROR;
+    }
+
+    // The copy below is field-by-field, so a callback added to SQLEXTENSION_HOST_CALLBACKS
+    // is dropped silently: the Extension still builds and simply never forwards it. Fail
+    // the build instead, whether that callback claims a Reserved slot or is appended.
+    //
+    static_assert(
+        sizeof(SQLEXTENSION_HOST_CALLBACKS::Reserved1) == sizeof(void *) &&
+        sizeof(SQLEXTENSION_HOST_CALLBACKS::Reserved2) == sizeof(void *) &&
+        sizeof(SQLEXTENSION_HOST_CALLBACKS) ==
+            offsetof(SQLEXTENSION_HOST_CALLBACKS, Reserved2) + sizeof(void *),
+        "SQLEXTENSION_HOST_CALLBACKS layout changed: copy the new callback in "
+        "SetHostCallbacks and cover it in the forwarded SizeInBytes below.");
+
+    // Cap the forwarded extent at what this build can hold. Reporting
+    // sizeof(g_hostCallbacksCopy) unconditionally would tell the managed layer that
+    // trailing fields a smaller, older host never supplied are valid.
+    //
+    g_hostCallbacksCopy = {};
+    g_hostCallbacksCopy.Version     = hostCallbacks->Version;
+    g_hostCallbacksCopy.Reserved0   = hostCallbacks->Reserved0;
+    g_hostCallbacksCopy.SizeInBytes = hostCallbacks->SizeInBytes < sizeof(g_hostCallbacksCopy)
+        ? hostCallbacks->SizeInBytes
+        : static_cast<SQLUINTEGER>(sizeof(g_hostCallbacksCopy));
+    g_hostCallbacksCopy.LogXEvent   = hostCallbacks->LogXEvent;
     g_hostCallbacks     = &g_hostCallbacksCopy;
 
     return g_dotnet_runtime->call_managed_method<decltype(&SetHostCallbacks)>(
         nameof(SetHostCallbacks),
-        hostCallbacks);
-}
+        &g_hostCallbacksCopy);
+}

language-extensions/dotnet-core-CSharp/test/src/managed/CSharpTestExecutor.cs

@@ -36,6 +36,19 @@ internal static class CSharpTestExecutorConstants
         /// so both sides must stay in sync.
         /// </summary>
         public const string LogEventMessage = "CSharpTestExecutorLogInformation emitted event";
+
+        /// <summary>
+        /// Marker message emitted through the SDK ExtensionEventLogger extensionName
+        /// overload by CSharpTestExecutorLogNamedExtension. Matched by the native test
+        /// (CSharpExecuteTests.cpp) alongside <see cref="LogEventExtensionName"/>.
+        /// </summary>
+        public const string LogNamedEventMessage = "CSharpTestExecutorLogNamedExtension emitted event";
+
+        /// <summary>
+        /// Extension name CSharpTestExecutorLogNamedExtension attributes its event to,
+        /// proving the ExtensionEventLogger.Log extensionName overload is honored end to end.
+        /// </summary>
+        public const string LogEventExtensionName = "TestExtension";
     }
 
     public class CSharpTestExecutor: AbstractSqlServerExtensionExecutor
@@ -60,6 +73,23 @@ public override DataFrame Execute(DataFrame input, Dictionary<string, dynamic> s
         }
     }
 
+    /// <summary>
+    /// Test executor that emits an event through the ExtensionEventLogger extensionName
+    /// overload, attributing it to a named extension. Lets the native test assert the
+    /// forwarded XEvent carries the caller-supplied extension name rather than the default.
+    /// </summary>
+    public class CSharpTestExecutorLogNamedExtension: AbstractSqlServerExtensionExecutor
+    {
+        public override DataFrame Execute(DataFrame input, Dictionary<string, dynamic> sqlParams){
+            ExtensionEventLogger.Log(
+                ExtensionTraceLevel.Information,
+                CSharpTestExecutorConstants.LogNamedEventMessage,
+                errorCode: 0,
+                extensionName: CSharpTestExecutorConstants.LogEventExtensionName);
+            return input;
+        }
+    }
+
     public class CSharpTestExecutorIntParam: AbstractSqlServerExtensionExecutor
     {
         public override DataFrame Execute(DataFrame input, Dictionary<string, dynamic> sqlParams){

language-extensions/dotnet-core-CSharp/test/src/native/CSharpExecuteTests.cpp

@@ -547,6 +547,56 @@ namespace ExtensionApiTest
         EXPECT_EQ(tagged->traceLevel, static_cast<SQLUSMALLINT>(Extension_Information));
     }
 
+    //----------------------------------------------------------------------------------------------
+    // Name: ExecuteForwardsNamedExtensionLogEvent
+    //
+    // Description:
+    //  Drives an executor that logs through the ExtensionEventLogger.Log extensionName
+    //  overload and verifies the forwarded XEvent carries the caller-supplied extension
+    //  name rather than the extension's default. Covers the SDK facade path an in-process
+    //  library (e.g. MssqlAI) uses to attribute its own activity events.
+    //
+    TEST_F(CSharpExtensionApiTests, ExecuteForwardsNamedExtensionLogEvent)
+    {
+        ASSERT_EQ(RegisterTestLogXEventCallback(sm_libHandle), SQL_SUCCESS);
+
+        const string namedExecutor =
+            "Microsoft.SqlServer.CSharpExtensionTest.CSharpTestExecutorLogNamedExtension";
+        const string namedMessage = "CSharpTestExecutorLogNamedExtension emitted event";
+        const string expectedExtensionName = "TestExtension";
+
+        string scriptString = m_UserLibName + m_Separator + namedExecutor;
+        InitializeSession(0, 0, scriptString);
+
+        g_capturedLogEvents.clear();
+
+        SQLUSMALLINT outputSchemaColumnsNumber = 0;
+        SQLRETURN result = (*sm_executeFuncPtr)(
+            *m_sessionId,
+            m_taskId,
+            0,       // rowsNumber
+            nullptr, // dataSet
+            nullptr, // strLenOrInd
+            &outputSchemaColumnsNumber);
+        ASSERT_EQ(result, SQL_SUCCESS);
+
+        const CapturedLogEvent *named = nullptr;
+        for (const CapturedLogEvent &ev : g_capturedLogEvents)
+        {
+            if (ev.message.find(namedMessage) != string::npos)
+            {
+                named = &ev;
+                break;
+            }
+        }
+
+        ASSERT_NE(named, nullptr)
+            << "Named-extension ExtensionEventLogger event was not forwarded to the host callback";
+        EXPECT_EQ(named->extensionName, expectedExtensionName)
+            << "Forwarded event did not carry the caller-supplied extension name";
+        EXPECT_EQ(named->traceLevel, static_cast<SQLUSMALLINT>(Extension_Information));
+    }
+
     //----------------------------------------------------------------------------------------------
     // Name: ExecuteIsolatesSessionTaggingBetweenSessions
     //

language-extensions/dotnet-core-CSharp/test/src/native/CSharpSetHostCallbacksTests.cpp

@@ -14,6 +14,7 @@
 #include "CSharpExtensionApiTests.h"
 #include "LogXEventTestHarness.h"
 
+#include <cstddef>
 #include <cstring>
 #include <string>
 #include <vector>
@@ -58,6 +59,43 @@ namespace ExtensionApiTest
         EXPECT_EQ(rc, SQL_ERROR);
     }
 
+    TEST_F(CSharpExtensionApiTests, SetHostCallbacks_IncompleteStructureReturnsError)
+    {
+        FN_setHostCallbacks *fn = RESOLVE_SET_HOST_CALLBACKS();
+        ASSERT_NE(fn, nullptr);
+
+        SQLEXTENSION_HOST_CALLBACKS hostCallbacks{};
+        hostCallbacks.Version = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
+        hostCallbacks.SizeInBytes = offsetof(SQLEXTENSION_HOST_CALLBACKS, LogXEvent);
+        hostCallbacks.LogXEvent = &TestLogXEventCallback;
+
+        EXPECT_EQ(fn(&hostCallbacks), SQL_ERROR);
+    }
+
+    TEST_F(CSharpExtensionApiTests, SetHostCallbacks_MinimumV1StructureIsAccepted)
+    {
+        FN_setHostCallbacks *fn = RESOLVE_SET_HOST_CALLBACKS();
+        ASSERT_NE(fn, nullptr);
+
+        struct V1HostCallbacks
+        {
+            SQLUSMALLINT             Version;
+            SQLUSMALLINT             Reserved0;
+            SQLUINTEGER              SizeInBytes;
+            PFunc_ExtensionLogXEvent LogXEvent;
+        };
+
+        g_capturedLogEvents.clear();
+
+        V1HostCallbacks hostCallbacks{};
+        hostCallbacks.Version = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
+        hostCallbacks.SizeInBytes = sizeof(hostCallbacks);
+        hostCallbacks.LogXEvent = &TestLogXEventCallback;
+
+        EXPECT_EQ(fn(reinterpret_cast<SQLEXTENSION_HOST_CALLBACKS *>(&hostCallbacks)), SQL_SUCCESS);
+        EXPECT_FALSE(g_capturedLogEvents.empty());
+    }
+
     //----------------------------------------------------------------------------------------------
     // Name: SetHostCallbacks_RegistersAndForwardsLogXEvent
     //
@@ -75,8 +113,9 @@ namespace ExtensionApiTest
         g_capturedLogEvents.clear();
 
         SQLEXTENSION_HOST_CALLBACKS hostCallbacks{};
-        hostCallbacks.Version    = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
-        hostCallbacks.LogXEvent  = &TestLogXEventCallback;
+        hostCallbacks.Version     = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
+        hostCallbacks.SizeInBytes = sizeof(hostCallbacks);
+        hostCallbacks.LogXEvent   = &TestLogXEventCallback;
 
         SQLRETURN rc = fn(&hostCallbacks);
         EXPECT_EQ(rc, SQL_SUCCESS);
@@ -109,8 +148,9 @@ namespace ExtensionApiTest
         g_capturedLogEvents.clear();
 
         SQLEXTENSION_HOST_CALLBACKS hostCallbacks{};
-        hostCallbacks.Version   = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
-        hostCallbacks.LogXEvent = nullptr;
+        hostCallbacks.Version     = SQLEXTENSION_HOST_CALLBACKS_VERSION_1;
+        hostCallbacks.SizeInBytes = sizeof(hostCallbacks);
+        hostCallbacks.LogXEvent   = nullptr;
 
         SQLRETURN rc = fn(&hostCallbacks);
         EXPECT_EQ(rc, SQL_SUCCESS);