Add CompatTeamsSSOTokenExchangeMiddleware

core/src/Microsoft.Teams.Bot.Compat/CompatTeamsSSOTokenExchangeMiddleware.cs

@@ -0,0 +1,195 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Net;
+using Microsoft.Bot.Builder;
+using Microsoft.Bot.Connector;
+using Microsoft.Bot.Connector.Authentication;
+using Microsoft.Bot.Schema;
+using Newtonsoft.Json.Linq;
+
+namespace Microsoft.Teams.Bot.Compat;
+
+/// <summary>
+/// If the activity name is signin/tokenExchange, this middleware will attempt to
+/// exchange the token, and deduplicate the incoming call, ensuring only one
+/// exchange request is processed.
+/// </summary>
+/// <remarks>
+/// This is a compatibility reimplementation of
+/// <c>Microsoft.Bot.Builder.Teams.TeamsSSOTokenExchangeMiddleware</c> that works with
+/// the Teams Bot Core SDK via <see cref="CompatUserTokenClient"/>.
+///
+/// If a user is signed into multiple Teams clients, the Bot could receive a
+/// "signin/tokenExchange" from each client. Each token exchange request for a
+/// specific user login will have an identical Activity.Value.Id.
+///
+/// Only one of these token exchange requests should be processed by the bot.
+/// The others return <see cref="HttpStatusCode.PreconditionFailed"/>.
+/// For a distributed bot in production, this requires a distributed storage
+/// ensuring only one token exchange is processed. This middleware supports
+/// CosmosDb storage found in Microsoft.Bot.Builder.Azure, or MemoryStorage for
+/// local development. IStorage's ETag implementation for token exchange activity
+/// deduplication.
+/// </remarks>
+public class CompatTeamsSSOTokenExchangeMiddleware : IMiddleware
+{
+    private readonly IStorage _storage;
+    private readonly string _oAuthConnectionName;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="CompatTeamsSSOTokenExchangeMiddleware"/> class.
+    /// </summary>
+    /// <param name="storage">The <see cref="IStorage"/> to use for deduplication.</param>
+    /// <param name="connectionName">The connection name to use for the single sign on token exchange.</param>
+    public CompatTeamsSSOTokenExchangeMiddleware(IStorage storage, string connectionName)
+    {
+        ArgumentNullException.ThrowIfNull(storage);
+        ArgumentException.ThrowIfNullOrEmpty(connectionName);
+
+        _oAuthConnectionName = connectionName;
+        _storage = storage;
+    }
+
+    /// <inheritdoc/>
+    public async Task OnTurnAsync(ITurnContext turnContext, NextDelegate next, CancellationToken cancellationToken = default)
+    {
+        ArgumentNullException.ThrowIfNull(turnContext);
+        ArgumentNullException.ThrowIfNull(next);
+
+        if (string.Equals(Channels.Msteams, turnContext.Activity.ChannelId, StringComparison.OrdinalIgnoreCase)
+            && string.Equals(SignInConstants.TokenExchangeOperationName, turnContext.Activity.Name, StringComparison.OrdinalIgnoreCase))
+        {
+            // If the TokenExchange is NOT successful, the response will have already been sent by ExchangedTokenAsync
+            if (!await ExchangedTokenAsync(turnContext, cancellationToken).ConfigureAwait(false))
+            {
+                return;
+            }
+
+            // Only one token exchange should proceed from here. Deduplication is performed second because in the case
+            // of failure due to consent required, every caller needs to receive the response.
+            if (!await DeduplicatedTokenExchangeIdAsync(turnContext, cancellationToken).ConfigureAwait(false))
+            {
+                // If the token is not exchangeable, do not process this activity further.
+                return;
+            }
+        }
+
+        await next(cancellationToken).ConfigureAwait(false);
+    }
+
+    private async Task<bool> DeduplicatedTokenExchangeIdAsync(ITurnContext turnContext, CancellationToken cancellationToken)
+    {
+        // Create a StoreItem with Etag of the unique 'signin/tokenExchange' request
+        var storeItem = new TokenStoreItem
+        {
+            ETag = (turnContext.Activity.Value as JObject)!.Value<string>("id")
+        };
+
+        var storeItems = new Dictionary<string, object> { { TokenStoreItem.GetStorageKey(turnContext), storeItem } };
+        try
+        {
+            // Writing the IStoreItem with ETag of unique id will succeed only once
+            await _storage.WriteAsync(storeItems, cancellationToken).ConfigureAwait(false);
+        }
+        catch (Exception ex)
+
+            // Memory storage throws a generic exception with a Message of 'Etag conflict. [other error info]'
+            // CosmosDbPartitionedStorage throws: ex.Message.Contains("pre-condition is not met")
+            when (ex.Message.StartsWith("Etag conflict", StringComparison.OrdinalIgnoreCase) || ex.Message.Contains("pre-condition is not met", StringComparison.OrdinalIgnoreCase))
+        {
+            // Do NOT proceed processing this message, some other thread or machine already has processed it.
+
+            // Send 200 invoke response.
+            await SendInvokeResponseAsync(turnContext, cancellationToken: cancellationToken).ConfigureAwait(false);
+            return false;
+        }
+
+        return true;
+    }
+
+    private static async Task SendInvokeResponseAsync(ITurnContext turnContext, object? body = null, HttpStatusCode httpStatusCode = HttpStatusCode.OK, CancellationToken cancellationToken = default)
+    {
+        await turnContext.SendActivityAsync(
+            new Activity
+            {
+                Type = ActivityTypesEx.InvokeResponse,
+                Value = new InvokeResponse
+                {
+                    Status = (int)httpStatusCode,
+                    Body = body,
+                },
+            }, cancellationToken).ConfigureAwait(false);
+    }
+
+    private async Task<bool> ExchangedTokenAsync(ITurnContext turnContext, CancellationToken cancellationToken)
+    {
+        TokenResponse? tokenExchangeResponse = null;
+        var tokenExchangeRequest = ((JObject)turnContext.Activity.Value)?.ToObject<TokenExchangeInvokeRequest>();
+
+        try
+        {
+            var userTokenClient = turnContext.TurnState.Get<UserTokenClient>();
+            if (userTokenClient != null)
+            {
+                tokenExchangeResponse = await userTokenClient.ExchangeTokenAsync(
+                    turnContext.Activity.From.Id,
+                    _oAuthConnectionName,
+                    turnContext.Activity.ChannelId,
+                    new TokenExchangeRequest { Token = tokenExchangeRequest!.Token },
+                    cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                throw new NotSupportedException("Token Exchange is not supported by the current adapter.");
+            }
+        }
+#pragma warning disable CA1031 // Do not catch general exception types (matching BF SDK behavior)
+        catch
+#pragma warning restore CA1031
+        {
+            // Ignore Exceptions
+            // If token exchange failed for any reason, tokenExchangeResponse above stays null,
+            // and hence we send back a failure invoke response to the caller.
+        }
+
+        if (string.IsNullOrEmpty(tokenExchangeResponse?.Token))
+        {
+            // The token could not be exchanged (which could be due to a consent requirement)
+            // Notify the sender that PreconditionFailed so they can respond accordingly.
+
+            var invokeResponse = new TokenExchangeInvokeResponse
+            {
+                Id = tokenExchangeRequest!.Id,
+                ConnectionName = _oAuthConnectionName,
+                FailureDetail = "The bot is unable to exchange token. Proceed with regular login.",
+            };
+
+            await SendInvokeResponseAsync(turnContext, invokeResponse, HttpStatusCode.PreconditionFailed, cancellationToken).ConfigureAwait(false);
+
+            return false;
+        }
+
+        return true;
+    }
+
+    private class TokenStoreItem : IStoreItem
+    {
+        public string? ETag { get; set; }
+
+        public static string GetStorageKey(ITurnContext turnContext)
+        {
+            var activity = turnContext.Activity;
+            var channelId = activity.ChannelId ?? throw new InvalidOperationException("invalid activity-missing channelId");
+            var conversationId = activity.Conversation?.Id ?? throw new InvalidOperationException("invalid activity-missing Conversation.Id");
+
+            var value = activity.Value as JObject;
+            if (value == null || !value.ContainsKey("id"))
+            {
+                throw new InvalidOperationException("Invalid signin/tokenExchange. Missing activity.Value.Id.");
+            }
+
+            return $"{channelId}/{conversationId}/{value.Value<string>("id")}";
+        }
+    }
+}