import req from './req.js'; import CryptoJS from 'crypto-js'; import {join} from 'path'; import fs from 'fs'; import {PassThrough} from 'stream'; export async function testSupport(url, headers) { const resp = await req .get(url, { responseType: 'stream', headers: Object.assign( { Range: 'bytes=0-0', }, headers, ), }) .catch((err) => { console.error(err); return err.response || {status: 500, data: {}}; }); if (resp && resp.status === 206) { const isAccept = resp.headers['accept-ranges'] === 'bytes'; const contentRange = resp.headers['content-range']; const contentLength = parseInt(resp.headers['content-length']); const isSupport = isAccept || !!contentRange || contentLength === 1; const length = contentRange ? parseInt(contentRange.split('/')[1]) : contentLength; delete resp.headers['content-range']; delete resp.headers['content-length']; if (length) resp.headers['content-length'] = length.toString(); return [isSupport, resp.headers]; } else { return [false, null]; } } const urlHeadCache = {}; let currentUrlKey = ''; const cacheRoot = (process.env['NODE_PATH'] || '.') + '/vod_cache'; const maxCache = 1024 * 1024 * 100; function delAllCache(keepKey) { try { fs.readdir(cacheRoot, (_, files) => { if (files) for (const file of files) { if (file === keepKey) continue; const dir = join(cacheRoot, file); fs.stat(dir, (_, stats) => { if (stats && stats.isDirectory()) { fs.readdir(dir, (_, subFiles) => { if (subFiles) for (const subFile of subFiles) { if (!subFile.endsWith('.p')) { fs.rm(join(dir, subFile), {recursive: true}, () => { }); } } }); } }); } }); } catch (error) { console.error(error); } } async function chunkStream(inReq, outResp, url, urlKey, headers, option) { urlKey = urlKey || CryptoJS.enc.Hex.stringify(CryptoJS.MD5(url)).toString(); if (currentUrlKey !== urlKey) { delAllCache(urlKey); currentUrlKey = urlKey; } if (!urlHeadCache[urlKey]) { const [isSupport, urlHeader] = await testSupport(url, headers); if (!isSupport || !urlHeader['content-length']) { console.log(`[chunkStream] 获取content-length失败,执行重定向到: ${url}`); outResp.redirect(url); return; } urlHeadCache[urlKey] = urlHeader; } let exist = true; await fs.promises.access(join(cacheRoot, urlKey)).catch((_) => (exist = false)); if (!exist) { await fs.promises.mkdir(join(cacheRoot, urlKey), {recursive: true}); } const contentLength = parseInt(urlHeadCache[urlKey]['content-length']); let byteStart = 0; let byteEnd = contentLength - 1; const streamHeader = {}; if (inReq.headers.range) { // console.log(inReq.id, inReq.headers.range); const ranges = inReq.headers.range.trim().split(/=|-/); if (ranges.length > 2 && ranges[2]) { byteEnd = parseInt(ranges[2]); } byteStart = parseInt(ranges[1]); Object.assign(streamHeader, urlHeadCache[urlKey]); streamHeader['content-length'] = (byteEnd - byteStart + 1).toString(); streamHeader['content-range'] = `bytes ${byteStart}-${byteEnd}/${contentLength}`; // console.log('streamHeader:', streamHeader); outResp.code(206); } else { Object.assign(streamHeader, urlHeadCache[urlKey]); outResp.code(200); } option = option || {chunkSize: 1024 * 256, poolSize: 5, timeout: 1000 * 10}; // console.log(`[chunkStream] option: `, option); const chunkSize = option.chunkSize; const poolSize = option.poolSize; const timeout = option.timeout; let chunkCount = Math.ceil(contentLength / chunkSize); let chunkDownIdx = Math.floor(byteStart / chunkSize); let chunkReadIdx = chunkDownIdx; let stop = false; const dlFiles = {}; for (let i = 0; i < poolSize && i < chunkCount; i++) { new Promise((resolve) => { (async function doDLTask(spChunkIdx) { if (stop || chunkDownIdx >= chunkCount) { resolve(); return; } if (spChunkIdx === undefined && (chunkDownIdx - chunkReadIdx) * chunkSize >= maxCache) { setTimeout(doDLTask, 5); return; } const chunkIdx = spChunkIdx || chunkDownIdx++; const taskId = `${inReq.id}-${chunkIdx}`; try { const dlFile = join(cacheRoot, urlKey, `${inReq.id}-${chunkIdx}.p`); let exist = true; await fs.promises.access(dlFile).catch((_) => (exist = false)); if (!exist) { const start = chunkIdx * chunkSize; const end = Math.min(contentLength - 1, (chunkIdx + 1) * chunkSize - 1); // console.log('[chunkIdx]:', inReq.id, chunkIdx); const dlResp = await req.get(url, { responseType: 'stream', timeout: timeout, headers: Object.assign( { Range: `bytes=${start}-${end}`, }, headers, ), }); const dlCache = join(cacheRoot, urlKey, `${inReq.id}-${chunkIdx}.dl`); const writer = fs.createWriteStream(dlCache); const readTimeout = setTimeout(() => { writer.destroy(new Error(`${taskId} read timeout`)); }, timeout); const downloaded = new Promise((resolve) => { writer.on('finish', async () => { if (stop) { await fs.promises.rm(dlCache).catch((e) => console.error(e)); } else { await fs.promises.rename(dlCache, dlFile).catch((e) => console.error(e)); dlFiles[taskId] = dlFile; } resolve(true); }); writer.on('error', async (e) => { console.error(e); await fs.promises.rm(dlCache).catch((e1) => console.error(e1)); resolve(false); }); }); dlResp.data.pipe(writer); const result = await downloaded; clearTimeout(readTimeout); if (!result) { setTimeout(() => { doDLTask(chunkIdx); }, 15); return; } } setTimeout(doDLTask, 5); } catch (error) { console.error(error); setTimeout(() => { doDLTask(chunkIdx); }, 15); } })(); }); } outResp.headers(streamHeader); const stream = new PassThrough(); new Promise((resolve) => { let writeMore = true; (async function waitReadFile() { try { if (chunkReadIdx >= chunkCount || stop) { stream.end(); resolve(); return; } if (!writeMore) { setTimeout(waitReadFile, 5); return; } const taskId = `${inReq.id}-${chunkReadIdx}`; if (!dlFiles[taskId]) { setTimeout(waitReadFile, 5); return; } const chunkByteStart = chunkReadIdx * chunkSize; const chunkByteEnd = Math.min(contentLength - 1, (chunkReadIdx + 1) * chunkSize - 1); const readFileStart = Math.max(byteStart, chunkByteStart) - chunkByteStart; const dlFile = dlFiles[taskId]; delete dlFiles[taskId]; const fd = await fs.promises.open(dlFile, 'r'); const buffer = Buffer.alloc(chunkByteEnd - chunkByteStart - readFileStart + 1); await fd.read(buffer, 0, chunkByteEnd - chunkByteStart - readFileStart + 1, readFileStart); await fd.close().catch((e) => console.error(e)); await fs.promises.rm(dlFile).catch((e) => console.error(e)); writeMore = stream.write(buffer); if (!writeMore) { stream.once('drain', () => { writeMore = true; }); } chunkReadIdx++; setTimeout(waitReadFile, 5); } catch (error) { setTimeout(waitReadFile, 5); } })(); }); stream.on('close', async () => { Object.keys(dlFiles).forEach((reqKey) => { if (reqKey.startsWith(inReq.id)) { fs.rm(dlFiles[reqKey], {recursive: true}, () => { }); delete dlFiles[reqKey]; } }); stop = true; }); return stream; } export default chunkStream;