Extension: Add model/resource for configuration + poké plugin to blac…

…klist domains

front/admin/db.ts

@@ -53,6 +53,7 @@ import {
   TrackerDataSourceConfigurationModel,
   TrackerGenerationModel,
 } from "@app/lib/models/doc_tracker";
+import { ExtensionConfigurationModel } from "@app/lib/models/extension";
 import { FeatureFlag } from "@app/lib/models/feature_flag";
 import { Plan, Subscription } from "@app/lib/models/plan";
 import {
@@ -133,6 +134,8 @@ async function main() {
   await TrackerDataSourceConfigurationModel.sync({ alter: true });
   await TrackerGenerationModel.sync({ alter: true });
 
+  await ExtensionConfigurationModel.sync({ alter: true });
+
   await Plan.sync({ alter: true });
   await Subscription.sync({ alter: true });
   await TemplateModel.sync({ alter: true });

front/lib/api/poke/plugins/workspaces/extension_blacklist_domains.ts

@@ -0,0 +1,79 @@
+import { Err, Ok } from "@dust-tt/types";
+
+import { createPlugin } from "@app/lib/api/poke/types";
+import { updateExtensionConfiguration } from "@app/lib/api/workspace";
+
+export const extensionBlacklistDomains = createPlugin(
+  {
+    id: "extension-blacklist-domains",
+    name: "Extension Blacklist Domains",
+    description: "Update the list of blacklisted domains for the extension",
+    resourceTypes: ["workspaces"],
+    args: {
+      domains: {
+        type: "string",
+        label: "Blacklisted domains",
+        description:
+          "Comma-separated list of domains to blacklist for the extension.",
+      },
+    },
+  },
+  async (auth, resourceId, args) => {
+    // Split by comma and remove any empty strings, and check domains are valid
+    const domains = args.domains
+      ? args.domains
+          .split(",")
+          .map((d) => d.trim())
+          .filter((d) => d)
+      : [];
+
+    if (!areDomainsValid(domains)) {
+      return new Err(
+        new Error(
+          "One or more domains are invalid. Please check the domain format."
+        )
+      );
+    }
+
+    const res = await updateExtensionConfiguration(auth, domains);
+    if (res.isErr()) {
+      return res;
+    }
+
+    return new Ok({
+      display: "text",
+      value: `Blacklisted domains updated.`,
+    });
+  }
+);
+
+function areDomainsValid(domains: string[]): boolean {
+  if (domains.length === 0) {
+    return true; // Empty domains array is valid
+  }
+
+  // Regular expression for domain validation
+  // - Starts with alphanumeric or hyphen
+  // - Can contain alphanumeric, hyphens
+  // - Must have at least one dot
+  // - TLD must be at least 2 characters
+  // - Cannot start or end with hyphen
+  // - Cannot have consecutive hyphens
+  const domainRegex =
+    /^(?!-)[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}(?!-)$/;
+
+  return domains.every((domain) => {
+    if (domain.length > 253) {
+      return false;
+    }
+    if (!domainRegex.test(domain)) {
+      return false;
+    }
+    const labels = domain.split(".");
+    if (labels.some((label) => label.length > 63)) {
+      return false;
+    }
+
+    return true;
+  });
+}

front/lib/api/poke/plugins/workspaces/index.ts

@@ -1,6 +1,7 @@
 export * from "./create_space";
 export * from "./disable_sso_enforcement";
 export * from "./extend_trial";
+export * from "./extension_blacklist_domains";
 export * from "./invite_user";
 export * from "./rename_workspace";
 export * from "./reset_message_rate_limit";

front/lib/api/workspace.ts

@@ -17,6 +17,7 @@ import type { Authenticator } from "@app/lib/auth";
 import { Subscription } from "@app/lib/models/plan";
 import { Workspace, WorkspaceHasDomain } from "@app/lib/models/workspace";
 import { getStripeSubscription } from "@app/lib/plans/stripe";
+import { ExtensionConfigurationResource } from "@app/lib/resources/extension";
 import { MembershipResource } from "@app/lib/resources/membership_resource";
 import { UserResource } from "@app/lib/resources/user_resource";
 import { renderLightWorkspaceType } from "@app/lib/workspace";
@@ -380,3 +381,21 @@ export async function disableSSOEnforcement(
 
   return new Ok(undefined);
 }
+
+export async function updateExtensionConfiguration(
+  auth: Authenticator,
+  blacklistedDomains: string[]
+): Promise<Result<void, Error>> {
+  const config = await ExtensionConfigurationResource.fetchForWorkspace(auth);
+
+  if (config) {
+    await config.updateBlacklistedDomains(auth, { blacklistedDomains });
+  } else {
+    await ExtensionConfigurationResource.makeNew(
+      { blacklistedDomains },
+      auth.getNonNullableWorkspace().id
+    );
+  }
+
+  return new Ok(undefined);
+}

front/lib/models/extension.ts

@@ -0,0 +1,49 @@
+import type { CreationOptional, ForeignKey, NonAttribute } from "sequelize";
+import { DataTypes } from "sequelize";
+
+import { Workspace } from "@app/lib/models/workspace";
+import { frontSequelize } from "@app/lib/resources/storage";
+import { BaseModel } from "@app/lib/resources/storage/wrappers";
+
+export class ExtensionConfigurationModel extends BaseModel<ExtensionConfigurationModel> {
+  declare createdAt: CreationOptional<Date>;
+  declare updatedAt: CreationOptional<Date>;
+
+  declare blacklistedDomains: string[];
+
+  declare workspaceId: ForeignKey<Workspace["id"]>;
+  declare workspace: NonAttribute<Workspace>;
+}
+ExtensionConfigurationModel.init(
+  {
+    createdAt: {
+      type: DataTypes.DATE,
+      allowNull: false,
+      defaultValue: DataTypes.NOW,
+    },
+    updatedAt: {
+      type: DataTypes.DATE,
+      allowNull: false,
+      defaultValue: DataTypes.NOW,
+    },
+    blacklistedDomains: {
+      type: DataTypes.ARRAY(DataTypes.STRING),
+      allowNull: false,
+      defaultValue: [],
+    },
+  },
+  {
+    modelName: "extension_configuration",
+    sequelize: frontSequelize,
+    indexes: [{ unique: true, fields: ["workspaceId"] }],
+  }
+);
+
+Workspace.hasOne(ExtensionConfigurationModel, {
+  foreignKey: { allowNull: false },
+  onDelete: "RESTRICT",
+});
+
+ExtensionConfigurationModel.belongsTo(Workspace, {
+  foreignKey: { allowNull: false },
+});

front/lib/resources/extension.ts

@@ -0,0 +1,126 @@
+import type {
+  ExtensionConfigurationType,
+  ModelId,
+  Result,
+} from "@dust-tt/types";
+import { Err, Ok } from "@dust-tt/types";
+import type {
+  Attributes,
+  CreationAttributes,
+  ModelStatic,
+  Transaction,
+} from "sequelize";
+
+import type { Authenticator } from "@app/lib/auth";
+import { ExtensionConfigurationModel } from "@app/lib/models/extension";
+import { BaseResource } from "@app/lib/resources/base_resource";
+import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types";
+import { makeSId } from "@app/lib/resources/string_ids";
+
+// Attributes are marked as read-only to reflect the stateless nature of our Resource.
+// This design will be moved up to BaseResource once we transition away from Sequelize.
+// eslint-disable-next-line @typescript-eslint/no-empty-interface, @typescript-eslint/no-unsafe-declaration-merging
+export interface ExtensionConfigurationResource
+  extends ReadonlyAttributesType<ExtensionConfigurationModel> {}
+// eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging
+export class ExtensionConfigurationResource extends BaseResource<ExtensionConfigurationModel> {
+  static model: ModelStatic<ExtensionConfigurationModel> =
+    ExtensionConfigurationModel;
+
+  constructor(
+    model: ModelStatic<ExtensionConfigurationModel>,
+    blob: Attributes<ExtensionConfigurationModel>
+  ) {
+    super(ExtensionConfigurationModel, blob);
+  }
+
+  get sId(): string {
+    return ExtensionConfigurationResource.modelIdToSId({
+      id: this.id,
+      workspaceId: this.workspaceId,
+    });
+  }
+
+  static modelIdToSId({
+    id,
+    workspaceId,
+  }: {
+    id: ModelId;
+    workspaceId: ModelId;
+  }): string {
+    return makeSId("extension", {
+      id,
+      workspaceId,
+    });
+  }
+
+  static async makeNew(
+    blob: Omit<CreationAttributes<ExtensionConfigurationModel>, "workspaceId">,
+    workspaceId: ModelId
+  ) {
+    const config = await ExtensionConfigurationModel.create({
+      ...blob,
+      workspaceId,
+    });
+
+    return new this(ExtensionConfigurationModel, config.get());
+  }
+
+  async delete(
+    auth: Authenticator,
+    { transaction }: { transaction?: Transaction }
+  ): Promise<Result<undefined, Error>> {
+    try {
+      await this.model.destroy({
+        where: {
+          id: this.id,
+        },
+        transaction,
+      });
+
+      return new Ok(undefined);
+    } catch (err) {
+      return new Err(err as Error);
+    }
+  }
+
+  static async fetchForWorkspace(
+    auth: Authenticator
+  ): Promise<ExtensionConfigurationResource | null> {
+    const workspaceId = auth.getNonNullableWorkspace().id;
+    const config = await this.model.findOne({
+      where: {
+        workspaceId,
+      },
+    });
+
+    return config ? new this(ExtensionConfigurationModel, config.get()) : null;
+  }
+
+  async updateBlacklistedDomains(
+    auth: Authenticator,
+    {
+      blacklistedDomains,
+    }: {
+      blacklistedDomains: string[];
+    }
+  ) {
+    if (this.workspaceId !== auth.getNonNullableWorkspace().id) {
+      throw new Error(
+        "Can't update extension configuration for another workspace."
+      );
+    }
+
+    await this.update({
+      blacklistedDomains,
+    });
+  }
+
+  toJSON(): ExtensionConfigurationType {
+    return {
+      id: this.id,
+      sId: this.sId,
+      blacklistedDomains: this.blacklistedDomains,
+    };
+  }
+}

front/lib/resources/string_ids.ts

@@ -27,6 +27,7 @@ const RESOURCES_PREFIX = {
   data_source_view: "dsv",
   tracker: "trk",
   template: "tpl",
+  extension: "ext",
 };
 
 export const CROSS_WORKSPACE_RESOURCES_WORKSPACE_ID: ModelId = 0;

front/migrations/db/migration_137.sql

@@ -0,0 +1,10 @@
+-- Migration created on Dec 30, 2024
+CREATE TABLE IF NOT EXISTS "extension_configurations" (
+    "id"  BIGSERIAL ,
+    "createdAt" TIMESTAMP WITH TIME ZONE NOT NULL,
+    "updatedAt" TIMESTAMP WITH TIME ZONE NOT NULL,
+    "blacklistedDomains" VARCHAR(255)[] NOT NULL DEFAULT ARRAY[]::VARCHAR(255)[],
+    "workspaceId" BIGINT NOT NULL REFERENCES "workspaces" ("id") ON DELETE RESTRICT ON UPDATE CASCADE,
+    PRIMARY KEY ("id")
+);
+CREATE UNIQUE INDEX "extension_configurations_workspace_id" ON "extension_configurations" ("workspaceId");

types/src/front/extension.ts

@@ -0,0 +1,7 @@
+import { ModelId } from "../shared/model_id";
+
+export type ExtensionConfigurationType = {
+  id: ModelId;
+  sId: string;
+  blacklistedDomains: string[];
+};

types/src/index.ts

@@ -44,6 +44,7 @@ export * from "./front/data_source_view";
 export * from "./front/dataset";
 export * from "./front/document";
 export * from "./front/dust_app_secret";
+export * from "./front/extension";
 export * from "./front/files";
 export * from "./front/groups";
 export * from "./front/key";