refactor: move assets to their own model (#2301)

* refactor: move assets to their own model

* move asset privacy checks to the model

4 files changed, 274 insertions, 196 deletions

diff --git a/packages/api/routes/assets.ts b/packages/api/routes/assets.ts
index 50d11c47..e7d1c35f 100644
--- a/packages/api/routes/assets.ts
+++ b/packages/api/routes/assets.ts
@@ -1,11 +1,8 @@
 import { zValidator } from "@hono/zod-validator";
-import { TRPCError } from "@trpc/server";
-import { eq } from "drizzle-orm";
 import { Hono } from "hono";
 import { z } from "zod";
 
-import { assets } from "@karakeep/db/schema";
-import { BareBookmark } from "@karakeep/trpc/models/bookmarks";
+import { Asset } from "@karakeep/trpc/models/assets";
 
 import { authMiddleware } from "../middlewares/auth";
 import { serveAsset } from "../utils/assets";
@@ -37,39 +34,11 @@ const app = new Hono()
   )
   .get("/:assetId", async (c) => {
     const assetId = c.req.param("assetId");
-    const assetDb = await c.var.ctx.db.query.assets.findFirst({
-      where: eq(assets.id, assetId),
-      columns: {
-        id: true,
-        userId: true,
-        bookmarkId: true,
-      },
-    });
 
-    if (!assetDb) {
-      return c.json({ error: "Asset not found" }, { status: 404 });
-    }
+    const asset = await Asset.fromId(c.var.ctx, assetId);
+    await asset.ensureCanView();
 
-    // If asset is not attached to a bookmark yet, only owner can access it
-    if (!assetDb.bookmarkId) {
-      if (assetDb.userId !== c.var.ctx.user.id) {
-        return c.json({ error: "Asset not found" }, { status: 404 });
-      }
-      return await serveAsset(c, assetId, assetDb.userId);
-    }
-
-    // If asset is attached to a bookmark, check bookmark access permissions
-    try {
-      // This throws if the user doesn't have access to the bookmark
-      await BareBookmark.bareFromId(c.var.ctx, assetDb.bookmarkId);
-    } catch (e) {
-      if (e instanceof TRPCError && e.code === "FORBIDDEN") {
-        return c.json({ error: "Asset not found" }, { status: 404 });
-      }
-      throw e;
-    }
-
-    return await serveAsset(c, assetId, assetDb.userId);
+    return await serveAsset(c, assetId, asset.asset.userId);
   });
 
 export default app;
