From 7b9ee30caa009789a5310a2793322d286828e48e Mon Sep 17 00:00:00 2001
From: Amanda Tarafa Mas <atarafamas@google.com>
Date: Thu, 9 Jan 2020 16:13:53 +0000
Subject: [PATCH] Support for ITokenAccessWithHeaders

---
 .../Grpc.Auth/GoogleAuthInterceptors.cs       | 28 ++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/csharp/Grpc.Auth/GoogleAuthInterceptors.cs b/src/csharp/Grpc.Auth/GoogleAuthInterceptors.cs
index e6c973431d7..b3ffe356972 100644
--- a/src/csharp/Grpc.Auth/GoogleAuthInterceptors.cs
+++ b/src/csharp/Grpc.Auth/GoogleAuthInterceptors.cs
@@ -16,7 +16,6 @@
 
 #endregion
 
-using System;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -43,6 +42,11 @@ namespace Grpc.Auth
         /// <returns>The interceptor.</returns>
         public static AsyncAuthInterceptor FromCredential(ITokenAccess credential)
         {
+            if (credential is ITokenAccessWithHeaders credentialWithHeaders)
+            {
+                return FromCredential(credentialWithHeaders);
+            }
+
             return new AsyncAuthInterceptor(async (context, metadata) =>
             {
                 var accessToken = await credential.GetAccessTokenForRequestAsync(context.ServiceUrl, CancellationToken.None).ConfigureAwait(false);
@@ -50,6 +54,28 @@ namespace Grpc.Auth
             });
         }
 
+        /// <summary>
+        /// Creates an <see cref="AsyncAuthInterceptor"/> that will obtain access token and associated information
+        /// from any credential type that implements <see cref="ITokenAccessWithHeaders"/>
+        /// </summary>
+        /// <param name="credential">The credential to use to obtain access tokens.</param>
+        /// <returns>The interceptor.</returns>
+        public static AsyncAuthInterceptor FromCredential(ITokenAccessWithHeaders credential)
+        {
+            return new AsyncAuthInterceptor(async (context, metadata) => 
+            {
+                AccessTokenWithHeaders tokenAndHeaders = await credential.GetAccessTokenWithHeadersForRequestAsync(context.ServiceUrl, CancellationToken.None).ConfigureAwait(false);
+                metadata.Add(CreateBearerTokenHeader(tokenAndHeaders.AccessToken));
+                foreach (var header in tokenAndHeaders.Headers)
+                {
+                    foreach (var headerValue in header.Value)
+                    {
+                        metadata.Add(new Metadata.Entry(header.Key, headerValue));
+                    }
+                }
+            });
+        }
+
         /// <summary>
         /// Creates an <see cref="AsyncAuthInterceptor"/> that will use given access token as authorization.
         /// </summary>