[codex] Harden desktop image data URL handling

apps/desktop/src/lib/trpc/routers/window.ts

@@ -1,8 +1,8 @@
 import fs from "node:fs/promises";
 import { homedir } from "node:os";
-import path from "node:path";
 import type { BrowserWindow } from "electron";
 import { dialog } from "electron";
+import { getImageMimeType } from "shared/file-types";
 import { z } from "zod";
 import { publicProcedure, router } from "..";
 
@@ -119,10 +119,9 @@ export const createWindowRouter = (getWindow: () => BrowserWindow | null) => {
 
 			const filePath = result.filePaths[0];
 			const buffer = await fs.readFile(filePath);
-			const ext = path.extname(filePath).slice(1).toLowerCase();
-			const mimeType = ext === "jpg" ? "jpeg" : ext;
+			const mimeType = getImageMimeType(filePath) ?? "image/png";
 			const base64 = buffer.toString("base64");
-			const dataUrl = `data:image/${mimeType};base64,${base64}`;
+			const dataUrl = `data:${mimeType};base64,${base64}`;
 
 			return { canceled: false, dataUrl };
 		}),

apps/desktop/src/main/lib/project-icons.test.ts

@@ -0,0 +1,57 @@
+import { describe, expect, test } from "bun:test";
+import { parseProjectIconDataUrl } from "./project-icons";
+
+const ICON_BASE64 = Buffer.from("icon").toString("base64");
+
+describe("parseProjectIconDataUrl", () => {
+	test("parses PNG data URLs", () => {
+		const result = parseProjectIconDataUrl(
+			`data:image/png;base64,${ICON_BASE64}`,
+		);
+
+		expect(result.ext).toBe("png");
+		expect(result.buffer.toString()).toBe("icon");
+	});
+
+	test("normalizes JPEG MIME types to the jpg extension", () => {
+		const result = parseProjectIconDataUrl(
+			`data:image/jpeg;base64,${ICON_BASE64}`,
+		);
+
+		expect(result.ext).toBe("jpg");
+		expect(result.buffer.toString()).toBe("icon");
+	});
+
+	test("parses SVG data URLs with extra MIME parameters", () => {
+		const result = parseProjectIconDataUrl(
+			`data:image/svg+xml;charset=utf-8;base64,${ICON_BASE64}`,
+		);
+
+		expect(result.ext).toBe("svg");
+		expect(result.buffer.toString()).toBe("icon");
+	});
+
+	test("maps ICO MIME types to the ico extension", () => {
+		const xIcon = parseProjectIconDataUrl(
+			`data:image/x-icon;base64,${ICON_BASE64}`,
+		);
+		const microsoftIcon = parseProjectIconDataUrl(
+			`data:image/vnd.microsoft.icon;base64,${ICON_BASE64}`,
+		);
+
+		expect(xIcon.ext).toBe("ico");
+		expect(microsoftIcon.ext).toBe("ico");
+	});
+
+	test("rejects unsupported image MIME types", () => {
+		expect(() =>
+			parseProjectIconDataUrl(`data:image/webp;base64,${ICON_BASE64}`),
+		).toThrow("Unsupported icon format");
+	});
+
+	test("rejects malformed data URLs", () => {
+		expect(() => parseProjectIconDataUrl("not-a-data-url")).toThrow(
+			"Invalid data URL format",
+		);
+	});
+});

apps/desktop/src/main/lib/project-icons.ts

@@ -2,12 +2,17 @@ import { randomUUID } from "node:crypto";
 import { existsSync, mkdirSync, readdirSync, unlinkSync } from "node:fs";
 import { copyFile, writeFile } from "node:fs/promises";
 import { extname, join } from "node:path";
+import {
+	getImageExtensionFromMimeType,
+	parseBase64DataUrl,
+} from "shared/file-types";
 import { SUPERSET_HOME_DIR } from "./app-environment";
 
 export const PROJECT_ICONS_DIR = join(SUPERSET_HOME_DIR, "project-icons");
 
 /** Max icon file size: 512KB */
 const MAX_ICON_SIZE = 512 * 1024;
