R524809
100 lines
9.7 KiB
TypeScript
100 lines
9.7 KiB
TypeScript
import { createHash } from 'node:crypto';
|
||
import { lookup } from 'node:dns/promises';
|
||
import { isIP } from 'node:net';
|
||
import { CreateBucketCommand, HeadBucketCommand, PutObjectCommand, type S3Client } from '@aws-sdk/client-s3';
|
||
import type { Database } from '@atlas/database';
|
||
import { workerConfig } from './config.js';
|
||
import type { OutboxJob } from './outbox-dispatcher.js';
|
||
|
||
const MAX_IMAGE_BYTES = 20 * 1024 * 1024;
|
||
const MAX_VIDEO_BYTES = 200 * 1024 * 1024;
|
||
type Transfer = { id: string; source_url: string; source_variant_id: string | null; source_snapshot_revision: string | number; media_type: 'image' | 'video' };
|
||
|
||
/** Downloads selected source media only after a public-address check. A failed
|
||
* transfer is terminal for that candidate, not for the product collection transaction. */
|
||
export async function materializeCollectionAssets(database: Database, storage: S3Client, bucket: string, job: OutboxJob) {
|
||
await ensureBucket(storage, bucket);
|
||
const snapshotRevision = Number(job.payload.source_snapshot_revision ?? 1);
|
||
if (!Number.isSafeInteger(snapshotRevision) || snapshotRevision < 1) return;
|
||
const transfers = await database.workspaceTransaction(job.workspaceId, async (tx) => {
|
||
const rows = await tx.query<Transfer>(`UPDATE atlas.collection_asset_transfers SET status = 'running', updated_at = now()
|
||
WHERE id IN (SELECT id FROM atlas.collection_asset_transfers WHERE collection_item_id = $1 AND source_snapshot_revision = $2 AND status = 'queued' ORDER BY created_at FOR UPDATE SKIP LOCKED)
|
||
RETURNING id, source_url, source_variant_id, source_snapshot_revision, media_type::text AS media_type`, [String(job.payload.collection_item_id), snapshotRevision]);
|
||
return rows.rows;
|
||
});
|
||
for (const transfer of transfers) await materializeOne(database, storage, bucket, job.workspaceId, transfer);
|
||
await database.workspaceTransaction(job.workspaceId, async (tx) => {
|
||
await tx.query(`UPDATE atlas.collection_items SET status = CASE WHEN EXISTS (SELECT 1 FROM atlas.collection_asset_transfers WHERE collection_item_id = $1 AND source_snapshot_revision = $2 AND status = 'running') THEN 'materializing'::atlas.collection_item_status ELSE 'ready_for_edit'::atlas.collection_item_status END, updated_at = now(), revision = revision + 1 WHERE id = $1 AND source_snapshot_revision = $2 AND deleted_at IS NULL`, [String(job.payload.collection_item_id), snapshotRevision]);
|
||
});
|
||
}
|
||
|
||
async function materializeOne(database: Database, storage: S3Client, bucket: string, workspaceId: string, transfer: Transfer) {
|
||
const snapshotRevision = Number(transfer.source_snapshot_revision);
|
||
try {
|
||
const downloaded = await downloadMedia(transfer.source_url, transfer.media_type);
|
||
const current = await database.workspaceTransaction(workspaceId, async (tx) => (await tx.query<{ id: string }>(`SELECT t.id FROM atlas.collection_asset_transfers t JOIN atlas.collection_items c ON c.id = t.collection_item_id
|
||
WHERE t.id = $1 AND t.status = 'running' AND t.source_snapshot_revision = $2 AND c.source_snapshot_revision = $2`, [transfer.id, snapshotRevision])).rows[0]);
|
||
if (!current) return;
|
||
const checksum = createHash('sha256').update(downloaded.bytes).digest('hex');
|
||
const objectKey = `collections/${workspaceId}/${checksum}`;
|
||
await storage.send(new PutObjectCommand({ Bucket: bucket, Key: objectKey, Body: downloaded.bytes, ContentType: downloaded.contentType }));
|
||
await database.workspaceTransaction(workspaceId, async (tx) => {
|
||
const existing = await tx.query<{ id: string }>('SELECT id FROM atlas.media_assets WHERE checksum_sha256 = $1', [checksum]);
|
||
const assetId = existing.rows[0]?.id;
|
||
if (assetId) {
|
||
await tx.query(`UPDATE atlas.collection_asset_transfers SET status = 'succeeded', asset_id = $2, completed_at = now(), updated_at = now(), error_payload = NULL WHERE id = $1 AND status = 'running' AND source_snapshot_revision = $3`, [transfer.id, assetId, snapshotRevision]);
|
||
return;
|
||
}
|
||
const asset = await tx.query<{ id: string }>(`INSERT INTO atlas.media_assets (workspace_id, storage_provider, bucket, object_key, content_type, byte_size, checksum_sha256, width_px, height_px)
|
||
VALUES (current_setting('app.workspace_id')::uuid, $8, $1, $2, $3, $4, $5, $6, $7) RETURNING id`, [bucket, objectKey, downloaded.contentType, downloaded.bytes.byteLength, checksum, downloaded.width, downloaded.height, workerConfig.objectStorageProvider]);
|
||
await tx.query(`UPDATE atlas.collection_asset_transfers SET status = 'succeeded', asset_id = $2, completed_at = now(), updated_at = now(), error_payload = NULL WHERE id = $1 AND status = 'running' AND source_snapshot_revision = $3`, [transfer.id, asset.rows[0].id, snapshotRevision]);
|
||
});
|
||
} catch (error) {
|
||
await database.workspaceTransaction(workspaceId, (tx) => tx.query(`UPDATE atlas.collection_asset_transfers SET status = 'failed', completed_at = now(), updated_at = now(), error_payload = jsonb_build_object('code', 'SOURCE_ASSET_DOWNLOAD_FAILED', 'message', $2) WHERE id = $1 AND status = 'running' AND source_snapshot_revision = $3`, [transfer.id, error instanceof Error ? error.message.slice(0, 300) : '素材下载失败', snapshotRevision]));
|
||
}
|
||
}
|
||
|
||
async function downloadMedia(source: string, mediaType: Transfer['media_type']) {
|
||
let url = new URL(source);
|
||
if (mediaType === 'video' && /\.m3u8(?:$|[?#])/i.test(url.pathname))
|
||
throw new Error('暂不支持 HLS 视频转存');
|
||
for (let redirect = 0; redirect < 4; redirect += 1) {
|
||
await assertPublicUrl(url);
|
||
const response = await fetch(url, { redirect: 'manual', signal: AbortSignal.timeout(60_000), headers: { Accept: mediaType === 'video' ? 'video/mp4,video/webm,*/*;q=0.8' : 'image/avif,image/webp,image/png,image/jpeg,*/*;q=0.8', 'User-Agent': 'Atlas-Collector/0.1' } });
|
||
if ([301, 302, 303, 307, 308].includes(response.status)) {
|
||
const location = response.headers.get('location'); if (!location) throw new Error('来源素材重定向缺少地址'); url = new URL(location, url); continue;
|
||
}
|
||
if (!response.ok || !response.body) throw new Error(`来源素材请求失败（HTTP ${response.status}）`);
|
||
const chunks: Buffer[] = []; let total = 0;
|
||
const limit = mediaType === 'video' ? MAX_VIDEO_BYTES : MAX_IMAGE_BYTES;
|
||
for await (const chunk of response.body as AsyncIterable<Uint8Array>) { total += chunk.byteLength; if (total > limit) throw new Error(mediaType === 'video' ? '来源视频超过 200MB' : '来源图片超过 20MB'); chunks.push(Buffer.from(chunk)); }
|
||
const bytes = Buffer.concat(chunks, total);
|
||
const media = mediaType === 'video' ? detectVideo(bytes) : detectImage(bytes);
|
||
if (!media) throw new Error(mediaType === 'video' ? '来源素材不是支持的 MP4 或 WebM 视频' : '来源素材不是支持的 PNG、JPEG 或 WebP 图片');
|
||
return { bytes, ...media };
|
||
}
|
||
throw new Error('来源素材重定向次数过多');
|
||
}
|
||
async function assertPublicUrl(url: URL) {
|
||
if (url.protocol !== 'https:') throw new Error('来源素材必须使用 HTTPS');
|
||
if (url.username || url.password || url.port || url.hostname === 'localhost' || isPrivateIp(url.hostname)) throw new Error('来源素材地址不可访问');
|
||
const addresses = await lookup(url.hostname, { all: true, verbatim: true }); if (!addresses.length || addresses.some((entry) => isPrivateIp(entry.address))) throw new Error('来源素材地址不可访问');
|
||
}
|
||
function isPrivateIp(value: string) {
|
||
if (isIP(value) === 4) return /^(10\.|127\.|0\.|169\.254\.|192\.168\.|172\.(1[6-9]|2\d|3[01])\.)/.test(value);
|
||
return value === '::1' || value.toLowerCase().startsWith('fc') || value.toLowerCase().startsWith('fd') || value.toLowerCase().startsWith('fe80:');
|
||
}
|
||
async function ensureBucket(storage: S3Client, bucket: string) { try { await storage.send(new HeadBucketCommand({ Bucket: bucket })); } catch { await storage.send(new CreateBucketCommand({ Bucket: bucket })); } }
|
||
function detectImage(bytes: Buffer): { contentType: 'image/png' | 'image/jpeg' | 'image/webp'; width: number; height: number } | null {
|
||
if (bytes.length >= 24 && bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]))) return dimensions('image/png', bytes.readUInt32BE(16), bytes.readUInt32BE(20));
|
||
if (bytes.length >= 30 && bytes.subarray(0, 4).toString() === 'RIFF' && bytes.subarray(8, 12).toString() === 'WEBP' && bytes.subarray(12, 16).toString() === 'VP8X') return dimensions('image/webp', 1 + bytes.readUIntLE(24, 3), 1 + bytes.readUIntLE(27, 3));
|
||
if (bytes.length > 9 && bytes[0] === 0xff && bytes[1] === 0xd8) for (let offset = 2; offset + 9 < bytes.length;) { if (bytes[offset] !== 0xff) return null; const marker = bytes[offset + 1], length = bytes.readUInt16BE(offset + 2); if (length < 2 || offset + 2 + length > bytes.length) return null; if ([0xc0, 0xc1, 0xc2, 0xc3, 0xc5, 0xc6, 0xc7, 0xc9, 0xca, 0xcb, 0xcd, 0xce, 0xcf].includes(marker)) return dimensions('image/jpeg', bytes.readUInt16BE(offset + 7), bytes.readUInt16BE(offset + 5)); offset += length + 2; }
|
||
return null;
|
||
}
|
||
function detectVideo(bytes: Buffer): { contentType: 'video/mp4' | 'video/webm'; width: null; height: null } | null {
|
||
if (bytes.length >= 12 && bytes.subarray(4, 8).toString('ascii') === 'ftyp') return { contentType: 'video/mp4', width: null, height: null };
|
||
if (bytes.length >= 4 && bytes.subarray(0, 4).equals(Buffer.from([0x1a, 0x45, 0xdf, 0xa3]))) return { contentType: 'video/webm', width: null, height: null };
|
||
return null;
|
||
}
|
||
function dimensions(contentType: 'image/png' | 'image/jpeg' | 'image/webp', width: number, height: number) { return width >= 1 && height >= 1 && width <= 8192 && height <= 8192 ? { contentType, width, height } : null; }
|