diff --git a/packages/trpc/models/assets.ts b/packages/trpc/models/assets.ts
new file mode 100644
index 00000000..98b89594
--- /dev/null
+++ b/packages/trpc/models/assets.ts
@@ -0,0 +1,252 @@
+import { TRPCError } from "@trpc/server";
+import { and, desc, eq, sql } from "drizzle-orm";
+import { z } from "zod";
+
+import { assets } from "@karakeep/db/schema";
+import { deleteAsset } from "@karakeep/shared/assetdb";
+import { zAssetTypesSchema } from "@karakeep/shared/types/bookmarks";
+
+import { AuthedContext } from "..";
+import {
+  isAllowedToAttachAsset,
+  isAllowedToDetachAsset,
+  mapDBAssetTypeToUserType,
+  mapSchemaAssetTypeToDB,
+} from "../lib/attachments";
+import { BareBookmark } from "./bookmarks";
+
+export class Asset {
+  constructor(
+    protected ctx: AuthedContext,
+    public asset: typeof assets.$inferSelect,
+  ) {}
+
+  static async fromId(ctx: AuthedContext, id: string): Promise<Asset> {
+    const assetdb = await ctx.db.query.assets.findFirst({
+      where: eq(assets.id, id),
+    });
+
+    if (!assetdb) {
+      throw new TRPCError({
+        code: "NOT_FOUND",
+        message: "Asset not found",
+      });
+    }
+
+    const asset = new Asset(ctx, assetdb);
+
+    if (!(await asset.canUserView())) {
+      throw new TRPCError({
+        code: "NOT_FOUND",
+        message: "Asset not found",
+      });
+    }
+
+    return asset;
+  }
+
+  static async list(
+    ctx: AuthedContext,
+    input: {
+      limit: number;
+      cursor: number | null;
+    },
+  ) {
+    const page = input.cursor ?? 1;
+    const [results, totalCount] = await Promise.all([
+      ctx.db
+        .select()
+        .from(assets)
+        .where(eq(assets.userId, ctx.user.id))
+        .orderBy(desc(assets.size))
+        .limit(input.limit)
+        .offset((page - 1) * input.limit),
+      ctx.db
+        .select({ count: sql<number>`count(*)` })
+        .from(assets)
+        .where(eq(assets.userId, ctx.user.id)),
+    ]);
+
+    return {
+      assets: results.map((a) => ({
+        ...a,
+        assetType: mapDBAssetTypeToUserType(a.assetType),
+      })),
+      nextCursor: page * input.limit < totalCount[0].count ? page + 1 : null,
+      totalCount: totalCount[0].count,
+    };
+  }
+
+  static async attachAsset(
+    ctx: AuthedContext,
+    input: {
+      bookmarkId: string;
+      asset: {
+        id: string;
+        assetType: z.infer<typeof zAssetTypesSchema>;
+      };
+    },
+  ) {
+    const [asset] = await Promise.all([
+      Asset.fromId(ctx, input.asset.id),
+      this.ensureBookmarkOwnership(ctx, input.bookmarkId),
+    ]);
+    asset.ensureOwnership();
+
+    if (!isAllowedToAttachAsset(input.asset.assetType)) {
+      throw new TRPCError({
+        code: "BAD_REQUEST",
+        message: "You can't attach this type of asset",
+      });
+    }
+
+    const [updatedAsset] = await ctx.db
+      .update(assets)
+      .set({
+        assetType: mapSchemaAssetTypeToDB(input.asset.assetType),
+        bookmarkId: input.bookmarkId,
+      })
+      .where(and(eq(assets.id, input.asset.id), eq(assets.userId, ctx.user.id)))
+      .returning();
+
+    return {
+      id: updatedAsset.id,
+      assetType: mapDBAssetTypeToUserType(updatedAsset.assetType),
+      fileName: updatedAsset.fileName,
+    };
+  }
+
+  static async replaceAsset(
+    ctx: AuthedContext,
+    input: {
+      bookmarkId: string;
+      oldAssetId: string;
+      newAssetId: string;
+    },
+  ) {
+    const [oldAsset, newAsset] = await Promise.all([
+      Asset.fromId(ctx, input.oldAssetId),
+      Asset.fromId(ctx, input.newAssetId),
+      this.ensureBookmarkOwnership(ctx, input.bookmarkId),
+    ]);
+    oldAsset.ensureOwnership();
+    newAsset.ensureOwnership();
+
+    if (
+      !isAllowedToAttachAsset(
+        mapDBAssetTypeToUserType(oldAsset.asset.assetType),
+      )
+    ) {
+      throw new TRPCError({
+        code: "BAD_REQUEST",
+        message: "You can't attach this type of asset",
+      });
+    }
+
+    await ctx.db.transaction(async (tx) => {
+      await tx.delete(assets).where(eq(assets.id, input.oldAssetId));
+      await tx
+        .update(assets)
+        .set({
+          bookmarkId: input.bookmarkId,
+          assetType: oldAsset.asset.assetType,
+        })
+        .where(eq(assets.id, input.newAssetId));
+    });
+
+    await deleteAsset({
+      userId: ctx.user.id,
+      assetId: input.oldAssetId,
+    }).catch(() => ({}));
+  }
+
+  static async detachAsset(
+    ctx: AuthedContext,
+    input: {
+      bookmarkId: string;
+      assetId: string;
+    },
+  ) {
+    const [asset] = await Promise.all([
+      Asset.fromId(ctx, input.assetId),
+      this.ensureBookmarkOwnership(ctx, input.bookmarkId),
+    ]);
+
+    if (
+      !isAllowedToDetachAsset(mapDBAssetTypeToUserType(asset.asset.assetType))
+    ) {
+      throw new TRPCError({
+        code: "BAD_REQUEST",
+        message: "You can't detach this type of asset",
+      });
+    }
+
+    const result = await ctx.db
+      .delete(assets)
+      .where(
+        and(
+          eq(assets.id, input.assetId),
+          eq(assets.bookmarkId, input.bookmarkId),
+        ),
+      );
+    if (result.changes == 0) {
+      throw new TRPCError({ code: "NOT_FOUND" });
+    }
+    await deleteAsset({ userId: ctx.user.id, assetId: input.assetId }).catch(
+      () => ({}),
+    );
+  }
+
+  private static async ensureBookmarkOwnership(
+    ctx: AuthedContext,
+    bookmarkId: string,
+  ) {
+    const bookmark = await BareBookmark.bareFromId(ctx, bookmarkId);
+    bookmark.ensureOwnership();
+  }
+
+  ensureOwnership() {
+    if (this.asset.userId != this.ctx.user.id) {
+      throw new TRPCError({
+        code: "FORBIDDEN",
+        message: "User is not allowed to access resource",
+      });
+    }
+  }
+
+  static async ensureOwnership(ctx: AuthedContext, assetId: string) {
+    return (await Asset.fromId(ctx, assetId)).ensureOwnership();
+  }
+
+  async canUserView(): Promise<boolean> {
+    // Asset owner can always view it
+    if (this.asset.userId === this.ctx.user.id) {
+      return true;
+    }
+
+    // If asset is attached to a bookmark, check bookmark access permissions
+    if (this.asset.bookmarkId) {
+      try {
+        // This throws if the user doesn't have access to the bookmark
+        await BareBookmark.bareFromId(this.ctx, this.asset.bookmarkId);
+        return true;
+      } catch (e) {
+        if (e instanceof TRPCError && e.code === "FORBIDDEN") {
+          return false;
+        }
+        throw e;
+      }
+    }
+
+    return false;
+  }
+
+  async ensureCanView() {
+    if (!(await this.canUserView())) {
+      throw new TRPCError({
+        code: "NOT_FOUND",
+        message: "Asset not found",
+      });
+    }
+  }
+}
diff --git a/packages/trpc/routers/assets.ts b/packages/trpc/routers/assets.ts
index 7be85446..c75f1e2e 100644
--- a/packages/trpc/routers/assets.ts
+++ b/packages/trpc/routers/assets.ts
@@ -1,57 +1,20 @@
-import { TRPCError } from "@trpc/server";
-import { and, desc, eq, sql } from "drizzle-orm";
 import { z } from "zod";
 
