feat(folders): adopted new strategy for rollback on folders

backend/src/controllers/v1/secretsFolderController.ts

@@ -51,6 +51,10 @@ export const createFolder = async (req: Request, res: Response) => {
       nodes: folder.nodes,
     });
     await folderVersion.save();
+    await EESecretService.takeSecretSnapshot({
+      workspaceId,
+      environment,
+    });
     return res.json({ folder: { id, name: folderName } });
   }
 
@@ -66,6 +70,12 @@ export const createFolder = async (req: Request, res: Response) => {
   });
   await folderVersion.save();
 
+  await EESecretService.takeSecretSnapshot({
+    workspaceId,
+    environment,
+    folderId: parentFolderId,
+  });
+
   return res.json({ folder });
 };
 
@@ -96,7 +106,6 @@ export const updateFolderById = async (req: Request, res: Response) => {
 
   parentFolder.version += 1;
   folder.name = name;
-  folder.version += 1;
 
   await Folder.findByIdAndUpdate(folders._id, folders);
   const folderVersion = new FolderVersion({
@@ -106,19 +115,11 @@ export const updateFolderById = async (req: Request, res: Response) => {
   });
   await folderVersion.save();
 
-  // decision to take snapshot ot not
+  await EESecretService.takeSecretSnapshot({
-  const secretCount = await Secret.count({
+    workspaceId,
-    workspace: workspaceId,
     environment,
-    folder: folderId,
+    folderId: parentFolder.id,
   });
-  if (secretCount || folder.children.length) {
-    await EESecretService.takeSecretSnapshot({
-      workspaceId,
-      environment,
-      folderId: parentFolder.id,
-    });
-  }
 
   return res.json({
     message: "Successfully updated folder",
@@ -148,23 +149,9 @@ export const deleteFolder = async (req: Request, res: Response) => {
   }
   const { deletedNode: delFolder, parent: parentFolder } = delOp;
 
+  parentFolder.version += 1;
   const delFolderIds = getAllFolderIds(delFolder);
 
-  // take a pre-delete snapshot if there are secrets inside the folders
-  const secretCount = await Secret.count({
-    workspace: workspaceId,
-    environment,
-    folder: folderId,
-  });
-
-  if (secretCount || delFolder.children.length) {
-    await EESecretService.takeSecretSnapshot({
-      workspaceId,
-      environment,
-      folderId: parentFolder.id,
-    });
-  }
-
   await Folder.findByIdAndUpdate(folders._id, folders);
   const folderVersion = new FolderVersion({
     workspace: workspaceId,
@@ -180,7 +167,6 @@ export const deleteFolder = async (req: Request, res: Response) => {
     });
   }
 
-  // post delete snapshot
   await EESecretService.takeSecretSnapshot({
     workspaceId,
     environment,

backend/src/ee/controllers/v1/workspaceController.ts

@@ -1,6 +1,6 @@
 import { Request, Response } from "express";
 import * as Sentry from "@sentry/node";
-import { Types } from "mongoose";
+import { PipelineStage, Types } from "mongoose";
 import { Secret } from "../../../models";
 import {
   SecretSnapshot,
@@ -12,11 +12,8 @@ import {
 } from "../../models";
 import { EESecretService } from "../../services";
 import { getLatestSecretVersionIds } from "../../helpers/secretVersion";
-import Folder from "../../../models/folder";
+import Folder, { TFolderSchema } from "../../../models/folder";
-import {
+import { searchByFolderId } from "../../../services/FolderService";
-  getAllFolderIds,
-  searchByFolderId,
-} from "../../../services/FolderService";
 
 /**
  * Return secret snapshots for workspace with id [workspaceId]
@@ -214,17 +211,135 @@ export const rollbackWorkspaceSecretSnapshot = async (
 
     if (!secretSnapshot) throw new Error("Failed to find secret snapshot");
 
-    // TODO: fix any
+    const snapshotFolderTree = secretSnapshot.folderVersion;
-    const oldSecretVersionsObj: any = secretSnapshot.secretVersions.reduce(
+    const latestFolderTree = await Folder.findOne({
-      (accumulator, s) => ({
+      workspace: workspaceId,
-        ...accumulator,
+      environment,
-        [`${s.secret.toString()}`]: s,
+    });
-      }),
-      {}
-    );
 
+    const latestFolderVersion = await FolderVersion.findOne({
+      environment,
+      workspace: workspaceId,
+      "nodes.id": folderId,
+    }).sort({ "nodes.version": -1 });
+
+    const oldSecretVersionsObj: Record<string, ISecretVersion> = {};
+    const secretIds: Types.ObjectId[] = [];
+    const folderIds: string[] = [folderId];
+
+    secretSnapshot.secretVersions.forEach((snapSecVer) => {
+      oldSecretVersionsObj[snapSecVer.secret.toString()] = snapSecVer;
+      secretIds.push(snapSecVer.secret);
+    });
+
+    // the parent node from current latest one
+    // this will be modified according to the snapshot and latest snapshots
+    const newFolderTree =
+      latestFolderTree && searchByFolderId(latestFolderTree.nodes, folderId);
+
+    if (newFolderTree) {
+      newFolderTree.children = snapshotFolderTree?.nodes?.children || [];
+      const queue = [newFolderTree];
+      // a bfs algorithm in which we take the latest snapshots of all the folders in a level
+      while (queue.length) {
+        const groupByFolderId: Record<string, TFolderSchema> = {};
+        // the original queue is popped out completely to get what ever in a level
+        // subqueue is filled with all the children thus next level folders
+        // subQueue will then be transfered to the oriinal queue
+        const subQueue: TFolderSchema[] = [];
+        // get everything inside a level
+        while (queue.length) {
+          const folder = queue.pop() as TFolderSchema;
+          folder.children.forEach((el) => {
+            folderIds.push(el.id); // push ids and data into queu
+            subQueue.push(el);
+            // to modify the original tree very fast we keep a reference object
+            // key with folder id and pointing to the various nodes
+            groupByFolderId[el.id] = el;
+          });
+        }
+        // get latest snapshots of all the folder
+        const matchWsFoldersPipeline = {
+          $match: {
+            workspace: new Types.ObjectId(workspaceId),
+            environment,
+            folderId: {
+              $in: Object.keys(groupByFolderId),
+            },
+          },
+        };
+        const sortByFolderIdAndVersion: PipelineStage = {
+          $sort: { folderId: 1, version: -1 },
+        };
+        const pickLatestVersionOfEachFolder = {
+          $group: {
+            _id: "$folderId",
+            latestVersion: { $first: "$version" },
+            doc: {
+              $first: "$$ROOT",
+            },
+          },
+        };
+        const populateSecVersion = {
+          $lookup: {
+            from: SecretVersion.collection.name,
+            localField: "doc.secretVersions",
+            foreignField: "_id",
+            as: "doc.secretVersions",
+          },
+        };
+        const populateFolderVersion = {
+          $lookup: {
+            from: FolderVersion.collection.name,
+            localField: "doc.folderVersion",
+            foreignField: "_id",
+            as: "doc.folderVersion",
+          },
+        };
+        const unwindFolderVerField = {
+          $unwind: {
+            path: "$doc.folderVersion",
+            preserveNullAndEmptyArrays: true,
+          },
+        };
+        const latestSnapshotsByFolders: Array<{ doc: typeof secretSnapshot }> =
+          await SecretSnapshot.aggregate([
+            matchWsFoldersPipeline,
+            sortByFolderIdAndVersion,
+            pickLatestVersionOfEachFolder,
+            populateSecVersion,
+            populateFolderVersion,
+            unwindFolderVerField,
+          ]);
+
+        // recursive snapshotting each level
+        latestSnapshotsByFolders.forEach((snap) => {
+          // mutate the folder tree to update the nodes to the latest version tree
+          // we are reconstructing the folder tree by latest snapshots here
+          if (groupByFolderId[snap.doc.folderId]) {
+            groupByFolderId[snap.doc.folderId].children =
+              snap.doc?.folderVersion?.nodes?.children || [];
+          }
+
+          // push all children of next level snapshots
+          if (snap.doc.folderVersion?.nodes?.children) {
+            queue.push(...snap.doc.folderVersion.nodes.children);
+          }
+
+          snap.doc.secretVersions.forEach((snapSecVer) => {
+            // record all the secrets
+            oldSecretVersionsObj[snapSecVer.secret.toString()] = snapSecVer;
+            secretIds.push(snapSecVer.secret);
+          });
+        });
+
+        queue.push(...subQueue);
+      }
+    }
+
+    // TODO: fix any
     const latestSecretVersionIds = await getLatestSecretVersionIds({
-      secretIds: secretSnapshot.secretVersions.map((sv) => sv.secret),
+      secretIds,
     });
 
     // TODO: fix any
@@ -245,27 +360,6 @@ export const rollbackWorkspaceSecretSnapshot = async (
       {}
     );
 
-    let folderIds: string[] = [];
-
-    const folders = await Folder.findOne({
-      workspace: workspaceId,
-      environment,
-    }).lean();
-    const latestFolderVersion = await FolderVersion.findOne({
-      environment,
-      workspace: workspaceId,
-      "nodes.id": folderId,
-    }).sort({ "nodes.version": -1 });
-
-    if (folders && folderId) {
-      const folder = searchByFolderId(folders.nodes, folderId);
-      if (folder) {
-        folderIds = getAllFolderIds(folder).map(({ id }) => id);
-        folder.children = secretSnapshot?.folderVersion?.nodes?.children || [];
-        folder.version = (latestFolderVersion?.nodes?.version || 0) + 1;
-      }
-    }
-
     const secDelQuery: Record<string, unknown> = {
       workspace: workspaceId,
       environment,
@@ -283,9 +377,9 @@ export const rollbackWorkspaceSecretSnapshot = async (
 
     // add secrets
     secrets = await Secret.insertMany(
-      secretSnapshot.secretVersions.map((sv) => {
+      Object.keys(oldSecretVersionsObj).map((sv) => {
-        const secretId = sv.secret;
         const {
+          secret: secretId,
           workspace,
           type,
           user,
@@ -294,16 +388,14 @@ export const rollbackWorkspaceSecretSnapshot = async (
           secretKeyCiphertext,
           secretKeyIV,
           secretKeyTag,
-          secretKeyHash,
           secretValueCiphertext,
           secretValueIV,
           secretValueTag,
-          secretValueHash,
           createdAt,
           algorithm,
           keyEncoding,
           folder: secFolderId,
-        } = oldSecretVersionsObj[secretId.toString()];
+        } = oldSecretVersionsObj[sv];
 
         return {
           _id: secretId,
@@ -316,11 +408,9 @@ export const rollbackWorkspaceSecretSnapshot = async (
           secretKeyCiphertext,
           secretKeyIV,
           secretKeyTag,
-          secretKeyHash,
           secretValueCiphertext,
           secretValueIV,
           secretValueTag,
-          secretValueHash,
           secretCommentCiphertext: "",
           secretCommentIV: "",
           secretCommentTag: "",
@@ -346,11 +436,9 @@ export const rollbackWorkspaceSecretSnapshot = async (
           secretKeyCiphertext,
           secretKeyIV,
           secretKeyTag,
-          secretKeyHash,
           secretValueCiphertext,
           secretValueIV,
           secretValueTag,
-          secretValueHash,
           algorithm,
           keyEncoding,
           folder: secFolderId,
@@ -367,11 +455,9 @@ export const rollbackWorkspaceSecretSnapshot = async (
           secretKeyCiphertext,
           secretKeyIV,
           secretKeyTag,
-          secretKeyHash,
           secretValueCiphertext,
           secretValueIV,
           secretValueTag,
-          secretValueHash,
           algorithm,
           keyEncoding,
           folder: secFolderId,
@@ -379,17 +465,18 @@ export const rollbackWorkspaceSecretSnapshot = async (
       )
     );
 
-    if (folders) {
+    if (newFolderTree && latestFolderTree) {
-      const newFolder = new Folder(folders);
+      // save the updated folder tree to the present one
-      newFolder._id = new Types.ObjectId();
+      newFolderTree.version = (latestFolderVersion?.nodes?.version || 0) + 1;
-      newFolder.isNew = true;
+      latestFolderTree._id = new Types.ObjectId();
-      // when there is no
+      latestFolderTree.isNew = true;
-      await newFolder.save();
+      await latestFolderTree.save();
+
       // create new folder version
       const newFolderVersion = new FolderVersion({
         workspace: workspaceId,
         environment,
-        nodes: newFolder.nodes,
+        nodes: newFolderTree,
       });
       await newFolderVersion.save();
     }
@@ -398,7 +485,9 @@ export const rollbackWorkspaceSecretSnapshot = async (
     await SecretVersion.updateMany(
       {
         secret: {
-          $in: secretSnapshot.secretVersions.map((sv) => sv.secret),
+          $in: Object.keys(oldSecretVersionsObj).map(
+            (sv) => oldSecretVersionsObj[sv].secret
+          ),
         },
       },
       {

backend/src/ee/helpers/secret.ts

@@ -1,10 +1,5 @@
 import { Types } from "mongoose";
 import { Secret, ISecret } from "../../models";
-import Folder from "../../models/folder";
-import {
-  getAllFolderIds,
-  searchByFolderId,
-} from "../../services/FolderService";
 import {
   SecretSnapshot,
   SecretVersion,
@@ -30,25 +25,17 @@ const takeSecretSnapshotHelper = async ({
   folderId?: string;
 }) => {
   // get all folder ids
-  let folderIds: string[] = [];
+  const secretIds = (
-  const folders = await Folder.findOne({
+    await Secret.find(
-    workspace: workspaceId,
+      {
-    environment,
+        workspace: workspaceId,
-  }).lean();
+        environment,
-  if (folders && folderId) {
+        folder: folderId,
-    const folder = searchByFolderId(folders.nodes, folderId);
+        // undefined means root thus collect all secrets
-    if (folder) folderIds = getAllFolderIds(folder).map(({ id }) => id);
+      },
-  }
+      "_id"
-
+    ).lean()
-  const secQuery: Record<string, unknown> = {
+  ).map((s) => s._id);
-    workspace: workspaceId,
-    environment,
-    // undefined means root thus collect all secrets
-  };
-  if (folderId !== "root") secQuery.folder = { $in: folderIds };
-  const secretIds = (await Secret.find(secQuery, "_id").lean()).map(
-    (s) => s._id
-  );
 
   const latestSecretVersions = (
     await SecretVersion.aggregate([
@@ -76,7 +63,7 @@ const takeSecretSnapshotHelper = async ({
   const latestFolderVersion = await FolderVersion.findOne({
     environment,
     workspace: workspaceId,
-    "nodes.id": folderId || "root",
+    "nodes.id": folderId,
   }).sort({ "nodes.version": -1 });
 
   const latestSecretSnapshot = await SecretSnapshot.findOne({
@@ -132,5 +119,5 @@ const markDeletedSecretVersionsHelper = async ({
 export {
   takeSecretSnapshotHelper,
   addSecretVersionsHelper,
-  markDeletedSecretVersionsHelper
+  markDeletedSecretVersionsHelper,
 };

backend/src/ee/models/secretVersion.ts

@@ -25,6 +25,7 @@ export interface ISecretVersion {
   secretValueTag: string;
   algorithm: "aes-256-gcm";
   keyEncoding: "utf8" | "base64";
+  createdAt: string;
   folder?: string;
 }
 
@@ -109,7 +110,7 @@ const secretVersionSchema = new Schema<ISecretVersion>(
     },
     folder: {
       type: String,
-      required: false,
+      required: true,
     },
   },
   {

backend/src/services/FolderService.ts

@@ -33,6 +33,18 @@ export const searchByFolderId = (
   }
 };
 
+export const folderBfsTraversal = async (
+  root: TFolderSchema,
+  callback: (data: TFolderSchema) => void | Promise<void>
+) => {
+  const queue = [root];
+  while (queue.length) {
+    const folder = queue.pop() as TFolderSchema;
+    await callback(folder);
+    queue.push(...folder.children);
+  }
+};
+
 // bfs and then append to the folder
 const appendChild = (folders: TFolderSchema, folderName: string) => {
   const folder = folders.children.find(({ name }) => name === folderName);
@@ -85,7 +97,6 @@ export const deleteFolderById = (folders: TFolderSchema, folderId: string) => {
   const queue = [folders];
   while (queue.length) {
     const folder = queue.pop() as TFolderSchema;
-    folder.version += 1;
     const index = folder.children.findIndex(({ id }) => folderId === id);
     if (index !== -1) {
       const deletedFolder = folder.children.splice(index, 1);

frontend/src/components/navigation/NavHeader.tsx

@@ -1,6 +1,6 @@
 import Link from 'next/link';
 import { useRouter } from 'next/router';
-import { faAngleRight, faHome } from '@fortawesome/free-solid-svg-icons';
+import { faAngleRight } from '@fortawesome/free-solid-svg-icons';
 import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
 
 import { useOrganization, useWorkspace } from '@app/context';
@@ -114,7 +114,7 @@ export default function NavHeader({
               ) : (
                 <Link passHref legacyBehavior href={{ pathname: '/dashboard/[id]', query }}>
                   <a className="text-sm font-semibold text-primary/80 hover:text-primary">
-                    {name === 'root' ? <FontAwesomeIcon icon={faHome} /> : name}
+                    {name === 'root' ? '/' : name}
                   </a>
                 </Link>
               )}

frontend/src/hooks/api/secretSnapshots/queries.tsx

@@ -180,6 +180,7 @@ export const usePerformSecretRollback = () => {
       return data;
     },
     onSuccess: (_, { workspaceId, environment, folderId }) => {
+      queryClient.invalidateQueries([{ workspaceId, environment }, 'secrets']);
       queryClient.invalidateQueries(secretSnapshotKeys.list(workspaceId, environment, folderId));
       queryClient.invalidateQueries(secretSnapshotKeys.count(workspaceId, environment, folderId));
     }