Associate slack user messages to the right user (#5187)

* Associate slack user messages to the right user * Add index + ✨ * ✨ * Better handle whitelisted email addresses. * Handle slack workflows unknown email address * Address comments from review * Add migration + fix script
dust-tt · May 21, 2024 · 06ee7b9 · 06ee7b9
1 parent ecc502c
commit 06ee7b9
Show file tree

Hide file tree

Showing 9 changed files with 151 additions and 16 deletions.
diff --git a/connectors/src/connectors/slack/bot.ts b/connectors/src/connectors/slack/bot.ts
@@ -437,6 +437,10 @@ async function botAnswerMessage(
     const mesasgeRes = await dustAPI.postUserMessage({
       conversationId: lastSlackChatBotMessage.conversationId,
       message: messageReqBody,
+      userEmailHeader:
+        slackChatBotMessage.slackEmail !== "unknown"
+          ? slackChatBotMessage.slackEmail
+          : undefined,
     });
     if (mesasgeRes.isErr()) {
       return new Err(new Error(mesasgeRes.error.message));
@@ -455,6 +459,10 @@ async function botAnswerMessage(
       visibility: "unlisted",
       message: messageReqBody,
       contentFragment: buildContentFragmentRes.value || undefined,
+      userEmailHeader:
+        slackChatBotMessage.slackEmail !== "unknown"
+          ? slackChatBotMessage.slackEmail
+          : undefined,
     });
     if (convRes.isErr()) {
       return new Err(new Error(convRes.error.message));

diff --git a/front/create_db_migration_file.sh b/front/create_db_migration_file.sh
@@ -56,7 +56,7 @@ diff --unified=0 --color=always main_output.txt current_output.txt
 
 # Run diff and extract only SQL statements, ensuring they end with a semicolon.
 echo "Running diff and extracting SQL statements..."
-diff_output=$(diff --unified=0 main_output.txt current_output.txt | awk '/^\+[^+]/ {print substr($0, 2)}') | sed 's/;*$/;/'
+diff_output=$(diff --unified=0 main_output.txt current_output.txt | awk '/^\+[^+]/ {print substr($0, 2)}' | sed 's/;*$/;/')
 if [ -n "$diff_output" ]; then
   echo "-- Migration created on $current_date" > diff_output.txt
   echo "$diff_output" >> diff_output.txt

diff --git a/front/lib/auth.ts b/front/lib/auth.ts
@@ -38,6 +38,7 @@ import { FREE_NO_PLAN_DATA } from "@app/lib/plans/free_plans";
 import { isUpgraded } from "@app/lib/plans/plan_codes";
 import { renderSubscriptionFromModels } from "@app/lib/plans/subscription";
 import { getTrialVersionForPlan, isTrial } from "@app/lib/plans/trial";
+import type { KeyAuthType } from "@app/lib/resources/key_resource";
 import { KeyResource } from "@app/lib/resources/key_resource";
 import { MembershipResource } from "@app/lib/resources/membership_resource";
 import { renderLightWorkspaceType } from "@app/lib/workspace";
@@ -60,11 +61,12 @@ const DUST_INTERNAL_EMAIL_REGEXP = /^[^@]+@dust\.tt$/;
  * workspace oriented. Use `getUserFromSession` if needed.
  */
 export class Authenticator {
-  _workspace: Workspace | null;
-  _user: User | null;
+  _flags: WhitelistableFeature[];
+  _key?: KeyAuthType;
   _role: RoleType;
   _subscription: SubscriptionType | null;
-  _flags: WhitelistableFeature[];
+  _user: User | null;
+  _workspace: Workspace | null;
 
   // Should only be called from the static methods below.
   constructor({
@@ -73,18 +75,21 @@ export class Authenticator {
     role,
     subscription,
     flags,
+    key,
   }: {
     workspace?: Workspace | null;
     user?: User | null;
     role: RoleType;
     subscription?: SubscriptionType | null;
     flags: WhitelistableFeature[];
+    key?: KeyAuthType;
   }) {
     this._workspace = workspace || null;
     this._user = user || null;
     this._role = role;
     this._subscription = subscription || null;
     this._flags = flags;
+    this._key = key;
   }
 
   /**
@@ -275,6 +280,7 @@ export class Authenticator {
         role,
         subscription,
         flags,
+        key: key.toAuthJSON(),
       }),
       keyWorkspaceId: keyWorkspace.sId,
     };
@@ -358,6 +364,62 @@ export class Authenticator {
     });
   }
 
+  /**
+   * Exchanges an Authenticator associated with a system key for one associated with a user.
+   *
+   * /!\ This function should only be used with Authenticators that are associated with a system key.
+   *
+   * @param auth
+   * @param param1
+   * @returns
+   */
+  async exchangeSystemKeyForUserAuthByEmail(
+    auth: Authenticator,
+    { userEmail }: { userEmail: string }
+  ): Promise<Authenticator | null> {
+    if (!auth.isSystemKey()) {
+      throw new Error("Provided authenticator does not have a system key.");
+    }
+
+    const owner = auth.workspace();
+    if (!owner) {
+      throw new Error("Workspace not found.");
+    }
+
+    const user = await User.findOne({
+      where: {
+        email: userEmail,
+      },
+    });
+    // If the user does not exist (e.g., whitelisted email addresses),
+    // simply ignore and return null.
+    if (!user) {
+      return null;
+    }
+
+    // Verify that the user has an active membership in the specified workspace.
+    const activeMembership =
+      await MembershipResource.getActiveMembershipOfUserInWorkspace({
+        user: renderUserType(user),
+        workspace: owner,
+      });
+    // If the user does not have an active membership in the workspace,
+    // simply ignore and return null.
+    if (!activeMembership) {
+      return null;
+    }
+
+    return new Authenticator({
+      flags: auth._flags,
+      key: auth._key,
+      // We limit scope to a user role.
+      role: "user",
+      user,
+      subscription: auth._subscription,
+      workspace: auth._workspace,
+    });
+  }
+
   role(): RoleType {
     return this._role;
   }
@@ -374,6 +436,10 @@ export class Authenticator {
     return isAdmin(this.workspace());
   }
 
+  isSystemKey(): boolean {
+    return !!this._key?.isSystem;
+  }
+
   workspace(): WorkspaceType | null {
     return this._workspace
       ? {

diff --git a/front/lib/models/user.ts b/front/lib/models/user.ts
@@ -103,6 +103,7 @@ User.init(
       { fields: ["provider", "providerId"] },
       { fields: ["auth0Sub"], unique: true, concurrently: true },
       { unique: true, fields: ["sId"] },
+      { fields: ["email"] },
     ],
   }
 );

diff --git a/front/lib/resources/key_resource.ts b/front/lib/resources/key_resource.ts
@@ -18,6 +18,12 @@ import { KeyModel } from "@app/lib/resources/storage/models/keys";
 import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types";
 import { new_id } from "@app/lib/utils";
 
+export interface KeyAuthType {
+  id: ModelId;
+  name: string | null;
+  isSystem: boolean;
+}
+
 // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging
 export interface KeyResource extends ReadonlyAttributesType<KeyModel> {}
 // eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging
@@ -167,6 +173,15 @@ export class KeyResource extends BaseResource<KeyModel> {
     };
   }
 
+  // Use to serialize a KeyResource in the Authenticator.
+  toAuthJSON(): KeyAuthType {
+    return {
+      id: this.id,
+      name: this.name,
+      isSystem: this.isSystem,
+    };
+  }
+
   get isActive() {
     return this.status === "active";
   }

diff --git a/front/migrations/db/migration_08.sql b/front/migrations/db/migration_08.sql
@@ -0,0 +1,2 @@
+-- Migration created on May 21, 2024
+CREATE INDEX "users_email" ON "users" ("email");
diff --git a/front/pages/api/v1/w/[wId]/assistant/conversations/[cId]/messages/index.ts b/front/pages/api/v1/w/[wId]/assistant/conversations/[cId]/messages/index.ts
@@ -27,10 +27,12 @@ async function handler(
     return apiError(req, res, keyRes.error);
   }
 
-  const { auth, keyWorkspaceId } = await Authenticator.fromKey(
+  const authenticator = await Authenticator.fromKey(
     keyRes.value,
     req.query.wId as string
   );
+  let { auth } = authenticator;
+  const { keyWorkspaceId } = authenticator;
 
   if (!auth.isBuilder() || keyWorkspaceId !== req.query.wId) {
     return apiError(req, res, {
@@ -71,6 +73,17 @@ async function handler(
 
       const { content, context, mentions, blocking } = bodyValidation.right;
 
+      // /!\ This is reserved for internal use!
+      // If the header "x-api-user-email" is present and valid,
+      // associate the message with the provided user email if it belongs to the same workspace.
+      const userEmailFromHeader = req.headers["x-api-user-email"];
+      if (typeof userEmailFromHeader === "string") {
+        auth =
+          (await auth.exchangeSystemKeyForUserAuthByEmail(auth, {
+            userEmail: userEmailFromHeader,
+          })) ?? auth;
+      }
+
       const messageRes = await postUserMessageWithPubSub(
         auth,
         {

diff --git a/front/pages/api/v1/w/[wId]/assistant/conversations/index.ts b/front/pages/api/v1/w/[wId]/assistant/conversations/index.ts
@@ -33,10 +33,12 @@ async function handler(
     return apiError(req, res, keyRes.error);
   }
 
-  const { auth, keyWorkspaceId } = await Authenticator.fromKey(
+  const authenticator = await Authenticator.fromKey(
     keyRes.value,
     req.query.wId as string
   );
+  let { auth } = authenticator;
+  const { keyWorkspaceId } = authenticator;
 
   if (!auth.isBuilder() || keyWorkspaceId !== req.query.wId) {
     return apiError(req, res, {
@@ -96,6 +98,17 @@ async function handler(
         }
       }
 
+      // /!\ This is reserved for internal use!
+      // If the header "x-api-user-email" is present and valid,
+      // associate the message with the provided user email if it belongs to the same workspace.
+      const userEmailFromHeader = req.headers["x-api-user-email"];
+      if (typeof userEmailFromHeader === "string") {
+        auth =
+          (await auth.exchangeSystemKeyForUserAuthByEmail(auth, {
+            userEmail: userEmailFromHeader,
+          })) ?? auth;
+      }
+
       let conversation = await createConversation(auth, {
         title,
         visibility,

diff --git a/types/src/front/lib/dust_api.ts b/types/src/front/lib/dust_api.ts
@@ -508,27 +508,36 @@ export class DustAPI {
   }
 
   // When creating a conversation with a user message, the API returns only after the user message
-  // was created (and if applicable the assocaited agent messages).
+  // was created (and if applicable the associated agent messages).
   async createConversation({
     title,
     visibility,
     message,
     contentFragment,
     blocking = false,
-  }: t.TypeOf<typeof PublicPostConversationsRequestBodySchema>): Promise<
+    userEmailHeader,
+  }: t.TypeOf<typeof PublicPostConversationsRequestBodySchema> & {
+    userEmailHeader?: string;
+  }): Promise<
     DustAPIResponse<{
       conversation: ConversationType;
       message: UserMessageType;
     }>
   > {
+    const headers: Record<string, string> = {
+      Authorization: `Bearer ${this._credentials.apiKey}`,
+      "Content-Type": "application/json",
+    };
+
+    if (userEmailHeader) {
+      headers["x-api-user-email"] = userEmailHeader;
+    }
+
     const res = await fetch(
       `${this.apiUrl()}/api/v1/w/${this.workspaceId()}/assistant/conversations`,
       {
         method: "POST",
-        headers: {
-          Authorization: `Bearer ${this._credentials.apiKey}`,
-          "Content-Type": "application/json",
-        },
+        headers,
         body: JSON.stringify({
           title,
           visibility,
@@ -545,18 +554,26 @@ export class DustAPI {
   async postUserMessage({
     conversationId,
     message,
+    userEmailHeader,
   }: {
     conversationId: string;
     message: t.TypeOf<typeof PublicPostMessagesRequestBodySchema>;
+    userEmailHeader?: string;
   }): Promise<DustAPIResponse<UserMessageType>> {
+    const headers: Record<string, string> = {
+      Authorization: `Bearer ${this._credentials.apiKey}`,
+      "Content-Type": "application/json",
+    };
+
+    if (userEmailHeader) {
+      headers["x-api-user-email"] = userEmailHeader;
+    }
+
     const res = await fetch(
       `${this.apiUrl()}/api/v1/w/${this.workspaceId()}/assistant/conversations/${conversationId}/messages`,
       {
         method: "POST",
-        headers: {
-          Authorization: `Bearer ${this._credentials.apiKey}`,
-          "Content-Type": "application/json",
-        },
+        headers,
         body: JSON.stringify({
           ...message,
         }),
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		-- Migration created on May 21, 2024
		CREATE INDEX "users_email" ON "users" ("email");