diff --git a/src/index.ts b/src/index.ts index 279cc2d..c76a989 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,6 +1,7 @@ import { serve } from 'bun'; import { config } from './config/index'; import { fileInfoCache } from './infrastructure/cache/index'; +import { clearQueue, getQueueStats, waitForQueue } from './infrastructure/telegram/upload-queue'; import { startBot } from './interfaces/bot/handler'; import { handleS3Request } from './interfaces/http/controllers/s3-controller'; import { cleanupRateLimitCache } from './interfaces/http/middleware/rate-limit'; @@ -63,9 +64,22 @@ logger.info('Server started', { port: config.port, url: config.baseUrl }); const gracefulShutdown = async (signal: string): Promise => { logger.info('Graceful shutdown signal received', { signal }); - logger.info('Closing HTTP server'); + logger.info('Closing HTTP server — no new requests accepted'); server.stop(); + // Drain pending upload queue with a timeout + const { pending, size } = getQueueStats(); + if (pending > 0 || size > 0) { + logger.info('Draining upload queue', { pending, size }); + const drainTimeout = setTimeout(() => { + logger.warn('Upload queue drain timeout — clearing remaining tasks'); + clearQueue(); + }, 30_000); + await waitForQueue(); + clearTimeout(drainTimeout); + logger.info('Upload queue drained'); + } + logger.info('Stopping Telegram bot'); bot.stop(signal); diff --git a/src/infrastructure/telegram/bot-pool.ts b/src/infrastructure/telegram/bot-pool.ts index aac67aa..dd20098 100644 --- a/src/infrastructure/telegram/bot-pool.ts +++ b/src/infrastructure/telegram/bot-pool.ts @@ -42,7 +42,9 @@ const sleep = (ms: number): Promise => { const isTransientError = (error: unknown): boolean => { const str = error instanceof Error ? error.message : String(error); const transientPatterns = [ - 'retry after', + // 'retry after' is deliberately omitted — 429 is handled by + // executeWithBotRetry at a deeper layer. Including it here would + // cause double-retry (up to 96 attempts per chunk). 'timeout', 'Timed out', 'etimedout', @@ -76,6 +78,12 @@ const isTransientError = (error: unknown): boolean => { */ const MAX_TRANSIENT_RETRIES = 3; +/** + * Timeout in milliseconds for individual Telegram API calls. + * 120 seconds to accommodate large document uploads. + */ +const TELEGRAM_API_TIMEOUT_MS = 120_000; + /** * Manages a pool of Telegram bots with automatic rotation and rate-limit handling. * @@ -129,7 +137,17 @@ export class BotPool implements ITelegramService { const currentBot = this.bots[botIndex]; const currentToken = this.botTokens[botIndex]; try { - return await action(currentBot, currentToken); + // Add timeout to prevent hung API calls from occupying queue slots + const result = await Promise.race([ + action(currentBot, currentToken), + new Promise((_, reject) => + setTimeout( + () => reject(new Error(`Telegram API timeout after ${TELEGRAM_API_TIMEOUT_MS}ms`)), + TELEGRAM_API_TIMEOUT_MS, + ), + ), + ]); + return result; } catch (error: unknown) { const errorStr = error instanceof Error ? error.message : String(error); const match = errorStr.match(/retry after (\d+)/i); diff --git a/src/infrastructure/telegram/upload-queue.ts b/src/infrastructure/telegram/upload-queue.ts index 913f53b..6642538 100644 --- a/src/infrastructure/telegram/upload-queue.ts +++ b/src/infrastructure/telegram/upload-queue.ts @@ -2,6 +2,13 @@ import PQueue from 'p-queue'; import { config } from '../../env'; import logger from '../../shared/logger/index'; +/** + * Maximum number of pending (queued + in-flight) upload tasks before + * new submissions are rejected. Prevents unbounded memory growth when + * Telegram is slow or unavailable. + */ +const MAX_QUEUE_PENDING = 1000; + /** * P-queue instance for serialising Telegram upload tasks. * @@ -16,7 +23,11 @@ const uploadQueue = new PQueue({ uploadQueue.on('add', () => { const stats = getQueueStats(); if (stats.size > 5) { - logger.warn('Upload queue building up', { pending: stats.pending, size: stats.size }); + logger.warn('Upload queue building up', { + pending: stats.pending, + size: stats.size, + max: MAX_QUEUE_PENDING, + }); } }); @@ -34,6 +45,14 @@ uploadQueue.on('next', () => { * @returns A promise that resolves with the task's result. */ export const enqueueUpload = (task: () => Promise): Promise => { + const stats = getQueueStats(); + if (stats.pending + stats.size > MAX_QUEUE_PENDING) { + return Promise.reject( + new Error( + `Upload queue full (${stats.pending + stats.size} pending, max ${MAX_QUEUE_PENDING})`, + ), + ); + } return uploadQueue.add(task); }; diff --git a/src/interfaces/http/controllers/s3-controller.ts b/src/interfaces/http/controllers/s3-controller.ts index b82134a..7a99a83 100644 --- a/src/interfaces/http/controllers/s3-controller.ts +++ b/src/interfaces/http/controllers/s3-controller.ts @@ -778,25 +778,34 @@ const handleHeadObject = async ( }; /** - * Streams the request body to a temporary file while computing its SHA-256 hash. + * Streams the request body to a temporary file while computing its SHA-256 + * and MD5 hashes. * * Unlike `req.arrayBuffer()`, this approach uses O(1) memory regardless of * file size, making it safe for multi-GB Docker registry layer blobs. * + * MD5 is computed alongside SHA-256 so that Content-MD5 verification (when + * the header is present) does not need to re-read the entire file. + * * @param body - The ReadableStream from the HTTP request body. - * @returns The temp file path, SHA-256 hash, total size, and signature bytes. + * @returns The temp file path, SHA-256 hash, MD5 hash (base64), total size, and signature bytes. */ const streamBodyToTemp = async ( body: ReadableStream | null, ): Promise<{ tempPath: string; fileHash: string; + md5Hash: string; sizeBytes: number; signatureBuffer: Buffer; }> => { const tempPath = `/tmp/filedrop-s3-${nanoid()}`; const writer = Bun.file(tempPath).writer(); - const hasher = new Bun.CryptoHasher('sha256'); + const sha256 = new Bun.CryptoHasher('sha256'); + const md5 = new Bun.CryptoHasher('md5'); + let writerFailed = false; + + // Handle body being null (GET/HEAD/DELETE or empty PUT) const reader = ( body ?? new ReadableStream({ @@ -816,7 +825,8 @@ const streamBodyToTemp = async ( if (done) break; const chunk = Buffer.from(value); sizeBytes += chunk.byteLength; - hasher.update(chunk); + sha256.update(chunk); + md5.update(chunk); writer.write(chunk); if (signatureBytes < SIGNATURE_BYTES) { @@ -827,16 +837,27 @@ const streamBodyToTemp = async ( } } - writer.end(); + try { + writer.end(); + } catch { + writerFailed = true; + } return { tempPath, - fileHash: hasher.digest('hex'), + fileHash: sha256.digest('hex'), + md5Hash: md5.digest('base64'), sizeBytes, signatureBuffer: Buffer.concat(signatureChunks, signatureBytes), }; } catch (error) { - writer.end(); + if (!writerFailed) { + try { + writer.end(); + } catch { + /* writer may already be errored */ + } + } await cleanupTempFile(tempPath); throw error; } finally { @@ -926,25 +947,17 @@ const handlePutObject = async ( } } - // Content-MD5 validation: verify MD5 when Content-MD5 header is present + // Content-MD5 validation: use pre-computed MD5 from streaming (no OOM re-read) const contentMd5 = headers['content-md5']; - if (contentMd5) { - const computedMd5 = Buffer.from( - await crypto.subtle.digest( - 'MD5', - new Uint8Array(await Bun.file(streamed.tempPath).arrayBuffer()), - ), - ).toString('base64'); - if (contentMd5 !== computedMd5) { - await cleanupTempFile(streamed.tempPath); - return s3ErrorResponse( - 'BadDigest', - 'The Content-MD5 you specified did not match what we received.', - `/${bucket}/${key}`, - 400, - reqId, - ); - } + if (contentMd5 && contentMd5 !== streamed.md5Hash) { + await cleanupTempFile(streamed.tempPath); + return s3ErrorResponse( + 'BadDigest', + 'The Content-MD5 you specified did not match what we received.', + `/${bucket}/${key}`, + 400, + reqId, + ); } // M12: Reject oversized bodies @@ -960,17 +973,17 @@ const handlePutObject = async ( } // Idempotent PUT: if the object already exists, skip upload - const existing = await findFileByBucketAndKey(bucketRecord.id, key); - if (existing) { - await cleanupTempFile(streamed.tempPath); - return s3Response(null, 200, reqId, { etag: `"${streamed.fileHash}"` }); - } - try { + const existing = await findFileByBucketAndKey(bucketRecord.id, key); + if (existing) { + await cleanupTempFile(streamed.tempPath); + return s3Response(null, 200, reqId, { etag: `"${streamed.fileHash}"` }); + } + return await storeFileFromTemp(streamed, key, bucketRecord, contentType, reqId); - } catch (uploadError) { + } catch (error) { await cleanupTempFile(streamed.tempPath); - throw uploadError; + throw error; } }; @@ -1022,11 +1035,15 @@ const storeFileFromTemp = async ( return s3Response(null, 200, reqId, { etag: `"${file.fileHash}"` }); } - const forwardResult = await botPool.forwardToStorage( - createReadStream(streamed.tempPath), - partFileNamePrefix, - 'document', - ); + const fileStream = createReadStream(streamed.tempPath); + let forwardResult: ForwardResult; + try { + forwardResult = await botPool.forwardToStorage(fileStream, partFileNamePrefix, 'document'); + } catch (error) { + fileStream.destroy(); + throw error; + } + fileStream.destroy(); const publicId = nanoid(); const { db, files: fileSchema } = await import('../../../db/index');