feat: add user agent to default KMS clients (#598)

aws · Oct 11, 2023 · f78858c · f78858c
1 parent 55196e0
commit f78858c
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 10 deletions.
diff --git a/aws-encryption-sdk-net/Source/AWSEncryptionSDK.csproj b/aws-encryption-sdk-net/Source/AWSEncryptionSDK.csproj
@@ -1,13 +1,15 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
-
   <PropertyGroup>
     <OutputType>Library</OutputType>
     <TargetFrameworks>netstandard2.1;net452</TargetFrameworks>
     <LangVersion>7.3</LangVersion>
     <GenerateAssemblyInfo>false</GenerateAssemblyInfo>
     <IsPackable>true</IsPackable>
+
+    <!-- This should be kept in sync with the version number in AssemblyInfo.cs -->
     <Version>3.0.0</Version>
+
     <AssemblyName>AWS.EncryptionSDK</AssemblyName>
     <PackageId>AWS.EncryptionSDK</PackageId>
     <Title>AWS Encryption SDK for .NET</Title>

diff --git a/aws-encryption-sdk-net/Source/AssemblyInfo.cs b/aws-encryption-sdk-net/Source/AssemblyInfo.cs
@@ -0,0 +1,6 @@
+using System.Reflection;
+
+[assembly: AssemblyTitle("AWS.EncryptionSDK")]
+
+// This should be kept in sync with the version number in AWSEncryptionSDK.csproj
+[assembly: AssemblyVersion("3.0.0")]
diff --git a/aws-encryption-sdk-net/Source/Extern/DefaultClientSupplier.cs b/aws-encryption-sdk-net/Source/Extern/DefaultClientSupplier.cs
@@ -1,9 +1,12 @@
 // Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+using System.Threading.Tasks;
 using Amazon;
 using Amazon.KeyManagementService;
 using Amazon.Runtime;
+using Amazon.Runtime.Internal;
+using Amazon.Util;
 // ReSharper disable once RedundantUsingDirective
 using AWS.EncryptionSDK.Core;
 
@@ -24,16 +27,14 @@ public partial class DefaultClientSupplier : Dafny.Aws.EncryptionSdk.Core.IClien
                 AWS.EncryptionSDK.Core.TypeConversion.FromDafny_N3_aws__N13_encryptionSdk__N4_core__S14_GetClientInput(input);
             try
             {
-                IAmazonKeyManagementService client;
-                if (convertedInput.Region != "")
+                var regionEndpoint = string.IsNullOrEmpty(convertedInput.Region)
+                    ? null
+                    : RegionEndpoint.GetBySystemName(convertedInput.Region);
+                var clientConfig = new AmazonKeyManagementServiceConfig
                 {
-                    var regionEndpoint = RegionEndpoint.GetBySystemName(convertedInput.Region);
-                    client = new AmazonKeyManagementServiceClient(regionEndpoint);
-                }
-                else
-                {
-                    client = new AmazonKeyManagementServiceClient();
-                }
+                    RegionEndpoint = regionEndpoint
+                };
+                var client = new DefaultKmsClient(clientConfig);
 
                 // ReSharper disable once RedundantNameQualifier
                 return Wrappers_Compile.Result<Dafny.Com.Amazonaws.Kms.IKeyManagementServiceClient,
@@ -50,4 +51,55 @@ public partial class DefaultClientSupplier : Dafny.Aws.EncryptionSdk.Core.IClien
             }
         }
     }
+
+    /// <summary>
+    /// A KMS client that adds the Encryption SDK version to the user agent.
+    /// </summary>
+    internal class DefaultKmsClient : AmazonKeyManagementServiceClient
+    {
+        public DefaultKmsClient(AmazonKeyManagementServiceConfig config) : base(config)
+        {
+        }
+
+        protected override void CustomizeRuntimePipeline(RuntimePipeline pipeline)
+        {
+            base.CustomizeRuntimePipeline(pipeline);
+            pipeline.AddHandlerAfter<Marshaller>(new UserAgentHandler());
+        }
+    }
+
+    /// <summary>
+    /// Adds the Encryption SDK version to the user agent.
+    /// </summary>
+    internal class UserAgentHandler : PipelineHandler
+    {
+        private static readonly string UserAgentSuffix;
+
+        static UserAgentHandler()
+        {
+            var version = typeof(UserAgentHandler).Assembly.GetName().Version;
+            var semver = $"{version.Major}.{version.Minor}.{version.Build}";
+            UserAgentSuffix = $" AwsEncryptionSdkNet/{semver}";
+        }
+
+        /// <inheritdoc />
+        public override void InvokeSync(IExecutionContext executionContext)
+        {
+            AddUserAgent(executionContext);
+            base.InvokeSync(executionContext);
+        }
+
+        /// <inheritdoc />
+        public override Task<T> InvokeAsync<T>(IExecutionContext executionContext)
+        {
+            AddUserAgent(executionContext);
+            return base.InvokeAsync<T>(executionContext);
+        }
+
+        private static void AddUserAgent(IExecutionContext executionContext)
+        {
+            var request = executionContext.RequestContext.Request;
+            request.Headers[AWSSDKUtils.UserAgentHeader] += UserAgentSuffix;
+        }
+    }
 }