@nightowne/tas-cli 2.0.0 → 2.3.0

package/src/index.js

@@ -5,9 +5,10 @@
 
 import fs from 'fs';
 import path from 'path';
+import { pipeline } from 'stream/promises';
 import { Encryptor, hashFile } from './crypto/encryption.js';
 import { Compressor } from './utils/compression.js';
-import { Chunker, createHeader, parseHeader, HEADER_SIZE } from './utils/chunker.js';
+import { createHeader, parseHeader, HEADER_SIZE } from './utils/chunker.js';
 import { TelegramClient } from './telegram/client.js';
 import { FileIndex } from './db/index.js';
 
@@ -23,10 +24,10 @@ export async function processFile(filePath, options) {
 
     onProgress?.('Reading file...');
 
-    // Read file
+    // Read file initially just to get size
     const filename = customName || path.basename(filePath);
-    const fileData = fs.readFileSync(filePath);
-    const originalSize = fileData.length;
+    const stats = fs.statSync(filePath);
+    const originalSize = stats.size;
 
     // Calculate hash
     onProgress?.('Calculating hash...');
@@ -41,115 +42,169 @@ export async function processFile(filePath, options) {
         throw new Error('File already uploaded (duplicate hash)');
     }
 
-    // Compress
-    onProgress?.('Compressing...');
+    // Prepare processing components
     const compressor = new Compressor();
-    const { data: compressedData, compressed } = await compressor.compress(fileData, filename);
+    const { stream: compressStream, compressed } = compressor.getCompressStream(filename);
+    const flags = compressed ? 1 : 0;
 
-    // Encrypt
-    onProgress?.('Encrypting...');
     const encryptor = new Encryptor(password);
-    const encryptedData = encryptor.encrypt(compressedData);
-
-    // Chunk if needed (Telegram bot limit ~50MB per file)
-    onProgress?.('Preparing chunks...');
-    const chunks = [];
-    const numChunks = Math.ceil(encryptedData.length / TELEGRAM_CHUNK_SIZE);
-
-    for (let i = 0; i < numChunks; i++) {
-        const start = i * TELEGRAM_CHUNK_SIZE;
-        const end = Math.min(start + TELEGRAM_CHUNK_SIZE, encryptedData.length);
-        chunks.push({
-            index: i,
-            total: numChunks,
-            data: encryptedData.subarray(start, end)
-        });
-    }
+    const encryptStream = encryptor.getEncryptStream();
 
-    // Prepare files with headers
-    const tempDir = path.join(dataDir, 'tmp');
+    const tempDir = process.env.TAS_TMP_DIR || path.join(dataDir, 'tmp');
     if (!fs.existsSync(tempDir)) {
         fs.mkdirSync(tempDir, { recursive: true });
     }
 
-    const chunkFiles = [];
-    const flags = compressed ? 1 : 0;
-
-    for (const chunk of chunks) {
-        const header = createHeader(filename, originalSize, chunk.index, chunk.total, flags);
-        const chunkData = Buffer.concat([header, chunk.data]);
-
-        const chunkFilename = chunks.length > 1
-            ? `${hash.substring(0, 12)}.part${chunk.index}.tas`
-            : `${hash.substring(0, 12)}.tas`;
-
-        const chunkPath = path.join(tempDir, chunkFilename);
-        fs.writeFileSync(chunkPath, chunkData);
-
-        chunkFiles.push({
-            index: chunk.index,
-            path: chunkPath,
-            size: chunkData.length
-        });
-    }
-
     // Connect to Telegram
     onProgress?.('Connecting to Telegram...');
     const client = new TelegramClient(dataDir);
     await client.initialize(config.botToken);
     client.setChatId(config.chatId);
 
-    // Upload chunks
-    const fileId = db.addFile({
-        filename,
-        hash,
-        originalSize,
-        storedSize: encryptedData.length,
-        chunks: chunks.length,
-        compressed
-    });
+    // We will stream through a custom Writable chunker
+    const { Writable } = await import('stream');
+
+    // First pass estimation (for calculating total chunks and progress)
+    // We don't know the exact final size due to compression and encryption overhead,
+    // so we'll estimate total chunks and update it if needed.
+    // For small files < 49MB we assume 1 chunk.
+    let estimatedSize = compressed ? originalSize : originalSize + 128; // Add encryption overhead
+    if (compressed && originalSize > 1024 * 1024) estimatedSize = originalSize * 0.8; // Rough guess
+    let estimatedChunks = Math.ceil(estimatedSize / TELEGRAM_CHUNK_SIZE) || 1;
+
+    // Register file in DB inside a transaction for atomicity
+    db.db.exec('BEGIN');
+    let fileId;
+    try {
+        fileId = db.addFile({
+            filename,
+            hash,
+            originalSize,
+            storedSize: 0, // Will update later
+            chunks: estimatedChunks,
+            compressed
+        });
+    } catch (err) {
+        db.db.exec('ROLLBACK');
+        db.close();
+        throw err;
+    }
 
+    onProgress?.('Processing and uploading streams...');
     let uploadedBytes = 0;
-    const totalBytes = chunkFiles.reduce((acc, c) => acc + c.size, 0);
+    let chunkIndex = 0;
+
+    let currentChunkBuffer = Buffer.alloc(0);
+    let totalStoredSize = 0;
+
+    // Helper to upload a single chunk
+    const uploadCurrentChunk = async (isFinal = false) => {
+        if (currentChunkBuffer.length === 0 && !isFinal) return; // Nothing to upload
+        if (currentChunkBuffer.length === 0 && isFinal && chunkIndex > 0) return; // Empty final chunk after perfect split
+
+        // At this point we know if it's the final chunk, so we know the total chunks
+        const totalChunks = isFinal ? chunkIndex + 1 : Math.max(estimatedChunks, chunkIndex + 1);
 
-    for (const chunk of chunkFiles) {
-        onProgress?.(`Uploading chunk ${chunk.index + 1}/${chunkFiles.length}...`);
+        const header = createHeader(filename, originalSize, chunkIndex, totalChunks, flags);
+        const chunkData = Buffer.concat([header, currentChunkBuffer]);
 
-        const caption = chunks.length > 1
-            ? `📦 ${filename} (${chunk.index + 1}/${chunks.length})`
+        const chunkFilename = totalChunks > 1
+            ? `${hash.substring(0, 12)}.part${chunkIndex}.tas`
+            : `${hash.substring(0, 12)}.tas`;
+
+        const chunkPath = path.join(tempDir, chunkFilename);
+        fs.writeFileSync(chunkPath, chunkData);
+
+        const caption = totalChunks > 1
+            ? `📦 ${filename} (${chunkIndex + 1}/${totalChunks})`
             : `📦 ${filename}`;
 
-        const result = await client.sendFile(chunk.path, caption);
+        onProgress?.(`Uploading chunk ${chunkIndex + 1}...`);
+
+        const result = await client.sendFile(chunkPath, caption);
 
-        uploadedBytes += chunk.size;
-        onByteProgress?.({ uploaded: uploadedBytes, total: totalBytes, chunk: chunk.index + 1, totalChunks: chunkFiles.length });
+        uploadedBytes += chunkData.length;
+        totalStoredSize += currentChunkBuffer.length;
 
-        // Store file_id instead of message_id for downloads
-        db.addChunk(fileId, chunk.index, result.messageId.toString(), chunk.size);
+        onByteProgress?.({ uploaded: uploadedBytes, total: estimatedSize, chunk: chunkIndex + 1, totalChunks });
 
-        // Also store file_id for easier retrieval
+        // Store file_id
+        db.addChunk(fileId, chunkIndex, result.messageId.toString(), chunkData.length);
         db.db.prepare('UPDATE chunks SET file_telegram_id = ? WHERE file_id = ? AND chunk_index = ?')
-            .run(result.fileId, fileId, chunk.index);
+            .run(result.fileId, fileId, chunkIndex);
+
+        // Clean up temp file immediately to save disk space
+        fs.unlinkSync(chunkPath);
+
+        chunkIndex++;
+        currentChunkBuffer = Buffer.alloc(0);
+    };
+
+    const chunkingStream = new Writable({
+        async write(chunk, encoding, callback) {
+            currentChunkBuffer = Buffer.concat([currentChunkBuffer, chunk]);
+
+            // If we exceeded the chunk limit, flush it
+            if (currentChunkBuffer.length >= TELEGRAM_CHUNK_SIZE) {
+                const overflow = currentChunkBuffer.subarray(TELEGRAM_CHUNK_SIZE);
+                currentChunkBuffer = currentChunkBuffer.subarray(0, TELEGRAM_CHUNK_SIZE);
+
+                try {
+                    await uploadCurrentChunk(false);
+                    currentChunkBuffer = overflow; // carry over
+                    callback();
+                } catch (err) {
+                    callback(err);
+                }
+            } else {
+                callback();
+            }
+        },
+        async final(callback) {
+            try {
+                await uploadCurrentChunk(true);
+                callback();
+            } catch (err) {
+                callback(err);
+            }
+        }
+    });
+
+    try {
+        const readStream = fs.createReadStream(filePath);
 
-        // Clean up temp file
-        fs.unlinkSync(chunk.path);
+        // Run the pipeline: Read -> Compress -> Encrypt -> Chunk & Upload
+        await pipeline(readStream, compressStream, encryptStream, chunkingStream);
+    } catch (pipelineErr) {
+        // Pipeline failed — roll back the DB transaction so no orphaned rows remain
+        try { db.db.exec('ROLLBACK'); } catch (e) { /* already rolled back */ }
+        db.close();
+        throw pipelineErr;
     }
 
+    // Update the DB with the final accurate values 
+    db.db.prepare('UPDATE files SET stored_size = ?, chunks = ? WHERE id = ?')
+        .run(totalStoredSize, chunkIndex, fileId);
+
+    // Commit the transaction — all DB rows are now permanent
+    db.db.exec('COMMIT');
+
     db.close();
 
-    // Clean up temp dir
+    // Clean up temp dir (only if empty)
     try {
-        fs.rmdirSync(tempDir);
+        const remaining = fs.readdirSync(tempDir);
+        if (remaining.length === 0) fs.rmdirSync(tempDir);
     } catch (e) {
-        // Ignore if not empty
+        // Ignore cleanup errors
     }
 
     return {
         filename,
         hash,
         originalSize,
-        storedSize: encryptedData.length,
-        chunks: chunks.length,
+        storedSize: totalStoredSize,
+        chunks: chunkIndex,
         compressed
     };
 }
@@ -173,57 +228,99 @@ export async function retrieveFile(fileRecord, options) {
         throw new Error('No chunk metadata found for this file');
     }
 
+    // Prepare components
+    const encryptor = new Encryptor(password);
+    const decryptStream = encryptor.getDecryptStream();
+
+    const tempDir = process.env.TAS_TMP_DIR || path.join(dataDir, 'tmp');
+    if (!fs.existsSync(tempDir)) {
+        fs.mkdirSync(tempDir, { recursive: true });
+    }
+
     // Connect to Telegram
     const client = new TelegramClient(dataDir);
     await client.initialize(config.botToken);
     client.setChatId(config.chatId);
 
-    // Download all chunks
-    const downloadedChunks = [];
-    let downloadedBytes = 0;
-    const totalBytes = fileRecord.stored_size || chunks.reduce((acc, c) => acc + (c.size || 0), 0);
+    // Get total size from first chunk's header, or from DB
+    const firstChunkData = await client.downloadFile(chunks[0].file_telegram_id);
+    const header = parseHeader(firstChunkData);
 
-    for (const chunk of chunks) {
-        onProgress?.(`Downloading chunk ${chunk.chunk_index + 1}/${chunks.length}...`);
+    // Total original uncompressed size
+    let expectedOriginalSize = header.originalSize;
+    let wasCompressed = header.compressed;
 
-        const data = await client.downloadFile(chunk.file_telegram_id);
-        downloadedBytes += data.length;
-        onByteProgress?.({ downloaded: downloadedBytes, total: totalBytes, chunk: chunk.chunk_index + 1, totalChunks: chunks.length });
+    const compressor = new Compressor();
+    const decompressStream = compressor.getDecompressStream(wasCompressed);
 
-        // Parse header
-        const header = parseHeader(data);
-        const payload = data.subarray(HEADER_SIZE);
+    // We need a Readable stream that will lazily fetch chunks from Telegram
+    // and push them into the decryption pipeline.
+    const { Readable } = await import('stream');
 
-        downloadedChunks.push({
-            index: header.chunkIndex,
-            total: header.totalChunks,
-            data: payload,
-            compressed: header.compressed
-        });
-    }
+    const totalBytes = fileRecord.stored_size || chunks.reduce((acc, c) => acc + (c.size || 0), 0);
+    let downloadedBytes = 0;
 
-    // Reassemble
-    onProgress?.('Reassembling...');
-    downloadedChunks.sort((a, b) => a.index - b.index);
-    const encryptedData = Buffer.concat(downloadedChunks.map(c => c.data));
+    // Pre-sort chunks by index so we download them in correct order
+    chunks.sort((a, b) => a.chunk_index - b.chunk_index);
+
+    let currentChunkIndex = 0;
+
+    // We already downloaded the first chunk to inspect its header, we shouldn't discard it.
+    let preloadedFirstChunk = firstChunkData;
+
+    const downloadStream = new Readable({
+        async read() {
+            try {
+                if (currentChunkIndex >= chunks.length) {
+                    this.push(null); // End of stream
+                    return;
+                }
+
+                const chunk = chunks[currentChunkIndex];
+                onProgress?.(`Downloading chunk ${chunk.chunk_index + 1}/${chunks.length}...`);
+
+                let data;
+                if (currentChunkIndex === 0 && preloadedFirstChunk) {
+                    data = preloadedFirstChunk;
+                    preloadedFirstChunk = null;
+                } else {
+                    data = await client.downloadFile(chunk.file_telegram_id);
+                }
+
+                downloadedBytes += data.length;
+                onByteProgress?.({ downloaded: downloadedBytes, total: totalBytes, chunk: chunk.chunk_index + 1, totalChunks: chunks.length });
+
+                // Strip header before pushing
+                const payload = data.subarray(HEADER_SIZE);
+                this.push(payload);
+
+                currentChunkIndex++;
+            } catch (err) {
+                this.destroy(err);
+            }
+        }
+    });
 
-    // Decrypt
-    onProgress?.('Decrypting...');
-    const encryptor = new Encryptor(password);
-    const compressedData = encryptor.decrypt(encryptedData);
+    const writeStream = fs.createWriteStream(outputPath);
+    const { pipeline } = await import('stream/promises');
 
-    // Decompress
-    onProgress?.('Decompressing...');
-    const compressor = new Compressor();
-    const wasCompressed = downloadedChunks[0].compressed;
-    const originalData = await compressor.decompress(compressedData, wasCompressed);
+    onProgress?.('Decrypting, decompressing, and writing file...');
 
-    // Write to output
-    onProgress?.('Writing file...');
-    fs.writeFileSync(outputPath, originalData);
+    // Pipeline: Download from Telegram -> Decrypt -> Decompress -> Disk
+    await pipeline(downloadStream, decryptStream, decompressStream, writeStream);
+
+    const finalStats = fs.statSync(outputPath);
+
+    // Verify file integrity by comparing hash
+    onProgress?.('Verifying file integrity...');
+    const downloadedHash = await hashFile(outputPath);
+    if (fileRecord.hash && downloadedHash !== fileRecord.hash) {
+        throw new Error(`Integrity check failed: expected hash ${fileRecord.hash.substring(0, 12)}..., got ${downloadedHash.substring(0, 12)}...`);
+    }
 
     return {
         path: outputPath,
-        size: originalData.length
+        size: finalStats.size,
+        verified: true
     };
 }

