@walkhi/code-relax 0.1.0-beta.3 → 0.1.0-beta.30

package/dist/src/lan-socket-server.mjs

@@ -1,35 +1,124 @@
 import { EventEmitter } from 'node:events';
 import crypto from 'node:crypto';
 import { WebSocketServer, WebSocket } from 'ws';
+import timelineSync from '../shared/timeline-sync.cjs';
+import { createFileUploadSession, FILE_CHUNK_BYTES, MAX_FILE_BYTES, MAX_FILE_COUNT, MAX_TOTAL_BYTES } from './file-upload-session.mjs';
+import { FILE_READ_LIMIT } from './file-transfer-range.mjs';
 
 // One authenticated connection carries requests and explicitly owned subscriptions.
 // The dispatcher is shared business routing, not an HTTP loopback proxy.
-export function attachLanSocket(server, { token, dispatch, authorizeDevice = () => false }) {
-  // Allow four individually valid images to reach the aggregate-size validator,
-  // so an oversized composition returns an error instead of dropping the connection.
-  const wss = new WebSocketServer({ noServer: true, maxPayload: 48 * 1024 * 1024, perMessageDeflate: false });
+export function attachLanSocket(server, { token, dispatch, authorizeDevice = () => false, onDeviceAuthorizationsChanged, uploadsRoot }) {
+  // Allow four individually valid images to reach the aggregate-size validator,
+  // so an oversized composition returns an error instead of dropping the connection.
+  const wss = new WebSocketServer({ noServer: true, maxPayload: 48 * 1024 * 1024, perMessageDeflate: false });
+  const uploadWss = new WebSocketServer({ noServer: true, maxPayload: 512 * 1024, perMessageDeflate: false });
+  const activeUploads = new Map();
+  const deviceConnections = new Map();
+  const stopWatching = onDeviceAuthorizationsChanged?.(() => {
+    for (const connection of deviceConnections.values()) {
+      if (!authorizeDevice(connection.clientId, connection.credential)) connection.revoke();
+    }
+  });
   server.on('upgrade', (request, socket, head) => {
     const origin = request.headers.origin;
-    let allowed = request.url === '/api/socket';
-    // Packaged ArkWeb pages have an opaque ("null") origin; device authentication remains required.
-    if (origin && origin !== 'null') {
-      try {
-        const originUrl = new URL(origin);
-        const requestUrl = new URL(`http://${request.headers.host}`);
-        const nativeClient = request.headers['x-code-relax-native'] === '1';
-        allowed &&= ['http:', 'https:'].includes(originUrl.protocol) &&
-          (originUrl.host === request.headers.host || (nativeClient && originUrl.hostname === requestUrl.hostname));
-      }
-      catch { allowed = false; }
+    const upload = request.url === '/api/file-transfer' && Boolean(uploadsRoot);
+    let allowed = request.url === '/api/socket' || upload;
+    // Packaged ArkWeb pages have an opaque ("null") origin; device authentication remains required.
+    if (origin && origin !== 'null') {
+      try {
+        const originUrl = new URL(origin);
+        const requestUrl = new URL(`http://${request.headers.host}`);
+        const nativeClient = request.headers['x-code-relax-native'] === '1';
+        allowed &&= ['http:', 'https:'].includes(originUrl.protocol) &&
+          (originUrl.host === request.headers.host || (nativeClient && originUrl.hostname === requestUrl.hostname));
+      }
+      catch { allowed = false; }
     }
     if (!allowed) { socket.end('HTTP/1.1 403 Forbidden\r\nConnection: close\r\n\r\n'); return; }
-    wss.handleUpgrade(request, socket, head, ws => wss.emit('connection', ws));
+    const target = upload ? uploadWss : wss;
+    target.handleUpgrade(request, socket, head, ws => target.emit('connection', ws));
+  });
+  uploadWss.on('connection', ws => {
+    let authenticated = false, stopped = false, identity = '', upload, purpose, chunk;
+    let queued = 0;
+    let pending = Promise.resolve();
+    const authTimer = setTimeout(() => ws.close(1008, 'Authentication required'), 5000);
+    const reply = value => { if (!stopped && ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(value)); };
+    ws.on('error', () => {});
+    const stop = () => {
+      if (stopped) return;
+      stopped = true;
+      deviceConnections.delete(ws);
+      clearTimeout(authTimer);
+      void pending.finally(async () => {
+        await upload?.close();
+        activeUploads.delete(ws);
+      });
+    };
+    ws.on('close', stop);
+    ws.on('message', (bytes, binary) => {
+      if (stopped || ws.readyState !== WebSocket.OPEN) return;
+      if (++queued > 4) { stop(); ws.close(1013, 'File request queue exceeded'); return; }
+      pending = pending.then(async () => {
+        if (stopped || ws.readyState !== WebSocket.OPEN) return;
+        if (binary) {
+          if (!authenticated || purpose !== 'upload' || !chunk || bytes.length !== chunk.length) throw new Error('无效的文件分片');
+          const current = chunk; chunk = undefined;
+          reply(await upload.receive({ ...current, data: bytes })); return;
+        }
+        if (bytes.length > 16 * 1024 || chunk) throw new Error('无效的文件请求');
+        const message = JSON.parse(bytes.toString());
+        if (!authenticated) {
+          purpose = message.purpose;
+          if (!['upload', 'download'].includes(purpose)) throw new Error('文件通道用途无效');
+          if (message?.type === 'device-auth' && typeof message.clientId === 'string' &&
+              authorizeDevice(message.clientId, message.credential) === true) identity = `device:${message.clientId}`;
+          else if (token && message?.type === 'auth' && typeof message.token === 'string') {
+            const supplied = Buffer.from(message.token), expected = Buffer.from(token || '');
+            if (supplied.length === expected.length && crypto.timingSafeEqual(supplied, expected)) identity = 'local-token';
+          }
+          const same = [...activeUploads.values()].filter(item => item.identity === identity && item.purpose === purpose);
+          const uploads = [...activeUploads.values()].filter(item => item.purpose === 'upload').length;
+          if (!identity || same.length >= (purpose === 'upload' ? 1 : 3) || activeUploads.size >= 8 ||
+              (purpose === 'upload' && uploads >= 2)) throw new Error('文件身份无效或文件任务过多');
+          if (message.type === 'device-auth') deviceConnections.set(ws, {
+            clientId: message.clientId, credential: message.credential,
+            revoke() { stop(); ws.close(4401, 'Device authorization revoked'); },
+          });
+          activeUploads.set(ws, { identity, purpose }); authenticated = true; clearTimeout(authTimer);
+          if (purpose === 'upload') upload = createFileUploadSession(uploadsRoot);
+          reply({ type: 'ready', version: 2, purpose, maxDownloadBytes: FILE_READ_LIMIT, maxFileBytes: MAX_FILE_BYTES, maxFileCount: MAX_FILE_COUNT,
+            maxTotalBytes: MAX_TOTAL_BYTES, chunkBytes: FILE_CHUNK_BYTES });
+          return;
+        }
+        if (purpose === 'download') {
+          if (message.type !== 'read' || typeof message.path !== 'string' || message.path.length > 12000 ||
+              !/^\/api\/(?:images\/[a-f0-9]{64}|files\/(?:local|[a-f0-9]{64})|code\/(?:image|binary)|desktop\/frames\/[a-f0-9]{32})(?:\?|$)/.test(message.path) || message.path.includes('#')) throw new Error('无效的文件读取请求');
+          await dispatch({ method: 'GET', url: message.path, headers: {}, socketRequest: true, fileTransfer: true }, {
+            json(status, body) { reply({ type: 'error', status, error: body?.error || '文件读取失败' }); },
+            filePart(metadata, data) {
+              if (stopped || ws.readyState !== WebSocket.OPEN) return;
+              if (!Buffer.isBuffer(data) || data.length > FILE_CHUNK_BYTES || ws.bufferedAmount > 1024 * 1024) throw new Error('文件接收过慢或分片无效');
+              reply({ type: 'part', ...metadata, length: data.length });
+              ws.send(data);
+            },
+          });
+          return;
+        }
+        if (message.type === 'chunk') {
+          if (!Number.isSafeInteger(message.length) || message.length < 1 || message.length > FILE_CHUNK_BYTES) throw new Error('文件分片大小无效');
+          chunk = message; return;
+        }
+        reply(await upload.receive(message));
+      }).catch(error => { reply({ type: 'error', error: error.message }); stop(); ws.close(1008, 'File request rejected'); })
+        .finally(() => { queued--; });
+    });
   });
   wss.on('connection', ws => {
-    let authenticated = false, alive = true;
+    let authenticated = false, alive = true, stopped = false;
     const active = new Map();
     const send = value => {
-      if (ws.readyState !== WebSocket.OPEN) return;
+      if (stopped || ws.readyState !== WebSocket.OPEN) return;
       if (ws.bufferedAmount > 64 * 1024 * 1024) { ws.close(1013, 'Client too slow'); return; }
       ws.send(JSON.stringify(value));
     };
@@ -40,62 +129,123 @@ export function attachLanSocket(server, { token, dispatch, authorizeDevice = ()
     }, 30000);
     ws.on('pong', () => { alive = true; });
     ws.on('error', () => {}); // close rejects requests and releases every subscription.
-    ws.on('close', () => {
+    const stop = () => {
+      if (stopped) return;
+      stopped = true;
+      deviceConnections.delete(ws);
       clearTimeout(authTimer); clearInterval(heartbeat);
       for (const reply of active.values()) reply.end();
       active.clear();
-    });
+    };
+    ws.on('close', stop);
     ws.on('message', async (bytes, binary) => {
+      if (stopped || ws.readyState !== WebSocket.OPEN) return;
       let message;
       try { if (binary) throw new Error(); message = JSON.parse(bytes.toString()); }
       catch { ws.close(1008, 'Invalid JSON'); return; }
       if (!message || typeof message !== 'object') { ws.close(1008, 'Invalid message'); return; }
-      if (!authenticated) {
-        let accepted = false;
-        if (message.type === 'device-auth') {
-          accepted = authorizeDevice(message.clientId, message.credential) === true;
-        } else if (message.type === 'auth') {
-          const supplied = Buffer.from(typeof message.token === 'string' ? message.token : '');
-          const expected = Buffer.from(token || '');
-          accepted = supplied.length === expected.length && crypto.timingSafeEqual(supplied, expected);
-        }
-        if (!accepted) {
-          ws.close(1008, 'Authentication failed'); return;
-        }
+      if (!authenticated) {
+        let accepted = false;
+        if (message.type === 'device-auth') {
+          accepted = authorizeDevice(message.clientId, message.credential) === true;
+        } else if (message.type === 'auth') {
+          const supplied = Buffer.from(typeof message.token === 'string' ? message.token : '');
+          const expected = Buffer.from(token || '');
+          accepted = supplied.length === expected.length && crypto.timingSafeEqual(supplied, expected);
+        }
+        if (!accepted) {
+          ws.close(1008, 'Authentication failed'); return;
+        }
+        if (message.type === 'device-auth') deviceConnections.set(ws, {
+          clientId: message.clientId, credential: message.credential,
+          revoke() { stop(); ws.close(4401, 'Device authorization revoked'); },
+        });
         authenticated = true; clearTimeout(authTimer); send({ type: 'ready', version: 1 }); return;
       }
       const { id, type } = message;
       if (!Number.isSafeInteger(id) || id < 1) { ws.close(1008, 'Invalid request id'); return; }
+      if (type === 'response-ack') { active.get(id)?.acknowledge?.(message.batch); return; }
+      if (type === 'cancel') { active.get(id)?.end(); return; }
       if (type === 'unsubscribe') { const reply = active.get(id); if (reply?.subscription) reply.end(); return; }
       if (!['request', 'subscribe'].includes(type) || active.has(id)) { ws.close(1008, 'Invalid request'); return; }
       const subscription = type === 'subscribe';
       const path = message.path;
-      const isStream = typeof path === 'string' && /^\/api\/(?:usage-events|shared-recovery\/events|threads\/[^/?]+\/events)(?:\?|$)/.test(path);
-      if (active.size >= 32 || (subscription && [...active.values()].filter(r => r.subscription).length >= 4)) {
+      const isStream = typeof path === 'string' && /^\/api\/(?:usage-events|thread-events|turn-events|shared-recovery\/events|threads\/[^/?]+\/(?:events|dots-activity-events))(?:\?|$)/.test(path);
+      // Native recovery, usage, catalog, current chat, and the main Dot's sidebar
+      // share one station connection (including when viewing a delegated task).
+      if (active.size >= 32 || (subscription && [...active.values()].filter(r => r.subscription).length >= 5)) {
         send({ type: 'response', id, status: 429, body: { error: '同时进行的请求过多' } }); return;
       }
       if (typeof path !== 'string' || !path.startsWith('/api/') || path.includes('#') ||
           path.startsWith('/api/attachments/') || path === '/api/socket' || path === '/api/uploads' ||
-          subscription !== isStream || !['GET', 'POST', 'DELETE'].includes(message.method || 'GET') ||
+          subscription !== isStream || !['GET', 'POST', 'PUT', 'DELETE'].includes(message.method || 'GET') ||
           (subscription && message.method && message.method !== 'GET') ||
           (message.body !== undefined && (!message.body || typeof message.body !== 'object' || Array.isArray(message.body)))) {
         send({ type: 'response', id, status: 400, body: { error: '无效的连接请求' } }); return;
       }
       const reply = new EventEmitter();
+      let waiting, batch = 0;
+      const sendPart = (pieces, pageEnd = false) => new Promise((resolve, reject) => {
+        const timer = setTimeout(() => reject(new Error('历史接收确认超时')), 20000);
+        waiting = { resolve, reject, timer, batch: ++batch };
+        send({ type: subscription ? 'event-part' : 'response-part', id, batch, pieces, ...(pageEnd ? { pageEnd: true } : {}) });
+      });
+      const sendTimeline = async (status, body, event) => {
+        if (event) batch = 0;
+        try {
+          for (const pieces of timelineSync.timelineBatches(body)) {
+            if (reply.writableEnded) return;
+            await sendPart(pieces);
+          }
+          if (!reply.writableEnded) {
+            send(event ? { type: 'event', id, event, streamed: true } : { type: 'response', id, status, streamed: true });
+            if (!event) reply.end();
+          }
+        } catch (error) {
+          if (!reply.writableEnded) {
+            send(event ? { type: 'event', id, event: 'bridge-error', data: { error: error.message } }
+              : { type: 'response', id, status: 500, body: { error: error.message } });
+            reply.end();
+          }
+        }
+      };
       Object.assign(reply, {
         subscription, headersSent: false, writableEnded: false, destroyed: false,
+        async historyPage(page) {
+          const parts = timelineSync.timelineBatches(page);
+          let part = parts.next();
+          while (!part.done) {
+            if (this.writableEnded) throw new Error('历史读取已取消');
+            const next = parts.next();
+            await sendPart(part.value, next.done);
+            part = next;
+          }
+        },
         json(status, body) {
           if (this.writableEnded) return;
+          if (status < 400 && body?.schemaVersion === 1 && Array.isArray(body.turns)) {
+            void sendTimeline(status, body); return;
+          }
           send({ type: 'response', id, status, body }); this.end();
         },
+        acknowledge(sequence) {
+          if (!waiting || sequence !== waiting.batch) return;
+          const completed = waiting; waiting = null; clearTimeout(completed.timer);
+          this.emit('progress'); completed.resolve();
+        },
         startEvents() {
           if (this.writableEnded) return;
           this.headersSent = true; send({ type: 'subscribed', id });
         },
-        event(event, data) { if (!this.writableEnded) send({ type: 'event', id, event, data }); },
+        event(event, data) {
+          if (this.writableEnded) return;
+          if (event === 'update' && data?.schemaVersion === 1) return sendTimeline(200, data, event);
+          send({ type: 'event', id, event, data });
+        },
         end() {
           if (this.writableEnded) return;
           this.writableEnded = this.destroyed = true; active.delete(id);
+          if (waiting) { clearTimeout(waiting.timer); waiting.reject(new Error('历史读取已取消')); waiting = null; }
           if (subscription) send({ type: 'end', id });
           this.emit('close');
         },
@@ -105,5 +255,9 @@ export function attachLanSocket(server, { token, dispatch, authorizeDevice = ()
       catch (error) { reply.json(500, { error: error.message }); }
     });
   });
-  return { close() { for (const ws of wss.clients) ws.terminate(); wss.close(); } };
+  return { close() {
+    stopWatching?.();
+    for (const ws of wss.clients) ws.terminate(); wss.close();
+    for (const ws of uploadWss.clients) ws.terminate(); uploadWss.close();
+  } };
 }