Add Feature Extension for User Agent (#3451)

* Add Feature Extension for User Agent

* resolve conflicts

* Add summary and review changes

* Add unit tests for WriteUserAgentFeatureRequest

* Update tests to avoid reflection

Author: samsharma2700

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs

@@ -207,6 +207,9 @@ internal bool IsDNSCachingBeforeRedirectSupported
         // Json Support Flag
         internal bool IsJsonSupportEnabled = false;
 
+        // User Agent Flag
+        internal bool IsUserAgentEnabled = true;
+
         // Vector Support Flag
         internal bool IsVectorSupportEnabled = false;
 
@@ -1430,6 +1433,10 @@ private void Login(ServerInfo server, TimeoutTimer timeout, string newPassword,
             requestedFeatures |= TdsEnums.FeatureExtension.JsonSupport;
             requestedFeatures |= TdsEnums.FeatureExtension.VectorSupport;
 
+        #if DEBUG
+            requestedFeatures |= TdsEnums.FeatureExtension.UserAgent;
+        #endif
+
             _parser.TdsLogin(login, requestedFeatures, _recoverySessionData, _fedAuthFeatureExtensionData, encrypt);
         }
 

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/TdsParser.cs

@@ -8579,6 +8579,49 @@ internal int WriteVectorSupportFeatureRequest(bool write)
             return len;
         }
 
+        /// <summary>
+        /// Writes the User Agent feature request to the physical state object.
+        /// The request includes the feature ID, feature data length, version number and encoded JSON payload.
+        /// </summary>
+        /// <param name="userAgentJsonPayload"> Byte array of UTF-8 encoded JSON payload for User Agent</param>
+        /// <param name="write">
+        /// If true, writes the feature request to the physical state object.
+        /// If false, just calculates the length.
+        /// </param>
+        /// <returns>The length of the feature request in bytes.</returns>
+        /// <remarks>
+        /// The feature request consists of:
+        /// - 1 byte for the feature ID.
+        /// - 4 bytes for the feature data length.
+        /// - 1 byte for the version number.
+        /// - N bytes for the JSON payload
+        /// </remarks>
+        internal int WriteUserAgentFeatureRequest(byte[] userAgentJsonPayload,
+                                                  bool write)
+        {
+            // 1byte (Feature Version) + size of UTF-8 encoded JSON payload 
+            int dataLen = 1 + userAgentJsonPayload.Length;
+            // 1byte (Feature ID) + 4bytes (Feature Data Length) + 1byte (Version) + N(JSON payload size)
+            int totalLen = 1 + 4 + dataLen;
+            
+            if (write)
+            {   
+                // Write Feature ID
+                _physicalStateObj.WriteByte(TdsEnums.FEATUREEXT_USERAGENT);
+
+                // Feature Data Length
+                WriteInt(dataLen, _physicalStateObj);
+
+                // Write Feature Version
+                _physicalStateObj.WriteByte(TdsEnums.SUPPORTED_USER_AGENT_VERSION);
+
+                // Write encoded JSON payload
+                _physicalStateObj.WriteByteArray(userAgentJsonPayload, userAgentJsonPayload.Length, 0);
+            }
+
+            return totalLen;
+        }
+
         private void WriteLoginData(SqlLogin rec,
                                      TdsEnums.FeatureExtension requestedFeatures,
                                      SessionData recoverySessionData,

src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs

@@ -211,6 +211,9 @@ internal bool IsDNSCachingBeforeRedirectSupported
         // Vector Support Flag
         internal bool IsVectorSupportEnabled = false;
 
+        // User Agent Flag
+        internal bool IsUserAgentEnabled = true;
+
         // TCE flags
         internal byte _tceVersionSupported;
 
@@ -1436,6 +1439,10 @@ private void Login(ServerInfo server, TimeoutTimer timeout, string newPassword,
             requestedFeatures |= TdsEnums.FeatureExtension.JsonSupport;
             requestedFeatures |= TdsEnums.FeatureExtension.VectorSupport;
 
+        #if DEBUG    
+            requestedFeatures |= TdsEnums.FeatureExtension.UserAgent;
+        #endif
+
             _parser.TdsLogin(login, requestedFeatures, _recoverySessionData, _fedAuthFeatureExtensionData, encrypt);
         }
 

src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/TdsParser.cs

@@ -8762,6 +8762,49 @@ internal int WriteVectorSupportFeatureRequest(bool write)
             return len;
         }
 
