From 2b49c22e7fd158677b63d37852445bfc4dc62b7f Mon Sep 17 00:00:00 2001
From: Flavien David <flavien.david74@gmail.com>
Date: Tue, 9 Jan 2024 16:06:07 +0100
Subject: [PATCH] Flav/whitelist slack domains (#3115)

* Add slack whitelistedDomains

* Assert slack user has permission to use chatbot

* :sparkles:

* :book:

* :shirt:

* :book:
---
 connectors/src/connectors/slack/bot.ts        | 32 +++++-----
 .../src/connectors/slack/lib/slack_client.ts  | 63 +++++++++++++++++++
 connectors/src/lib/models/slack.ts            |  5 ++
 3 files changed, 86 insertions(+), 14 deletions(-)

diff --git a/connectors/src/connectors/slack/bot.ts b/connectors/src/connectors/slack/bot.ts
index 2968cb608239..983fbba39fc8 100644
--- a/connectors/src/connectors/slack/bot.ts
+++ b/connectors/src/connectors/slack/bot.ts
@@ -19,7 +19,11 @@ import { ConversationsRepliesResponse } from "@slack/web-api/dist/response/Conve
 import * as t from "io-ts";
 import jaroWinkler from "talisman/metrics/jaro-winkler";
 
-import { getSlackClient } from "@connectors/connectors/slack/lib/slack_client";
+import {
+  getSlackClient,
+  getSlackUserInfo,
+  isUserAllowedToUseChatbot,
+} from "@connectors/connectors/slack/lib/slack_client";
 import { Connector } from "@connectors/lib/models";
 import {
   SlackChannel,
@@ -82,7 +86,8 @@ export async function botAnswerMessageWithErrorHandling(
       slackUserId,
       slackMessageTs,
       slackThreadTs,
-      connector
+      connector,
+      slackConfig
     );
     if (res.isErr()) {
       logger.error(
@@ -147,7 +152,8 @@ async function botAnswerMessage(
   slackUserId: string,
   slackMessageTs: string,
   slackThreadTs: string | null,
-  connector: Connector
+  connector: Connector,
+  slackConfig: SlackConfiguration
 ): Promise<Result<AgentGenerationSuccessEvent, Error>> {
   let lastSlackChatBotMessage: SlackChatBotMessage | null = null;
   if (slackThreadTs) {
@@ -164,22 +170,20 @@ async function botAnswerMessage(
 
   // We start by retrieving the slack user info.
   const slackClient = await getSlackClient(connector.id);
-  const slackUserInfo = await slackClient.users.info({
-    user: slackUserId,
-  });
+  const slackUserInfo = await getSlackUserInfo(slackClient, slackUserId);
 
   if (!slackUserInfo.ok || !slackUserInfo.user) {
     throw new Error(`Failed to get user info: ${slackUserInfo.error}`);
   }
 
-  // We check that the user is not restricted or a stranger.
-  // See incident: https://dust4ai.slack.com/archives/C05B529FHV1/p1704799263814619
-  if (
-    slackUserInfo.user.profile?.team !== slackTeamId ||
-    slackUserInfo.user.is_restricted ||
-    slackUserInfo.user.is_ultra_restricted ||
-    slackUserInfo.user.is_stranger
-  ) {
+  const hasChatbotAccess = await isUserAllowedToUseChatbot(
+    slackClient,
+    slackUserInfo,
+    slackChannel,
+    slackTeamId,
+    slackConfig.whitelistedDomains
+  );
+  if (!hasChatbotAccess) {
     return new Err(
       new SlackExternalUserError(
         "Hi there. Sorry, but I can only answer to members of the workspace where I am installed."
diff --git a/connectors/src/connectors/slack/lib/slack_client.ts b/connectors/src/connectors/slack/lib/slack_client.ts
index e9dd80c6bf1c..88af45ab0bab 100644
--- a/connectors/src/connectors/slack/lib/slack_client.ts
+++ b/connectors/src/connectors/slack/lib/slack_client.ts
@@ -2,6 +2,7 @@ import { ModelId } from "@dust-tt/types";
 import {
   CodedError,
   ErrorCode,
+  UsersInfoResponse,
   WebAPIHTTPError,
   WebAPIPlatformError,
   WebClient,
@@ -101,6 +102,68 @@ export async function getSlackClient(
   return proxied;
 }
 
+export async function getSlackUserInfo(slackClient: WebClient, userId: string) {
+  return slackClient.users.info({
+    user: userId,
+  });
+}
+
+async function getSlackConversationInfo(
+  slackClient: WebClient,
+  channelId: string
+) {
+  return slackClient.conversations.info({ channel: channelId });
+}
+
+// Verify the Slack user is not an external guest to the workspace.
+// An exception is made for users from domains on the whitelist,
+// allowing them to interact with the bot in public channels.
+// See incident: https://dust4ai.slack.com/archives/C05B529FHV1/p1704799263814619
+export async function isUserAllowedToUseChatbot(
+  slackClient: WebClient,
+  slackUserInfo: UsersInfoResponse,
+  slackChanneId: string,
+  slackTeamId: string,
+  whitelistedDomains?: readonly string[]
+): Promise<boolean> {
+  if (!slackUserInfo.user) {
+    return false;
+  }
+
+  const {
+    is_restricted,
+    is_stranger: isStranger,
+    is_ultra_restricted,
+    profile,
+  } = slackUserInfo.user;
+
+  const isInWorkspace = profile?.team === slackTeamId;
+  if (!isInWorkspace) {
+    return false;
+  }
+
+  const isGuest = is_restricted || is_ultra_restricted;
+  const isExternal = isGuest || isStranger;
+
+  if (isExternal) {
+    const userDomain = profile?.email?.split("@")[1];
+    // Ensure the domain matches exactly.
+    const isWhitelistedDomain = userDomain
+      ? whitelistedDomains?.includes(userDomain) ?? false
+      : false;
+
+    const slackConversationInfo = await getSlackConversationInfo(
+      slackClient,
+      slackChanneId
+    );
+
+    const isChannelPublic = !slackConversationInfo.channel?.is_private;
+    return isChannelPublic && isWhitelistedDomain;
+  }
+
+  return true;
+}
+
 export async function getSlackAccessToken(
   nangoConnectionId: string
 ): Promise<string> {
diff --git a/connectors/src/lib/models/slack.ts b/connectors/src/lib/models/slack.ts
index ee68ac97ed42..1e2a2cc0979e 100644
--- a/connectors/src/lib/models/slack.ts
+++ b/connectors/src/lib/models/slack.ts
@@ -20,6 +20,7 @@ export class SlackConfiguration extends Model<
   declare slackTeamId: string;
   declare botEnabled: boolean;
   declare connectorId: ForeignKey<Connector["id"]>;
+  declare whitelistedDomains?: readonly string[];
 }
 SlackConfiguration.init(
   {
@@ -47,6 +48,10 @@ SlackConfiguration.init(
       allowNull: false,
       defaultValue: false,
     },
+    whitelistedDomains: {
+      type: DataTypes.ARRAY(DataTypes.STRING),
+      allowNull: true,
+    },
   },
   {
     sequelize: sequelize_conn,