fix: fix bot-sso race condition issue

bot-sso/commands/index.ts

@@ -1,4 +1,4 @@
-import { BotCommand } from "../helpers/botCommand";
+import { BotCommand, SSOCommand } from "../helpers/botCommand";
 import { LearnCommand } from "./learn";
 import { ShowUserProfile } from "./showUserProfile";
 import { WelcomeCommand } from "./welcome";
@@ -8,3 +8,12 @@ export const commands: BotCommand[] = [
   new ShowUserProfile(),
   new WelcomeCommand(),
 ];
+
+export const SSOCommands: SSOCommand[] = [
+  new ShowUserProfile(),
+];
+
+export const SSOCommandMap: Map<string, any> = new Map();
+SSOCommands.forEach((command) => {
+  SSOCommandMap.set(command.constructor.name, command.operationWithSSOToken);
+});

bot-sso/helpers/botCommand.ts

@@ -50,7 +50,7 @@ export class SSOCommand extends BotCommand {
   async run(parameters: any): Promise<any> {
     this.validateParameters(parameters);
     const ssoDialog = parameters.ssoDialog;
-    ssoDialog.setSSOOperation(this.operationWithSSOToken);
+    await ssoDialog.setSSOOperation(parameters.context, this.constructor.name);
     await ssoDialog.run(parameters.context, parameters.dialogState);
   }
 }

bot-sso/helpers/ssoDialog.ts

@@ -7,14 +7,17 @@ import {
 } from "botbuilder-dialogs";
 import {
   ActivityTypes,
+  StatePropertyAccessor,
   Storage,
   tokenExchangeOperationName,
   TurnContext,
+  UserState,
 } from "botbuilder";
 import { TeamsBotSsoPrompt } from "@microsoft/teamsfx";
 import "isomorphic-fetch";
 import oboAuthConfig from "../authConfig";
 import config from "../config";
+import { SSOCommandMap } from "../commands";
 
 const DIALOG_NAME = "SSODialog";
 const MAIN_WATERFALL_DIALOG = "MainWaterfallDialog";
@@ -24,17 +27,14 @@ export class SSODialog extends ComponentDialog {
   private requiredScopes: string[] = ["User.Read"]; // hard code the scopes for demo purpose only
   private dedupStorage: Storage;
   private dedupStorageKeys: string[];
-  private operationWithSSO: (
-    arg0: any,
-    ssoToken: string
-  ) => Promise<any> | undefined;
+  private userStateAccessor: StatePropertyAccessor<any>;
 
   // Developer controlls the lifecycle of credential provider, as well as the cache in it.
   // In this sample the provider is shared in all conversations
-  constructor(dedupStorage: Storage) {
+  constructor(userState: UserState, dedupStorage: Storage) {
     super(DIALOG_NAME);
 
-    const initialLoginEndpoint =`https://${config.botDomain}/auth-start.html` ;
+    const initialLoginEndpoint = `https://${config.botDomain}/auth-start.html`;
 
     const dialog = new TeamsBotSsoPrompt(
       oboAuthConfig,
@@ -57,17 +57,18 @@ export class SSODialog extends ComponentDialog {
 
     this.initialDialogId = MAIN_WATERFALL_DIALOG;
     this.dedupStorage = dedupStorage;
+    this.userStateAccessor = userState.createProperty<any>('operationWithSSO');
     this.dedupStorageKeys = [];
   }
 
-  setSSOOperation(
-    handler: (arg0: any, arg1: string) => Promise<any> | undefined
+  async setSSOOperation(context: TurnContext,
+    handler: string
   ) {
-    this.operationWithSSO = handler;
+    await this.userStateAccessor.set(context, handler);
   }
 
-  resetSSOOperation() {
-    this.operationWithSSO = undefined;
+  async resetSSOOperation(context: TurnContext) {
+    await this.userStateAccessor.delete(context);
   }
 
   /**
@@ -106,9 +107,11 @@ export class SSODialog extends ComponentDialog {
         "There is an issue while trying to sign you in and retrieve your profile photo, please type \"show\" command to login and consent permissions again."
       );
     } else {
+      const SSOCommandName = (await this.userStateAccessor.get(stepContext.context));
       // Once got ssoToken, run operation that depends on ssoToken
-      if (this.operationWithSSO) {
-        await this.operationWithSSO(stepContext.context, tokenResponse.ssoToken);
+      if (SSOCommandName) {
+        const operationWithSSO = SSOCommandMap.get(SSOCommandName);
+        await operationWithSSO(stepContext.context, tokenResponse.ssoToken);
       }
     }
     return await stepContext.endDialog();
@@ -123,7 +126,7 @@ export class SSODialog extends ComponentDialog {
     this.dedupStorageKeys = this.dedupStorageKeys.filter(
       (key) => key.indexOf(conversationId) < 0
     );
-    this.resetSSOOperation();
+    this.resetSSOOperation(context);
   }
 
   // If a user is signed into multiple Teams clients, the Bot might receive a "signin/tokenExchange" from each client.

bot-sso/package.json

@@ -41,4 +41,4 @@
     "nodemon": "^2.0.7",
     "shx": "^0.3.3"
   }
-}
+}

bot-sso/teamsBot.ts

@@ -18,7 +18,7 @@ const rawLearnCard = require("./adaptiveCards/learn.json");
 export class TeamsBot extends TeamsActivityHandler {
   likeCountObj: { likeCount: number };
   conversationState: BotState;
-  userState: BotState;
+  userState: UserState;
   dialog: SSODialog;
   dialogState: any;
   commandsHelper: CommandsHelper;
@@ -37,7 +37,7 @@ export class TeamsBot extends TeamsActivityHandler {
     // Create conversation and user state with in-memory storage provider.
     this.conversationState = new ConversationState(memoryStorage);
     this.userState = new UserState(memoryStorage);
-    this.dialog = new SSODialog(new MemoryStorage());
+    this.dialog = new SSODialog(this.userState, new MemoryStorage());
     this.dialogState = this.conversationState.createProperty("DialogState");
 
     this.onMessage(async (context, next) => {