+        /// <summary>
+        /// Writes the User Agent feature request to the physical state object.
+        /// The request includes the feature ID, feature data length, version number and encoded JSON payload.
+        /// </summary>
+        /// <param name="userAgentJsonPayload"> Byte array of UTF-8 encoded JSON payload for User Agent</param>
+        /// <param name="write">
+        /// If true, writes the feature request to the physical state object.
+        /// If false, just calculates the length.
+        /// </param>
+        /// <returns>The length of the feature request in bytes.</returns>
+        /// <remarks>
+        /// The feature request consists of:
+        /// - 1 byte for the feature ID.
+        /// - 4 bytes for the feature data length.
+        /// - 1 byte for the version number.
+        /// - N bytes for the JSON payload
+        /// </remarks>
+        internal int WriteUserAgentFeatureRequest(byte[] userAgentJsonPayload,
+                                                  bool write)
+        {
+            // 1byte (Feature Version) + size of UTF-8 encoded JSON payload 
+            int dataLen = 1 + userAgentJsonPayload.Length;
+            // 1byte (Feature ID) + 4bytes (Feature Data Length) + 1byte (Version) + N(JSON payload size)
+            int totalLen = 1 + 4 + dataLen;
+
+            if (write)
+            {
+                // Write Feature ID
+                _physicalStateObj.WriteByte(TdsEnums.FEATUREEXT_USERAGENT);
+
+                // Feature Data Length
+                WriteInt(dataLen, _physicalStateObj);
+
+                // Write Feature Version
+                _physicalStateObj.WriteByte(TdsEnums.SUPPORTED_USER_AGENT_VERSION);
+
+                // Write encoded JSON payload
+                _physicalStateObj.WriteByteArray(userAgentJsonPayload, userAgentJsonPayload.Length, 0);
+            }
+
+            return totalLen;
+        }
+
         private void WriteLoginData(SqlLogin rec,
                                     TdsEnums.FeatureExtension requestedFeatures,
                                     SessionData recoverySessionData,

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/TdsEnums.cs

@@ -241,6 +241,8 @@ public enum EnvChangeType : byte
         public const byte FEATUREEXT_SQLDNSCACHING = 0x0B;
         public const byte FEATUREEXT_JSONSUPPORT = 0x0D;
         public const byte FEATUREEXT_VECTORSUPPORT = 0x0E;
+        // TODO: re-verify if this byte competes with another feature
+        public const byte FEATUREEXT_USERAGENT = 0x0F;
 
         [Flags]
         public enum FeatureExtension : uint
@@ -255,7 +257,8 @@ public enum FeatureExtension : uint
             UTF8Support = 1 << (TdsEnums.FEATUREEXT_UTF8SUPPORT - 1),
             SQLDNSCaching = 1 << (TdsEnums.FEATUREEXT_SQLDNSCACHING - 1),
             JsonSupport = 1 << (TdsEnums.FEATUREEXT_JSONSUPPORT - 1),
-            VectorSupport = 1 << (TdsEnums.FEATUREEXT_VECTORSUPPORT - 1)
+            VectorSupport = 1 << (TdsEnums.FEATUREEXT_VECTORSUPPORT - 1),
+            UserAgent = 1 << (TdsEnums.FEATUREEXT_USERAGENT - 1)
         }
 
         public const uint UTF8_IN_TDSCOLLATION = 0x4000000;
@@ -985,6 +988,9 @@ internal enum FedAuthInfoId : byte
         internal const byte MAX_SUPPORTED_VECTOR_VERSION = 0x01;
         internal const int VECTOR_HEADER_SIZE = 8;
 
+        // User Agent constants
+        internal const byte SUPPORTED_USER_AGENT_VERSION = 0x01;
+
         // TCE Related constants
         internal const byte MAX_SUPPORTED_TCE_VERSION = 0x03; // max version
         internal const byte MIN_TCE_VERSION_WITH_ENCLAVE_SUPPORT = 0x02; // min version with enclave support

src/Microsoft.Data.SqlClient/tests/UnitTests/TdsParserInternalsTest.cs

