Merge pull request dubinc#2337 from dubinc/fix-upload-security

Add base64 image validation

Author: steven-tey

apps/web/app/api/user/route.ts

@@ -2,6 +2,7 @@ import { DubApiError } from "@/lib/api/errors";
 import { hashToken, withSession } from "@/lib/auth";
 import { storage } from "@/lib/storage";
 import { ratelimit, redis } from "@/lib/upstash";
+import { base64ImageSchema } from "@/lib/zod/schemas/misc";
 import { sendEmail } from "@dub/email";
 import { unsubscribe } from "@dub/email/resend/unsubscribe";
 import ConfirmEmailChange from "@dub/email/templates/confirm-email-change";
@@ -15,7 +16,7 @@ import { z } from "zod";
 const updateUserSchema = z.object({
   name: z.preprocess(trim, z.string().min(1).max(64)).optional(),
   email: z.preprocess(trim, z.string().email()).optional(),
-  image: z.string().url().optional(),
+  image: base64ImageSchema.nullish(),
   source: z.preprocess(trim, z.string().min(1).max(32)).optional(),
   defaultWorkspace: z.preprocess(trim, z.string().min(1)).optional(),
 });

apps/web/lib/actions/partners/update-partner-profile.ts

@@ -13,12 +13,13 @@ import { PartnerProfileType } from "@prisma/client";
 import { waitUntil } from "@vercel/functions";
 import { stripe } from "../../stripe";
 import z from "../../zod";
+import { base64ImageSchema } from "../../zod/schemas/misc";
 import { authPartnerActionClient } from "../safe-action";
 
 const updatePartnerProfileSchema = z
   .object({
     name: z.string(),
-    image: z.string().nullable(),
+    image: base64ImageSchema.nullish(),
     description: z.string().nullable(),
     country: z.enum(Object.keys(COUNTRIES) as [string, ...string[]]).nullable(),
     profileType: z.nativeEnum(PartnerProfileType),

apps/web/lib/zod/schemas/domains.ts

@@ -1,6 +1,10 @@
 import { normalizeWorkspaceId } from "@/lib/api/workspace-id";
 import z from "@/lib/zod";
-import { booleanQuerySchema, getPaginationQuerySchema } from "./misc";
+import {
+  base64ImageSchema,
+  booleanQuerySchema,
+  getPaginationQuerySchema,
+} from "./misc";
 import { parseUrlSchemaAllowEmpty } from "./utils";
 
 export const RegisteredDomainSchema = z.object({
@@ -140,12 +144,7 @@ export const createDomainBodySchema = z.object({
       "Provide context to your teammates in the link creation modal by showing them an example of a link to be shortened.",
     )
     .openapi({ example: "https://dub.co/help/article/what-is-dub" }),
-  logo: z
-    .string()
-    .trim()
-    .nullish()
-    .transform((v) => v || null)
-    .describe("The logo of the domain."),
+  logo: base64ImageSchema.nullish().describe("The logo of the domain."),
   assetLinks: z
     .string()
     .nullish()

apps/web/lib/zod/schemas/integration.ts

@@ -1,5 +1,6 @@
 import { R2_URL } from "@dub/utils";
 import { z } from "zod";
+import { base64ImageSchema } from "./misc";
 
 export const integrationSchema = z.object({
   id: z.string(),
@@ -37,12 +38,7 @@ export const createIntegrationSchema = z.object({
       message: "installUrl must be a valid URL",
     })
     .nullish(),
-  logo: z
-    .string()
-    .url({
-      message: "Please provide a valid URL for the logo",
-    })
-    .nullish(),
+  logo: base64ImageSchema.nullish().describe("The logo of the integration."),
   description: z
     .string()
     .max(120, {

apps/web/lib/zod/schemas/misc.ts

@@ -49,3 +49,29 @@ export const maxDurationSchema = z.coerce
     message: `Max duration must be ${RECURRING_MAX_DURATIONS.join(", ")}`,
   })
   .nullish();
+
+// Base64 encoded image
+export const base64ImageSchema = z
+  .string()
+  .regex(/^data:image\/(png|jpeg|jpg|gif|webp);base64,/, {
+    message:
+      "Invalid image format. Must be a base64 encoded image with valid image type (png, jpeg, jpg, gif, webp).",
+  })
+  .refine(
+    (str) => {
+      const base64Data = str.split(",")[1];
+
+      if (!base64Data) {
+        return false;
+      }
+
+      return /^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(
+        base64Data,
+      );
+    },
+    {
+      message:
+        "Invalid base64 content. The image data is not properly base64 encoded.",
+    },
+  )
+  .transform((v) => v || null);

apps/web/lib/zod/schemas/workspaces.ts

@@ -2,7 +2,7 @@ import z from "@/lib/zod";
 import { DEFAULT_REDIRECTS, RESERVED_SLUGS, validSlugRegex } from "@dub/utils";
 import slugify from "@sindresorhus/slugify";
 import { DomainSchema } from "./domains";
-import { planSchema, roleSchema } from "./misc";
+import { base64ImageSchema, planSchema, roleSchema } from "./misc";
 
 export const workspaceIdSchema = z.object({
   workspaceId: z
@@ -139,7 +139,7 @@ export const createWorkspaceSchema = z.object({
         message: "Cannot use reserved slugs",
       },
     ),
-  logo: z.string().optional(),
+  logo: base64ImageSchema.nullish(),
   conversionEnabled: z.boolean().optional(),
 });
 

apps/web/tests/misc/base64.test.ts

@@ -0,0 +1,53 @@
+import { base64ImageSchema } from "@/lib/zod/schemas/misc";
+import { describe, expect, it } from "vitest";
+
+describe("base64ImageSchema", () => {
+  it("should validate a correct base64 PNG image", () => {
+    const validBase64Image =
+      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==";
+    expect(() => base64ImageSchema.parse(validBase64Image)).not.toThrow();
+  });
+
+  it("should reject an invalid image type", () => {
+    const invalidImageType =
+      "data:image/invalid;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==";
+    expect(() => base64ImageSchema.parse(invalidImageType)).toThrow(
+      "Invalid image format",
+    );
+  });
+
+  it("should reject malformed base64 data", () => {
+    const malformedBase64 = "data:image/png;base64,invalid-base64-data";
+    expect(() => base64ImageSchema.parse(malformedBase64)).toThrow(
+      "Invalid base64 content",
+    );
+  });
+
+  it("should reject a string without data URI prefix", () => {
+    const noPrefix =
+      "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==";
+    expect(() => base64ImageSchema.parse(noPrefix)).toThrow(
+      "Invalid image format",
+    );
+  });
+
+  it("should reject URLs", () => {
+    const url = "https://github.com/dubinc/dub";
+    expect(() => base64ImageSchema.parse(url)).toThrow("Invalid image format");
+  });
+
+  it("should reject URLs with image extensions", () => {
+    const imageUrl = "https://example.com/image.png";
+    expect(() => base64ImageSchema.parse(imageUrl)).toThrow(
+      "Invalid image format",
+    );
+  });
+
+  it("should reject URLs with data URI prefix but invalid format", () => {
+    const invalidDataUri =
+      "data:image/png;base64,https://example.com/image.png";
+    expect(() => base64ImageSchema.parse(invalidDataUri)).toThrow(
+      "Invalid base64 content",
+    );
+  });
+});