package/src/share/server.js

@@ -57,6 +57,25 @@ function formatTimeLeft(expiresAt) {
     return `${minutes}m`;
 }
 
+/**
+ * Escape HTML special characters to prevent XSS
+ */
+function escapeHtml(str) {
+    return String(str)
+        .replace(/&/g, '&')
+        .replace(/</g, '&lt;')
+        .replace(/>/g, '&gt;')
+        .replace(/"/g, '&quot;')
+        .replace(/'/g, '&#x27;');
+}
+
+/**
+ * Sanitize filename for Content-Disposition header (RFC 6266)
+ */
+function sanitizeFilenameForHeader(filename) {
+    return filename.replace(/["\\\r\n]/g, '_');
+}
+
 /**
  * Generate the download HTML page
  */
@@ -69,6 +88,7 @@ function generateDownloadPage(share, fileRecord) {
         : fileSize > 1024
             ? `${(fileSize / 1024).toFixed(1)} KB`
             : `${fileSize} B`;
+    const safeFilename = escapeHtml(fileRecord.filename);
 
     return `<!DOCTYPE html>
 <html lang="en">
@@ -146,7 +166,7 @@ function generateDownloadPage(share, fileRecord) {
     <div class="card">
         <div class="icon">🔐</div>
         <h1>Secure File Share</h1>
-        <div class="filename">${fileRecord.filename}</div>
+        <div class="filename">${safeFilename}</div>
         <div class="meta">
             <span>📦 ${sizeStr}</span>
             <span>⏳ ${timeLeft}</span>
@@ -241,33 +261,64 @@ export class ShareServer {
     }
 
     /**
-     * Download and decrypt a file from Telegram
+     * Prepare download streams from Telegram — fetches and verifies
+     * the first chunk before returning, so the caller can decide
+     * whether to commit to a 200 response.
+     *
+     * Returns: { downloadStream, decryptStream, decompressStream }
      */
-    async downloadAndDecrypt(fileRecord) {
+    async prepareDownloadStreams(fileRecord) {
         const chunks = this.db.getChunks(fileRecord.id);
         if (chunks.length === 0) throw new Error('No chunks found');
 
-        const downloadedChunks = [];
-
-        for (const chunk of chunks) {
-            const data = await this.client.downloadFile(chunk.file_telegram_id);
-            const header = parseHeader(data);
-            const payload = data.subarray(HEADER_SIZE);
-
-            downloadedChunks.push({
-                index: header.chunkIndex,
-                data: payload,
-                compressed: header.compressed
-            });
-        }
-
-        downloadedChunks.sort((a, b) => a.index - b.index);
-        const encryptedData = Buffer.concat(downloadedChunks.map(c => c.data));
-
-        const compressedData = this.encryptor.decrypt(encryptedData);
+        // Pre-sort chunks by index so we download them in correct order
+        chunks.sort((a, b) => a.chunk_index - b.chunk_index);
+
+        // Fetch and verify the first chunk eagerly
+        const firstChunkData = await this.client.downloadFile(chunks[0].file_telegram_id);
+        const header = parseHeader(firstChunkData);
+        let wasCompressed = header.compressed;
+
+        // Prepare streams
+        const decryptStream = this.encryptor.getDecryptStream();
+        const decompressStream = this.compressor.getDecompressStream(wasCompressed);
+
+        const { Readable } = await import('stream');
+
+        const self = this;
+        let currentChunkIndex = 0;
+        let preloadedFirstChunk = firstChunkData;
+
+        const downloadStream = new Readable({
+            async read() {
+                try {
+                    if (currentChunkIndex >= chunks.length) {
+                        this.push(null); // End of stream
+                        return;
+                    }
+
+                    const chunk = chunks[currentChunkIndex];
+                    let data;
+
+                    if (currentChunkIndex === 0 && preloadedFirstChunk) {
+                        data = preloadedFirstChunk;
+                        preloadedFirstChunk = null;
+                    } else {
+                        data = await self.client.downloadFile(chunk.file_telegram_id);
+                    }
+
+                    // Strip header before pushing
+                    const payload = data.subarray(HEADER_SIZE);
+                    this.push(payload);
+
+                    currentChunkIndex++;
+                } catch (err) {
+                    this.destroy(err);
+                }
+            }
+        });
 
-        const wasCompressed = downloadedChunks[0].compressed;
-        return await this.compressor.decompress(compressedData, wasCompressed);
+        return { downloadStream, decryptStream, decompressStream };
     }
 
     /**
@@ -330,9 +381,6 @@ export class ShareServer {
                 return;
             }
 
-            // Download the file from Telegram, decrypt, and serve
-            const data = await this.downloadAndDecrypt(fileRecord);
-
             // Increment download count
             this.db.incrementShareDownload(token);
 
@@ -351,13 +399,33 @@ export class ShareServer {
                 '.mp3': 'audio/mpeg'
             };
             const contentType = contentTypes[ext] || 'application/octet-stream';
+            const safeName = sanitizeFilenameForHeader(fileRecord.filename);
+
+            try {
+                // Verify first chunk can be fetched and decrypted BEFORE sending 200 OK.
+                // This prevents sending a truncated/corrupt file on Telegram or decrypt failure.
+                const { downloadStream, decryptStream, decompressStream } = await this.prepareDownloadStreams(fileRecord);
+
+                // First chunk verified — commit to the 200 response
+                res.writeHead(200, {
+                    'Content-Type': contentType,
+                    'Content-Disposition': `attachment; filename="${safeName}"; filename*=UTF-8''${encodeURIComponent(fileRecord.filename)}`,
+                    'Transfer-Encoding': 'chunked'
+                });
 
-            res.writeHead(200, {
-                'Content-Type': contentType,
-                'Content-Disposition': `attachment; filename="${fileRecord.filename}"`,
-                'Content-Length': data.length
-            });
-            res.end(data);
+                const { pipeline } = await import('stream/promises');
+                await pipeline(downloadStream, decryptStream, decompressStream, res);
+            } catch (streamErr) {
+                console.error('Share stream error:', streamErr.message);
+                // Headers not yet sent — we can still return a proper error
+                if (!res.headersSent) {
+                    res.writeHead(500, { 'Content-Type': 'text/plain' });
+                    res.end('Download failed — file could not be decrypted or fetched.');
+                } else {
+                    // Headers already sent (shouldn't happen now, but safety net)
+                    res.end();
+                }
+            }
 
         } catch (err) {
             console.error('Share server error:', err.message);