@@ -0,0 +1,107 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+using System;
+using System.Linq;
+using System.Reflection;
+using System.Text;
+using Xunit;
+
+#nullable enable
+
+namespace Microsoft.Data.SqlClient.UnitTests
+{
+    public class TdsParserInternalsTest
+    {
+        private readonly TdsParser _parser = new(false, false);
+
+        // TODO(ADO-37888): Avoid reflection by exposing a way for tests to intercept outbound TDS packets.
+        // Helper function to extract private _physicalStateObj fields raw buffer and no. of bytes written so far 
+        private static (byte[] buffer, int count) ExtractOutputBuffer(TdsParser parser)
+        {
+            FieldInfo stateField = typeof(TdsParser)
+                .GetField("_physicalStateObj", BindingFlags.Instance | BindingFlags.NonPublic)
+                ?? throw new InvalidOperationException("_physicalStateObj not found");
+
+            object stateObj = stateField.GetValue(parser)
+                ?? throw new InvalidOperationException("physical state object is null");
+
+            Type stateType = stateObj.GetType();
+
+            FieldInfo buffField = stateType
+                .GetField("_outBuff", BindingFlags.Instance | BindingFlags.NonPublic)
+                ?? throw new InvalidOperationException("_outBuff not found");
+
+            byte[] buffer = (byte[])buffField.GetValue(stateObj)!;
+
+            FieldInfo usedField = stateType
+                .GetField("_outBytesUsed", BindingFlags.Instance | BindingFlags.NonPublic)
+                ?? throw new InvalidOperationException("_outBytesUsed not found");
+
+            int count = (int)usedField.GetValue(stateObj)!;
+
+            return (buffer, count);
+        }
+
+        [Fact]
+        public void WriteUserAgentFeatureRequest_WriteFalse_LengthOnlyReturn()
+        {
+            byte[] payload = Encoding.UTF8.GetBytes("{\"kel\":\"sier\"}");
+            var (_, countBefore) = ExtractOutputBuffer(_parser);
+
+            int lengthOnly = _parser.WriteUserAgentFeatureRequest(payload, write: false);
+
+            var (_, countAfter) = ExtractOutputBuffer(_parser);
+
+            // assert: total = 1 (feat-ID) + 4 (len field) + [1 (version) + payload.Length]
+            int expectedDataLen = 1 + payload.Length;
+            int expectedTotalLen = 1 + 4 + expectedDataLen;
+            Assert.Equal(expectedTotalLen, lengthOnly);
+
+            // assert: no bytes were written when write == false
+            Assert.Equal(countBefore, countAfter);
+        }
+
+        [Fact]
+        public void WriteUserAgentFeatureRequest_WriteTrue_AppendsOnlyExtensionBytes()
+        {
+            byte[] payload = Encoding.UTF8.GetBytes("{\"kel\":\"sier\"}");
+            var (bufferBefore, countBefore) = ExtractOutputBuffer(_parser);
+
+            int returnedLength = _parser.WriteUserAgentFeatureRequest(payload, write: true);
+
+            var (bufferAfter, countAfter) = ExtractOutputBuffer(_parser);
+
+            // We expect both of these to be the same object
+            Assert.Same(bufferBefore, bufferAfter);
+            int appended = countAfter - countBefore;
+            Assert.Equal(returnedLength, appended);
+
+            int start = countBefore;
+
+            Assert.Equal(
+                TdsEnums.FEATUREEXT_USERAGENT,
+                bufferAfter[start]);
+
+            int dataLenFromStream = BitConverter.ToInt32(bufferAfter, start + 1);
+            int expectedDataLen = 1 + payload.Length;
+            Assert.Equal(expectedDataLen, dataLenFromStream);
+
+            Assert.Equal(
+                TdsEnums.SUPPORTED_USER_AGENT_VERSION,
+                bufferAfter[start + 5]);
+
+            // slice into the existing buffer
+            ReadOnlySpan<byte> writtenSpan = new ReadOnlySpan<byte>(
+                bufferAfter,
+                start + 6,
+                appended - 6);
+
+            Assert.True(
+                writtenSpan.SequenceEqual(payload),
+                "Payload bytes did not match");
+
+        }
+
+    }
+}