+const PROJECT_ICON_EXTENSIONS = new Set(["png", "jpg", "svg", "ico"]);
 
 /**
  * Ensures the project icons directory exists.
@@ -52,6 +57,25 @@ export function getProjectIconProtocolUrl(projectId: string): string {
 	return `superset-icon://projects/${projectId}?v=${encodeURIComponent(randomUUID())}`;
 }
 
+export function parseProjectIconDataUrl(dataUrl: string): {
+	buffer: Buffer;
+	ext: string;
+} {
+	const { base64Data, mimeType } = parseBase64DataUrl(dataUrl);
+	const ext = getImageExtensionFromMimeType(mimeType);
+
+	if (!ext || !PROJECT_ICON_EXTENSIONS.has(ext)) {
+		throw new Error(
+			"Unsupported icon format. Supported formats are PNG, JPEG, SVG, and ICO.",
+		);
+	}
+
+	return {
+		buffer: Buffer.from(base64Data, "base64"),
+		ext,
+	};
+}
+
 /**
  * Saves an icon file for a project from a local file path.
  * Copies the file to PROJECT_ICONS_DIR/{projectId}.{ext}.
@@ -89,14 +113,7 @@ export async function saveProjectIconFromDataUrl({
 	ensureProjectIconsDir();
 	removeExistingIcon(projectId);
 
-	// Parse data URL: data:image/png;base64,<data>
-	const match = dataUrl.match(/^data:image\/(\w+);base64,(.+)$/);
-	if (!match) {
-		throw new Error("Invalid data URL format");
-	}
-
-	const ext = match[1] === "jpeg" ? "jpg" : match[1];
-	const buffer = Buffer.from(match[2], "base64");
+	const { buffer, ext } = parseProjectIconDataUrl(dataUrl);
 
 	if (buffer.length > MAX_ICON_SIZE) {
 		throw new Error(

apps/desktop/src/renderer/routes/_authenticated/settings/account/components/AccountSettings/AccountSettings.tsx

@@ -10,6 +10,10 @@ import { apiTrpcClient } from "renderer/lib/api-trpc-client";
 import { authClient } from "renderer/lib/auth-client";
 import { electronTrpc } from "renderer/lib/electron-trpc";
 import { useCollections } from "renderer/routes/_authenticated/providers/CollectionsProvider";
+import {
+	getImageExtensionFromMimeType,
+	parseBase64DataUrl,
+} from "shared/file-types";
 import {
 	isItemVisible,
 	SETTING_ITEM_ID,
@@ -64,9 +68,8 @@ export function AccountSettings({ visibleItems }: AccountSettingsProps) {
 			const result = await selectImageMutation.mutateAsync();
 			if (result.canceled || !result.dataUrl) return;
 
-			const mimeMatch = result.dataUrl.match(/^data:([^;]+);/);
-			const mimeType = mimeMatch?.[1] || "image/png";
-			const ext = mimeType.split("/")[1] || "png";
+			const { mimeType } = parseBase64DataUrl(result.dataUrl);
+			const ext = getImageExtensionFromMimeType(mimeType) ?? "png";
 
 			const uploadResult = await apiTrpcClient.user.uploadAvatar.mutate({
 				fileData: result.dataUrl,

apps/desktop/src/renderer/routes/_authenticated/settings/organization/components/OrganizationSettings/OrganizationSettings.tsx

@@ -26,6 +26,10 @@ import { apiTrpcClient } from "renderer/lib/api-trpc-client";
 import { authClient } from "renderer/lib/auth-client";
 import { electronTrpc } from "renderer/lib/electron-trpc";
 import { useCollections } from "renderer/routes/_authenticated/providers/CollectionsProvider";
+import {
+	getImageExtensionFromMimeType,
+	parseBase64DataUrl,
+} from "shared/file-types";
 import { MemberActions } from "../../../members/components/MembersSettings/components/MemberActions";
 import { PendingInvitations } from "../../../members/components/PendingInvitations";
 import type { TeamMember } from "../../../members/types";
@@ -140,9 +144,8 @@ export function OrganizationSettings({
 			const result = await selectImageMutation.mutateAsync();
 			if (result.canceled || !result.dataUrl) return;
 
-			const mimeMatch = result.dataUrl.match(/^data:([^;]+);/);
-			const mimeType = mimeMatch?.[1] || "image/png";
-			const ext = mimeType.split("/")[1] || "png";
+			const { mimeType } = parseBase64DataUrl(result.dataUrl);
+			const ext = getImageExtensionFromMimeType(mimeType) ?? "png";
 
 			const uploadResult = await apiTrpcClient.organization.uploadLogo.mutate({
 				organizationId: organization.id,

apps/desktop/src/renderer/routes/_authenticated/settings/project/$projectId/components/ProjectSettings/ProjectSettings.tsx

@@ -133,6 +133,7 @@ export function ProjectSettings({
 	const setProjectIcon = electronTrpc.projects.setProjectIcon.useMutation({
 		onError: (err) => {
 			console.error("[project-settings/setProjectIcon] Failed:", err);
+			toast.error(err.message || "Failed to update project icon");
 		},
 		onSettled: () => {
 			utils.projects.get.invalidate({ id: projectId });
@@ -621,7 +622,7 @@ export function ProjectSettings({
 							<input
 								ref={fileInputRef}
 								type="file"
-								accept="image/png,image/jpeg,image/svg+xml,image/x-icon"
+								accept="image/png,image/jpeg,image/svg+xml,image/x-icon,image/vnd.microsoft.icon,.ico"
 								className="hidden"
 								onChange={handleFileChange}
 							/>

apps/desktop/src/shared/file-types.test.ts

@@ -0,0 +1,42 @@
+import { describe, expect, test } from "bun:test";
+import {
+	getImageExtensionFromMimeType,
+	getImageMimeType,
+	parseBase64DataUrl,
+} from "./file-types";
+
+const PNG_BASE64 = Buffer.from("png").toString("base64");
+
+describe("file-types", () => {
+	test("maps image file paths to MIME types", () => {
+		expect(getImageMimeType("logo.svg")).toBe("image/svg+xml");
+		expect(getImageMimeType("logo.ico")).toBe("image/x-icon");
+		expect(getImageMimeType("logo.unknown")).toBeNull();
+	});
+
+	test("maps image MIME types to preferred extensions", () => {
+		expect(getImageExtensionFromMimeType("image/jpeg")).toBe("jpg");
+		expect(getImageExtensionFromMimeType("image/vnd.microsoft.icon")).toBe(
+			"ico",
+		);
+		expect(getImageExtensionFromMimeType("image/webp")).toBe("webp");
+		expect(getImageExtensionFromMimeType("image/avif")).toBeNull();
+	});
+
+	test("parses base64 data URLs with extra MIME parameters", () => {
+		expect(
+			parseBase64DataUrl(
+				`data:image/svg+xml;charset=utf-8;base64,${PNG_BASE64}`,
+			),
+		).toEqual({
+			base64Data: PNG_BASE64,
+			mimeType: "image/svg+xml",
+		});
+	});
+
+	test("rejects malformed base64 data URLs", () => {
+		expect(() => parseBase64DataUrl("not-a-data-url")).toThrow(
+			"Invalid data URL format",
+		);
+	});
+});

apps/desktop/src/shared/file-types.ts

@@ -27,6 +27,19 @@ const IMAGE_MIME_TYPES: Record<string, string> = {
 	ico: "image/x-icon",
 };
 
+/** Extensions for supported image MIME types */
+const IMAGE_MIME_TYPE_EXTENSIONS: Record<string, string> = {
+	"image/png": "png",
+	"image/jpeg": "jpg",
+	"image/jpg": "jpg",
+	"image/gif": "gif",
+	"image/webp": "webp",
+	"image/svg+xml": "svg",
+	"image/bmp": "bmp",
+	"image/x-icon": "ico",
+	"image/vnd.microsoft.icon": "ico",
+};
+
 /** Markdown extensions */
 const MARKDOWN_EXTENSIONS = new Set(["md", "markdown", "mdx"]);
 
