From 19a45f9e5f3e789503c7f7ddce4597aa13c8d2fc Mon Sep 17 00:00:00 2001
From: digua <i@digua.me>
Date: Mon, 25 May 2026 23:23:21 +0800
Subject: [PATCH] fix: protect self update and migration retry

---
 apps/cli/src/http/auth.test.ts  | 17 +++++++++++++++++
 apps/cli/src/http/auth.ts       | 11 +++++++++++
 apps/desktop/main/paths.test.ts | 10 ++++++++++
 apps/desktop/main/paths.ts      |  8 +++++++-
 src/services/platform/web.ts    |  7 ++++++-
 5 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 apps/cli/src/http/auth.test.ts
 create mode 100644 apps/desktop/main/paths.test.ts

diff --git a/apps/cli/src/http/auth.test.ts b/apps/cli/src/http/auth.test.ts
new file mode 100644
index 00000000..b743bc7e
--- /dev/null
+++ b/apps/cli/src/http/auth.test.ts
@@ -0,0 +1,17 @@
+import { describe, it } from 'node:test'
+import assert from 'node:assert/strict'
+import { createServer } from './server'
+import { setAuthToken } from './auth'
+
+describe('CLI HTTP auth hook', () => {
+  it('requires bearer token for the CLI self-update endpoint under /_web', async () => {
+    setAuthToken('clb_test_token')
+    const app = createServer()
+    app.post('/_web/system/update', async () => ({ success: true }))
+
+    const response = await app.inject({ method: 'POST', url: '/_web/system/update' })
+    await app.close()
+
+    assert.equal(response.statusCode, 401)
+  })
+})
diff --git a/apps/cli/src/http/auth.ts b/apps/cli/src/http/auth.ts
index 5ae63289..a7af8de6 100644
--- a/apps/cli/src/http/auth.ts
+++ b/apps/cli/src/http/auth.ts
@@ -10,6 +10,7 @@ import { unauthorized, errorResponse } from './errors'
 
 let cachedToken: string | null = null
 let webModeEnabled = false
+const protectedWebPaths = new Set(['/_web/system/update'])
 
 /**
  * 设置 auth hook 使用的 token（由 server 启动时注入）
@@ -36,12 +37,22 @@ function safeTokenCompare(a: string, b: string): boolean {
 export async function authHook(request: FastifyRequest, reply: FastifyReply): Promise<void> {
   if (!cachedToken) return
 
+  if (protectedWebPaths.has(request.url.split('?')[0])) {
+    return requireBearerToken(request, reply)
+  }
+
   // /_web/ internal API: no auth required (same-origin Web UI only)
   if (request.url.startsWith('/_web/')) return
 
   // Web mode: only /api/ routes require auth; static files and SPA are public
   if (webModeEnabled && !request.url.startsWith('/api/')) return
 
+  return requireBearerToken(request, reply)
+}
+
+function requireBearerToken(request: FastifyRequest, reply: FastifyReply): void {
+  if (!cachedToken) return
+
   const authHeader = request.headers.authorization
   if (!authHeader || !authHeader.startsWith('Bearer ')) {
     const err = unauthorized()
diff --git a/apps/desktop/main/paths.test.ts b/apps/desktop/main/paths.test.ts
new file mode 100644
index 00000000..46cb93bd
--- /dev/null
+++ b/apps/desktop/main/paths.test.ts
@@ -0,0 +1,10 @@
+import { describe, it } from 'node:test'
+import assert from 'node:assert/strict'
+import { shouldMarkUnifiedDirMigrationDone } from './paths'
+
+describe('desktop path migration', () => {
+  it('marks unified directory migration done only when no directory failed', () => {
+    assert.equal(shouldMarkUnifiedDirMigrationDone([]), true)
+    assert.equal(shouldMarkUnifiedDirMigrationDone(['settings']), false)
+  })
+})
diff --git a/apps/desktop/main/paths.ts b/apps/desktop/main/paths.ts
index 73132af7..12986af0 100644
--- a/apps/desktop/main/paths.ts
+++ b/apps/desktop/main/paths.ts
@@ -573,6 +573,10 @@ export function removeLegacyDir(): boolean {
 
 const SYSTEM_SUBDIRS = ['ai', 'settings', 'cache', 'logs', 'temp', 'nlp']
 
+export function shouldMarkUnifiedDirMigrationDone(failedDirs: string[]): boolean {
+  return failedDirs.length === 0
+}
+
 /**
  * 检测是否需要从 Electron 旧目录结构迁移到新的双根目录结构
  *
@@ -712,7 +716,9 @@ export function migrateToUnifiedDirs(): { success: boolean; error?: string } {
     writeMigrationLog(getLogsDir(), summary, ensureDir)
     console.log(`[Migration] ${summary}`)
 
-    writeConfigField('data', 'electron_migration_done', true)
+    if (shouldMarkUnifiedDirMigrationDone(failedDirs)) {
+      writeConfigField('data', 'electron_migration_done', true)
+    }
 
     return { success: failedDirs.length === 0 }
   } catch (error) {
diff --git a/src/services/platform/web.ts b/src/services/platform/web.ts
index 6a459521..ae2b639e 100644
--- a/src/services/platform/web.ts
+++ b/src/services/platform/web.ts
@@ -74,7 +74,12 @@ export class WebPlatformAdapter implements PlatformAdapter {
 
   async performUpdate(): Promise<PerformUpdateResult> {
     try {
-      const resp = await fetch('/_web/system/update', { method: 'POST' })
+      const configResp = await fetch('/_web/automation/config')
+      const config = (await configResp.json()) as { token?: string }
+      const headers: Record<string, string> = {}
+      if (config.token) headers.Authorization = `Bearer ${config.token}`
+
+      const resp = await fetch('/_web/system/update', { method: 'POST', headers })
       return await resp.json()
     } catch (err) {
       return { success: false, error: err instanceof Error ? err.message : String(err) }