mirror of
https://github.com/NapNeko/NapCatQQ.git
synced 2026-02-05 23:19:37 +00:00
fe878687ee
Introduces new OneBot actions for streaming file upload and download, including chunked file transfer with memory/disk management and SHA256 verification. Adds CleanStreamTempFile, DownloadFileStream, UploadFileStream, and TestStreamDownload actions, updates action routing and network adapters to support streaming via HTTP and WebSocket, and provides Python test scripts for concurrent upload testing.
385 lines
13 KiB
TypeScript
385 lines
13 KiB
TypeScript
import { ActionName } from '@/onebot/action/router';
|
||
import { OneBotAction } from '@/onebot/action/OneBotAction';
|
||
import { Static, Type } from '@sinclair/typebox';
|
||
import fs from 'fs';
|
||
import { join as joinPath } from 'node:path';
|
||
import { randomUUID } from 'crypto';
|
||
import { createHash } from 'crypto';
|
||
|
||
// 简化配置
|
||
const CONFIG = {
|
||
TIMEOUT: 10 * 60 * 1000, // 10分钟超时
|
||
MEMORY_THRESHOLD: 10 * 1024 * 1024, // 10MB，超过使用磁盘
|
||
MEMORY_LIMIT: 100 * 1024 * 1024 // 100MB内存总限制
|
||
} as const;
|
||
|
||
const SchemaData = Type.Object({
|
||
stream_id: Type.String(),
|
||
chunk_data: Type.Optional(Type.String()),
|
||
chunk_index: Type.Optional(Type.Number()),
|
||
total_chunks: Type.Optional(Type.Number()),
|
||
file_size: Type.Optional(Type.Number()),
|
||
expected_sha256: Type.Optional(Type.String()),
|
||
is_complete: Type.Optional(Type.Boolean()),
|
||
filename: Type.Optional(Type.String()),
|
||
reset: Type.Optional(Type.Boolean()),
|
||
verify_only: Type.Optional(Type.Boolean())
|
||
});
|
||
|
||
type Payload = Static<typeof SchemaData>;
|
||
|
||
// 简化流状态接口
|
||
interface StreamState {
|
||
id: string;
|
||
filename: string;
|
||
totalChunks: number;
|
||
receivedChunks: number;
|
||
missingChunks: Set<number>;
|
||
|
||
// 可选属性
|
||
fileSize?: number;
|
||
expectedSha256?: string;
|
||
|
||
// 存储策略
|
||
useMemory: boolean;
|
||
memoryChunks?: Map<number, Buffer>;
|
||
tempDir?: string;
|
||
finalPath?: string;
|
||
|
||
// 管理
|
||
createdAt: number;
|
||
timeoutId: NodeJS.Timeout;
|
||
}
|
||
|
||
interface StreamResult {
|
||
stream_id: string;
|
||
status: 'receiving' | 'completed' | 'error' | 'ready';
|
||
received_chunks: number;
|
||
total_chunks?: number;
|
||
missing_chunks?: number[];
|
||
file_path?: string;
|
||
file_size?: number;
|
||
sha256?: string;
|
||
message?: string;
|
||
}
|
||
|
||
export class UploadFileStream extends OneBotAction<Payload, StreamResult> {
|
||
override actionName = ActionName.UploadFileStream;
|
||
override payloadSchema = SchemaData;
|
||
|
||
private static streams = new Map<string, StreamState>();
|
||
private static memoryUsage = 0;
|
||
|
||
async _handle(payload: Payload): Promise<StreamResult> {
|
||
const { stream_id, reset, verify_only } = payload;
|
||
|
||
try {
|
||
if (reset) return this.resetStream(stream_id);
|
||
if (verify_only) return this.verifyStream(stream_id);
|
||
|
||
const stream = this.getOrCreateStream(payload);
|
||
|
||
if (payload.chunk_data && payload.chunk_index !== undefined) {
|
||
const result = await this.processChunk(stream, payload.chunk_data, payload.chunk_index);
|
||
if (result.status === 'error') return result;
|
||
}
|
||
|
||
if (payload.is_complete || stream.receivedChunks === stream.totalChunks) {
|
||
return await this.completeStream(stream);
|
||
}
|
||
|
||
return this.getStreamStatus(stream);
|
||
|
||
} catch (error) {
|
||
// 确保在任何错误情况下都清理资源
|
||
this.cleanupStream(stream_id, true);
|
||
return this.errorResult(stream_id, error);
|
||
}
|
||
}
|
||
|
||
private resetStream(streamId: string): StreamResult {
|
||
this.cleanupStream(streamId);
|
||
return {
|
||
stream_id: streamId,
|
||
status: 'ready',
|
||
received_chunks: 0,
|
||
message: 'Stream reset'
|
||
};
|
||
}
|
||
|
||
private verifyStream(streamId: string): StreamResult {
|
||
const stream = UploadFileStream.streams.get(streamId);
|
||
if (!stream) {
|
||
return this.errorResult(streamId, new Error('Stream not found'));
|
||
}
|
||
return this.getStreamStatus(stream);
|
||
}
|
||
|
||
private getOrCreateStream(payload: Payload): StreamState {
|
||
let stream = UploadFileStream.streams.get(payload.stream_id);
|
||
|
||
if (!stream) {
|
||
if (!payload.total_chunks) {
|
||
throw new Error('total_chunks required for new stream');
|
||
}
|
||
stream = this.createStream(payload);
|
||
}
|
||
|
||
return stream;
|
||
}
|
||
|
||
private createStream(payload: Payload): StreamState {
|
||
const { stream_id, total_chunks, file_size, filename, expected_sha256 } = payload;
|
||
|
||
const useMemory = this.shouldUseMemory(file_size);
|
||
if (useMemory && file_size && (UploadFileStream.memoryUsage + file_size) > CONFIG.MEMORY_LIMIT) {
|
||
throw new Error('Memory limit exceeded');
|
||
}
|
||
|
||
const stream: StreamState = {
|
||
id: stream_id,
|
||
filename: filename || `upload_${randomUUID()}`,
|
||
totalChunks: total_chunks!,
|
||
receivedChunks: 0,
|
||
missingChunks: new Set(Array.from({ length: total_chunks! }, (_, i) => i)),
|
||
fileSize: file_size,
|
||
expectedSha256: expected_sha256,
|
||
useMemory,
|
||
createdAt: Date.now(),
|
||
timeoutId: this.setupTimeout(stream_id)
|
||
};
|
||
|
||
try {
|
||
if (useMemory) {
|
||
stream.memoryChunks = new Map();
|
||
if (file_size) UploadFileStream.memoryUsage += file_size;
|
||
} else {
|
||
this.setupDiskStorage(stream);
|
||
}
|
||
|
||
UploadFileStream.streams.set(stream_id, stream);
|
||
return stream;
|
||
} catch (error) {
|
||
// 如果设置存储失败，清理已创建的资源
|
||
clearTimeout(stream.timeoutId);
|
||
if (stream.tempDir && fs.existsSync(stream.tempDir)) {
|
||
try {
|
||
fs.rmSync(stream.tempDir, { recursive: true, force: true });
|
||
} catch (cleanupError) {
|
||
console.error(`Failed to cleanup temp dir during creation error:`, cleanupError);
|
||
}
|
||
}
|
||
throw error;
|
||
}
|
||
}
|
||
|
||
private shouldUseMemory(fileSize?: number): boolean {
|
||
return fileSize !== undefined && fileSize <= CONFIG.MEMORY_THRESHOLD;
|
||
}
|
||
|
||
private setupDiskStorage(stream: StreamState): void {
|
||
const tempDir = joinPath(this.core.NapCatTempPath, `upload_${stream.id}`);
|
||
const finalPath = joinPath(this.core.NapCatTempPath, stream.filename);
|
||
|
||
fs.mkdirSync(tempDir, { recursive: true });
|
||
|
||
stream.tempDir = tempDir;
|
||
stream.finalPath = finalPath;
|
||
}
|
||
|
||
private setupTimeout(streamId: string): NodeJS.Timeout {
|
||
return setTimeout(() => {
|
||
console.log(`Stream ${streamId} timeout`);
|
||
this.cleanupStream(streamId);
|
||
}, CONFIG.TIMEOUT);
|
||
}
|
||
|
||
private async processChunk(stream: StreamState, chunkData: string, chunkIndex: number): Promise<StreamResult> {
|
||
// 验证索引
|
||
if (chunkIndex < 0 || chunkIndex >= stream.totalChunks) {
|
||
throw new Error(`Invalid chunk index: ${chunkIndex}`);
|
||
}
|
||
|
||
// 检查重复
|
||
if (!stream.missingChunks.has(chunkIndex)) {
|
||
return this.getStreamStatus(stream, `Chunk ${chunkIndex} already received`);
|
||
}
|
||
|
||
try {
|
||
const buffer = Buffer.from(chunkData, 'base64');
|
||
|
||
// 存储分片
|
||
if (stream.useMemory) {
|
||
stream.memoryChunks!.set(chunkIndex, buffer);
|
||
} else {
|
||
const chunkPath = joinPath(stream.tempDir!, `${chunkIndex}.chunk`);
|
||
await fs.promises.writeFile(chunkPath, buffer);
|
||
}
|
||
|
||
// 更新状态
|
||
stream.missingChunks.delete(chunkIndex);
|
||
stream.receivedChunks++;
|
||
this.refreshTimeout(stream);
|
||
|
||
return this.getStreamStatus(stream);
|
||
|
||
} catch (error) {
|
||
throw new Error(`Chunk processing failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
|
||
}
|
||
}
|
||
|
||
private refreshTimeout(stream: StreamState): void {
|
||
clearTimeout(stream.timeoutId);
|
||
stream.timeoutId = this.setupTimeout(stream.id);
|
||
}
|
||
|
||
private getStreamStatus(stream: StreamState, message?: string): StreamResult {
|
||
const missingChunks = Array.from(stream.missingChunks).sort();
|
||
|
||
return {
|
||
stream_id: stream.id,
|
||
status: 'receiving',
|
||
received_chunks: stream.receivedChunks,
|
||
total_chunks: stream.totalChunks,
|
||
missing_chunks: missingChunks.length > 0 ? missingChunks : undefined,
|
||
file_size: stream.fileSize,
|
||
message
|
||
};
|
||
}
|
||
|
||
private async completeStream(stream: StreamState): Promise<StreamResult> {
|
||
try {
|
||
// 合并分片
|
||
const finalBuffer = stream.useMemory ?
|
||
await this.mergeMemoryChunks(stream) :
|
||
await this.mergeDiskChunks(stream);
|
||
|
||
// 验证SHA256
|
||
const sha256 = this.validateSha256(stream, finalBuffer);
|
||
|
||
// 保存文件
|
||
const finalPath = stream.finalPath || joinPath(this.core.NapCatTempPath, stream.filename);
|
||
await fs.promises.writeFile(finalPath, finalBuffer);
|
||
|
||
// 清理资源但保留文件
|
||
this.cleanupStream(stream.id, false);
|
||
|
||
return {
|
||
stream_id: stream.id,
|
||
status: 'completed',
|
||
received_chunks: stream.receivedChunks,
|
||
total_chunks: stream.totalChunks,
|
||
file_path: finalPath,
|
||
file_size: finalBuffer.length,
|
||
sha256,
|
||
message: 'Upload completed'
|
||
};
|
||
|
||
} catch (error) {
|
||
throw new Error(`Stream completion failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
|
||
}
|
||
}
|
||
|
||
private async mergeMemoryChunks(stream: StreamState): Promise<Buffer> {
|
||
const chunks: Buffer[] = [];
|
||
for (let i = 0; i < stream.totalChunks; i++) {
|
||
const chunk = stream.memoryChunks!.get(i);
|
||
if (!chunk) throw new Error(`Missing memory chunk ${i}`);
|
||
chunks.push(chunk);
|
||
}
|
||
return Buffer.concat(chunks);
|
||
}
|
||
|
||
private async mergeDiskChunks(stream: StreamState): Promise<Buffer> {
|
||
const chunks: Buffer[] = [];
|
||
for (let i = 0; i < stream.totalChunks; i++) {
|
||
const chunkPath = joinPath(stream.tempDir!, `${i}.chunk`);
|
||
if (!fs.existsSync(chunkPath)) throw new Error(`Missing chunk file ${i}`);
|
||
chunks.push(await fs.promises.readFile(chunkPath));
|
||
}
|
||
return Buffer.concat(chunks);
|
||
}
|
||
|
||
private validateSha256(stream: StreamState, buffer: Buffer): string | undefined {
|
||
if (!stream.expectedSha256) return undefined;
|
||
|
||
const actualSha256 = createHash('sha256').update(buffer).digest('hex');
|
||
if (actualSha256 !== stream.expectedSha256) {
|
||
throw new Error(`SHA256 mismatch. Expected: ${stream.expectedSha256}, Got: ${actualSha256}`);
|
||
}
|
||
return actualSha256;
|
||
}
|
||
|
||
private cleanupStream(streamId: string, deleteFinalFile = true): void {
|
||
const stream = UploadFileStream.streams.get(streamId);
|
||
if (!stream) return;
|
||
|
||
try {
|
||
// 清理超时
|
||
clearTimeout(stream.timeoutId);
|
||
|
||
// 清理内存
|
||
if (stream.useMemory) {
|
||
if (stream.fileSize) {
|
||
UploadFileStream.memoryUsage = Math.max(0, UploadFileStream.memoryUsage - stream.fileSize);
|
||
}
|
||
stream.memoryChunks?.clear();
|
||
}
|
||
|
||
// 清理临时文件夹及其所有内容
|
||
if (stream.tempDir) {
|
||
try {
|
||
if (fs.existsSync(stream.tempDir)) {
|
||
fs.rmSync(stream.tempDir, { recursive: true, force: true });
|
||
console.log(`Cleaned up temp directory: ${stream.tempDir}`);
|
||
}
|
||
} catch (error) {
|
||
console.error(`Failed to cleanup temp directory ${stream.tempDir}:`, error);
|
||
}
|
||
}
|
||
|
||
// 删除最终文件（如果需要）
|
||
if (deleteFinalFile && stream.finalPath) {
|
||
try {
|
||
if (fs.existsSync(stream.finalPath)) {
|
||
fs.unlinkSync(stream.finalPath);
|
||
console.log(`Deleted final file: ${stream.finalPath}`);
|
||
}
|
||
} catch (error) {
|
||
console.error(`Failed to delete final file ${stream.finalPath}:`, error);
|
||
}
|
||
}
|
||
|
||
} catch (error) {
|
||
console.error(`Cleanup error for stream ${streamId}:`, error);
|
||
} finally {
|
||
UploadFileStream.streams.delete(streamId);
|
||
console.log(`Stream ${streamId} cleaned up`);
|
||
}
|
||
}
|
||
|
||
private errorResult(streamId: string, error: any): StreamResult {
|
||
return {
|
||
stream_id: streamId,
|
||
status: 'error',
|
||
received_chunks: 0,
|
||
message: error instanceof Error ? error.message : 'Unknown error'
|
||
};
|
||
}
|
||
|
||
// 全局状态查询
|
||
static getGlobalStatus() {
|
||
return {
|
||
activeStreams: this.streams.size,
|
||
memoryUsageMB: Math.round(this.memoryUsage / 1024 / 1024 * 100) / 100,
|
||
streams: Array.from(this.streams.values()).map(stream => ({
|
||
streamId: stream.id,
|
||
filename: stream.filename,
|
||
progress: `${stream.receivedChunks}/${stream.totalChunks}`,
|
||
useMemory: stream.useMemory,
|
||
createdAt: new Date(stream.createdAt).toISOString()
|
||
}))
|
||
};
|
||
}
|
||
}
|