@@ -53,6 +66,38 @@ export function getImageMimeType(filePath: string): string | null {
 	return IMAGE_MIME_TYPES[ext] ?? null;
 }
 
+/**
+ * Gets the preferred file extension for an image MIME type.
+ * Returns null if not a supported image type.
+ */
+export function getImageExtensionFromMimeType(mimeType: string): string | null {
+	return IMAGE_MIME_TYPE_EXTENSIONS[mimeType.toLowerCase()] ?? null;
+}
+
+/**
+ * Parses a base64 data URL and returns its MIME type and base64 payload.
+ */
+export function parseBase64DataUrl(dataUrl: string): {
+	base64Data: string;
+	mimeType: string;
+} {
+	const separatorIndex = dataUrl.indexOf(",");
+	if (separatorIndex === -1) {
+		throw new Error("Invalid data URL format");
+	}
+
+	const header = dataUrl.slice(0, separatorIndex);
+	const base64Data = dataUrl.slice(separatorIndex + 1);
+	const mimeMatch = header.match(/^data:([^;,]+)(?:;[^,]*)*;base64$/i);
+	const mimeType = mimeMatch?.[1]?.toLowerCase();
+
+	if (!mimeType) {
+		throw new Error("Invalid data URL format");
+	}
+
+	return { base64Data, mimeType };
+}
+
 /**
  * Checks if a file is markdown based on extension
  */