Deduplicate artifact bytes and deliver files by authenticated hash

2026-10-06 14:35:13 -04:00
parent 18283c63a7
commit 62df14b180
45 changed files with 510 additions and 139 deletions
@@ -0,0 +1,85 @@
+import Database from 'better-sqlite3';
+import { createHash } from 'node:crypto';
+import { afterEach, beforeEach, expect, it, vi } from 'vitest';
+import schema from '../../src/lib/server/db/schema.sql?raw';
+import { newFixture, type Fixture } from '../helpers/fixture.js';
+
+let fixture: Fixture;
+beforeEach(async () => { vi.resetModules(); fixture = await newFixture([]); });
+afterEach(async () => { vi.restoreAllMocks(); await fixture.cleanup(); });
+
+it('shares bytes across versions, staged copies and commands, and scopes downloads to the assigned host', async () => {
+    const { db, closeDb } = await import('../../src/lib/server/db/index.js');
+    const { createScript } = await import('../../src/lib/server/fleet/scripts.js');
+    const { queueFileChange } = await import('../../src/lib/server/fleet/staged-files.js');
+    const hosts = await import('../../src/lib/server/hosts/index.js');
+    const { pruneFileBlobs } = await import('../../src/lib/server/db/file-blobs.js');
+    const content = Buffer.alloc(1024 * 1024, 71);
+    const hash = createHash('sha256').update(content).digest('hex');
+    for (let i = 0; i < 20; i++) {
+        createScript({ name: `plugin-${i}`, kind: 'plugin', targetRelPath: 'plugin.dll', content });
+        queueFileChange({ sourceKind: 'campaign', sourceId: 'campaign', server: `t5_server_${i}`, kind: 'plugin', relPath: 'plugin.dll', content, inventorySource: 'test' });
+    }
+    const host = hosts.createHost({ name: 'Owner' });
+    const credential = hosts.enrollHost(hosts.createEnrollment(host.id).token, { capabilities: ['files.blobs'] }).credential;
+    const other = hosts.createHost({ name: 'Other' });
+    const otherCredential = hosts.enrollHost(hosts.createEnrollment(other.id).token, {}).credential;
+    const payload = { stageId: 'one', files: [{ path: 'plugins/plugin.dll', contentBase64: content.toString('base64'), sha256: hash }] };
+    const op = hosts.queueHostOperation(host.id, 'files.stage', payload, { idempotencyKey: 'one' });
+    expect(hosts.queueHostOperation(host.id, 'files.stage', payload, { idempotencyKey: 'one' }).id).toBe(op.id);
+    expect(() => hosts.queueHostOperation(host.id, 'files.stage', { ...payload, stageId: 'two' }, { idempotencyKey: 'one' })).toThrow('different request');
+    expect(db().prepare('SELECT COUNT(*) n,SUM(size) bytes FROM file_blobs').get()).toEqual({ n: 1, bytes: content.length });
+    const packed = JSON.parse(hosts.getHostOperationPayload(host.id, op.id).toString());
+    expect(packed.files[0]).toEqual({ path: 'plugins/plugin.dll', contentSha256: hash, sha256: hash });
+    expect(pruneFileBlobs()).toBe(0);
+    const { GET } = await import('../../src/routes/api/agent/blobs/[sha256]/+server.js');
+    const request = (token: string) => new Request(`https://ec.test/api/agent/blobs/${hash}`, { headers: { authorization: `Bearer ${token}` } });
+    expect(Buffer.from(await GET({ request: request(credential), params: { sha256: hash } }).arrayBuffer())).toEqual(content);
+    expect(() => GET({ request: request(otherCredential), params: { sha256: hash } })).toThrow();
+    closeDb();
+    expect(hosts.getHostOperation(op.id)?.payload).toEqual(payload);
+    // Previous releases receive their original inline protocol during a rolling upgrade.
+    hosts.heartbeat(host.id, {});
+    expect(JSON.parse(hosts.getHostOperationPayload(host.id, op.id).toString())).toEqual(payload);
+});
+
+it('backfills legacy bytes once while preserving foreign keys, IDs and pending work', async () => {
+    const legacy = new Database(fixture.dbPath);
+    const content = Buffer.from([0, 255, 1, 88]);
+    const hash = createHash('sha256').update(content).digest('hex');
+    try {
+        legacy.exec(schema.replace(/\b(content|content_before|content_after|prior_content)_hash(\s+)TEXT REFERENCES file_blobs\(sha256\)/g, '$1$2BLOB'));
+        legacy.exec("INSERT INTO scripts(id,name,target_rel_path) VALUES(7,'legacy','plugin.dll')");
+        legacy.prepare('INSERT INTO script_versions(id,script_id,version,sha256,size,content) VALUES(9,7,1,?,?,?)').run(hash, content.length, content);
+        legacy.prepare("INSERT INTO staged_file_changes(source_kind,source_id,server,kind,rel_path,content,inventory_source) VALUES('campaign','old','t5_one','plugin','plugin.dll',?,'test')").run(content);
+        legacy.exec("INSERT INTO deployments(id,action,started_at) VALUES('old','promote','2026-10-01T00:00:00.000Z')");
+        legacy.prepare("INSERT INTO deployment_files(deployment_id,server,verb,rel_path,content_before,content_after) VALUES('old','t5_one','updated','plugin.dll',?,?)").run(content, content);
+    } finally { legacy.close(); }
+    const { db, closeDb } = await import('../../src/lib/server/db/index.js');
+    const { getVersionContent } = await import('../../src/lib/server/fleet/scripts.js');
+    for (let pass = 0; pass < 2; pass++) {
+        expect(getVersionContent(9)).toEqual(content);
+        expect(db().prepare('SELECT content_hash,state FROM staged_file_changes').get()).toEqual({ content_hash: hash, state: 'pending' });
+        expect(db().prepare('SELECT COUNT(*) n FROM file_blobs').get()).toEqual({ n: 1 });
+        expect(db().pragma('foreign_key_check')).toEqual([]);
+        expect((db().pragma('table_info(script_versions)') as { name: string }[]).some(column => column.name === 'content')).toBe(false);
+        closeDb();
+    }
+});
+
+it('migrates legacy command and receipt bytes without changing replay identity', async () => {
+    const { db, closeDb } = await import('../../src/lib/server/db/index.js');
+    const hosts = await import('../../src/lib/server/hosts/index.js');
+    const host = hosts.createHost({ name: 'Legacy' });
+    hosts.enrollHost(hosts.createEnrollment(host.id).token, {});
+    const payload = { stageId: 'legacy', files: [{ path: 'plugin.dll', contentBase64: 'AAEC' }] };
+    const result = { before: [{ path: 'plugin.dll', contentBase64: 'AAEC' }] };
+    db().prepare(`INSERT INTO ec_host_operations(id,host_id,kind,idempotency_key,payload,state,result,created_at,updated_at)
+        VALUES('legacy',?,'files.stage','legacy',?,'running',?, ?,?)`).run(host.id, JSON.stringify(payload), JSON.stringify(result), new Date().toISOString(), new Date().toISOString());
+    db().prepare("DELETE FROM ec_host_schema_migrations WHERE name IN ('staged-operation-payloads','operation-file-blobs')").run();
+    closeDb();
+    expect(hosts.getHostOperation('legacy')).toMatchObject({ payload, result, state: 'running' });
+    expect(hosts.queueHostOperation(host.id, 'files.stage', payload, { idempotencyKey: 'legacy' }).id).toBe('legacy');
+    expect(db().prepare('SELECT COUNT(*) n FROM file_blobs').get()).toEqual({ n: 1 });
+    expect(hosts.getHostOperationStatus('legacy')?.result).toEqual({ before: [{ path: 'plugin.dll', contentSha256: createHash('sha256').update(Buffer.from('AAEC', 'base64')).digest('hex') }] });
+});