-import { assets, bookmarks } from "@karakeep/db/schema";
-import { deleteAsset } from "@karakeep/shared/assetdb";
 import {
   zAssetSchema,
   zAssetTypesSchema,
 } from "@karakeep/shared/types/bookmarks";
 
-import { authedProcedure, Context, router } from "../index";
-import {
-  isAllowedToAttachAsset,
-  isAllowedToDetachAsset,
-  mapDBAssetTypeToUserType,
-  mapSchemaAssetTypeToDB,
-} from "../lib/attachments";
+import { authedProcedure, router } from "../index";
+import { Asset } from "../models/assets";
 import { ensureBookmarkOwnership } from "./bookmarks";
 
-export const ensureAssetOwnership = async (opts: {
-  ctx: Context;
-  assetId: string;
-}) => {
-  const asset = await opts.ctx.db.query.assets.findFirst({
-    where: eq(bookmarks.id, opts.assetId),
-  });
-  if (!opts.ctx.user) {
-    throw new TRPCError({
-      code: "UNAUTHORIZED",
-      message: "User is not authorized",
-    });
-  }
-  if (!asset) {
-    throw new TRPCError({
-      code: "NOT_FOUND",
-      message: "Asset not found",
-    });
-  }
-  if (asset.userId != opts.ctx.user.id) {
-    throw new TRPCError({
-      code: "FORBIDDEN",
-      message: "User is not allowed to access resource",
-    });
-  }
-  return asset;
-};
-
 export const assetsAppRouter = router({
   list: authedProcedure
     .input(
       z.object({
         limit: z.number().min(1).max(100).default(20),
-        cursor: z.number().nullish(), // page number
+        cursor: z.number().nullish(),
       }),
     )
     .output(
@@ -71,29 +34,10 @@ export const assetsAppRouter = router({
       }),
     )
     .query(async ({ input, ctx }) => {
-      const page = input.cursor ?? 1;
-      const [results, totalCount] = await Promise.all([
-        ctx.db
-          .select()
-          .from(assets)
-          .where(eq(assets.userId, ctx.user.id))
-          .orderBy(desc(assets.size))
-          .limit(input.limit)
-          .offset((page - 1) * input.limit),
-        ctx.db
-          .select({ count: sql<number>`count(*)` })
-          .from(assets)
-          .where(eq(assets.userId, ctx.user.id)),
-      ]);
-
-      return {
-        assets: results.map((a) => ({
-          ...a,
-          assetType: mapDBAssetTypeToUserType(a.assetType),
-        })),
-        nextCursor: page * input.limit < totalCount[0].count ? page + 1 : null,
-        totalCount: totalCount[0].count,
-      };
+      return await Asset.list(ctx, {
+        limit: input.limit,
+        cursor: input.cursor ?? null,
+      });
     }),
   attachAsset: authedProcedure
     .input(
@@ -108,29 +52,7 @@ export const assetsAppRouter = router({
     .output(zAssetSchema)
     .use(ensureBookmarkOwnership)
     .mutation(async ({ input, ctx }) => {
-      await ensureAssetOwnership({ ctx, assetId: input.asset.id });
-      if (!isAllowedToAttachAsset(input.asset.assetType)) {
-        throw new TRPCError({
-          code: "BAD_REQUEST",
-          message: "You can't attach this type of asset",
-        });
-      }
-      const [updatedAsset] = await ctx.db
-        .update(assets)
-        .set({
-          assetType: mapSchemaAssetTypeToDB(input.asset.assetType),
-          bookmarkId: input.bookmarkId,
-        })
-        .where(
-          and(eq(assets.id, input.asset.id), eq(assets.userId, ctx.user.id)),
-        )
-        .returning();
-
-      return {
-        id: updatedAsset.id,
-        assetType: mapDBAssetTypeToUserType(updatedAsset.assetType),
-        fileName: updatedAsset.fileName,
-      };
+      return await Asset.attachAsset(ctx, input);
     }),
   replaceAsset: authedProcedure
     .input(
@@ -143,41 +65,7 @@ export const assetsAppRouter = router({
     .output(z.void())
     .use(ensureBookmarkOwnership)
     .mutation(async ({ input, ctx }) => {
-      await Promise.all([
-        ensureAssetOwnership({ ctx, assetId: input.oldAssetId }),
-        ensureAssetOwnership({ ctx, assetId: input.newAssetId }),
-      ]);
-      const [oldAsset] = await ctx.db
-        .select()
-        .from(assets)
-        .where(
-          and(eq(assets.id, input.oldAssetId), eq(assets.userId, ctx.user.id)),
-        )
-        .limit(1);
-      if (
-        !isAllowedToAttachAsset(mapDBAssetTypeToUserType(oldAsset.assetType))
-      ) {
-        throw new TRPCError({
-          code: "BAD_REQUEST",
-          message: "You can't attach this type of asset",
-        });
-      }
-
-      await ctx.db.transaction(async (tx) => {
-        await tx.delete(assets).where(eq(assets.id, input.oldAssetId));
-        await tx
-          .update(assets)
-          .set({
-            bookmarkId: input.bookmarkId,
-            assetType: oldAsset.assetType,
-          })
-          .where(eq(assets.id, input.newAssetId));
-      });
-
-      await deleteAsset({
-        userId: ctx.user.id,
-        assetId: input.oldAssetId,
-      }).catch(() => ({}));
+      await Asset.replaceAsset(ctx, input);
     }),
   detachAsset: authedProcedure
     .input(
@@ -189,34 +77,6 @@ export const assetsAppRouter = router({
     .output(z.void())
     .use(ensureBookmarkOwnership)
     .mutation(async ({ input, ctx }) => {
-      await ensureAssetOwnership({ ctx, assetId: input.assetId });
-      const [oldAsset] = await ctx.db
-        .select()
-        .from(assets)
-        .where(
-          and(eq(assets.id, input.assetId), eq(assets.userId, ctx.user.id)),
-        );
-      if (
-        !isAllowedToDetachAsset(mapDBAssetTypeToUserType(oldAsset.assetType))
-      ) {
-        throw new TRPCError({
-          code: "BAD_REQUEST",
-          message: "You can't deattach this type of asset",
-        });
-      }
-      const result = await ctx.db
-        .delete(assets)
-        .where(
-          and(
-            eq(assets.id, input.assetId),
-            eq(assets.bookmarkId, input.bookmarkId),
-          ),
-        );
-      if (result.changes == 0) {
-        throw new TRPCError({ code: "NOT_FOUND" });
-      }
-      await deleteAsset({ userId: ctx.user.id, assetId: input.assetId }).catch(
-        () => ({}),
-      );
+      await Asset.detachAsset(ctx, input);
     }),
 });
diff --git a/packages/trpc/routers/bookmarks.ts b/packages/trpc/routers/bookmarks.ts
index 65d401e2..a9d0df38 100644
--- a/packages/trpc/routers/bookmarks.ts
+++ b/packages/trpc/routers/bookmarks.ts
@@ -48,9 +48,9 @@ import { normalizeTagName } from "@karakeep/shared/utils/tag";
 import type { AuthedContext } from "../index";
 import { authedProcedure, createRateLimitMiddleware, router } from "../index";
 import { getBookmarkIdsFromMatcher } from "../lib/search";
+import { Asset } from "../models/assets";
 import { BareBookmark, Bookmark } from "../models/bookmarks";
 import { ImportSession } from "../models/importSessions";
-import { ensureAssetOwnership } from "./assets";
 
 export const ensureBookmarkOwnership = experimental_trpcMiddleware<{
   ctx: AuthedContext;
@@ -178,10 +178,7 @@ export const bookmarksAppRouter = router({
                   .returning()
               )[0];
               if (input.precrawledArchiveId) {
-                await ensureAssetOwnership({
-                  ctx,
-                  assetId: input.precrawledArchiveId,
-                });
+                await Asset.ensureOwnership(ctx, input.precrawledArchiveId);
                 await tx
                   .update(assets)
                   .set({
@@ -232,13 +229,13 @@ export const bookmarksAppRouter = router({
                   sourceUrl: null,
                 })
                 .returning();
-              const uploadedAsset = await ensureAssetOwnership({
-                ctx,
-                assetId: input.assetId,
-              });
+              const uploadedAsset = await Asset.fromId(ctx, input.assetId);
+              uploadedAsset.ensureOwnership();
               if (
-                !uploadedAsset.contentType ||
-                !SUPPORTED_BOOKMARK_ASSET_TYPES.has(uploadedAsset.contentType)
+                !uploadedAsset.asset.contentType ||
+                !SUPPORTED_BOOKMARK_ASSET_TYPES.has(
+                  uploadedAsset.asset.contentType,
+                )
               ) {
                 throw new TRPCError({
                   code: "BAD_REQUEST",