diff --git a/addons/addon-image/src/kitty/KittyGraphicsHandler.ts b/addons/addon-image/src/kitty/KittyGraphicsHandler.ts index 810f114f..eb202908 100644 --- a/addons/addon-image/src/kitty/KittyGraphicsHandler.ts +++ b/addons/addon-image/src/kitty/KittyGraphicsHandler.ts @@ -65,6 +65,10 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler { private _decodedChunks: Uint8Array[] = []; private _totalDecodedSize = 0; + /** Pre-calculated encoded size limit */ + private _encodedSizeLimit = 0; + private _totalEncodedSize = 0; + /** Parsed command. These are the control data before semicolon. */ private _parsedCommand: IKittyCommand | null = null; @@ -100,6 +104,9 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler { this._decodedChunks = []; this._totalDecodedSize = 0; this._parsedCommand = null; + // Pre-calculate encoded limit once: base64 is 4 bytes encoded → 3 bytes decoded + this._encodedSizeLimit = Math.ceil(this._opts.kittySizeLimit * 4 / 3); + this._totalEncodedSize = 0; } public put(data: Uint32Array, start: number, end: number): void { @@ -158,9 +165,14 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler { private _streamPayload(data: Uint32Array, start: number, end: number): void { if (this._aborted) return; - // Check size limit - const estimatedTotal = this._totalDecodedSize + Math.ceil(this._shardBufferPos * 3 / 4) + Math.ceil((end - start) * 3 / 4); - if (estimatedTotal > this._opts.kittySizeLimit) { + // Check size limit (compare encoded bytes against pre-calculated limit) + // Include cumulative size from pending transmission for multi-chunk images + const pendingKey = this._parsedCommand?.id ?? 0; + const pending = this._pendingTransmissions.get(pendingKey); + const previousEncodedSize = pending?.totalEncodedSize ?? 0; + this._totalEncodedSize += end - start; + const cumulativeEncodedSize = previousEncodedSize + this._totalEncodedSize; + if (cumulativeEncodedSize > this._encodedSizeLimit) { this._aborted = true; return; } @@ -290,6 +302,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler { cmd: { ...cmd }, chunks: [], totalSize: 0, + totalEncodedSize: 0, leftover: new Uint32Array(4), leftoverLength: 0 }; @@ -466,6 +479,8 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler { pending.chunks.push(bytes); pending.totalSize += bytes.length; } + // Track cumulative encoded size for size limit enforcement + pending.totalEncodedSize += this._totalEncodedSize; if (isMoreComing) return true; @@ -498,6 +513,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler { cmd: { ...cmd }, chunks: bytes.length > 0 ? [bytes] : [], totalSize: bytes.length, + totalEncodedSize: this._totalEncodedSize, leftover: new Uint32Array(4), leftoverLength: 0 }); diff --git a/addons/addon-image/src/kitty/KittyGraphicsTypes.ts b/addons/addon-image/src/kitty/KittyGraphicsTypes.ts index ef91bf57..d58a51a7 100644 --- a/addons/addon-image/src/kitty/KittyGraphicsTypes.ts +++ b/addons/addon-image/src/kitty/KittyGraphicsTypes.ts @@ -103,8 +103,10 @@ export interface IPendingTransmission { cmd: IKittyCommand; /** Accumulated decoded payload chunks */ chunks: Uint8Array[]; - /** Total size of accumulated chunks */ + /** Total size of accumulated decoded chunks */ totalSize: number; + /** Total encoded (base64) bytes received across all chunks - for size limit enforcement */ + totalEncodedSize: number; /** Leftover base64 bytes (0-3) that weren't aligned in previous chunk */ leftover: Uint32Array; /** Number of valid bytes in leftover */ diff --git a/addons/addon-image/test/ImageAddon.test.ts b/addons/addon-image/test/ImageAddon.test.ts index d07990dc..55282d46 100644 --- a/addons/addon-image/test/ImageAddon.test.ts +++ b/addons/addon-image/test/ImageAddon.test.ts @@ -391,6 +391,39 @@ test.describe('ImageAddon', () => { deepStrictEqual(storedData, KITTY_BLACK_1X1_BYTES); }); + test('enforces size limit across chunked transmissions', async () => { + // Create a custom addon with very small size limit (100 bytes) + // The 1x1 PNG is ~164 bytes base64, so 2 chunks should exceed 100 + await ctx.page.evaluate(() => { + (window as any).smallLimitAddon = new ImageAddon({ + kittySupport: true, + kittySizeLimit: 100 // Very small limit + }); + (window as any).term.loadAddon((window as any).smallLimitAddon); + }); + + // Split the base64 data into two chunks + const half = Math.floor(KITTY_BLACK_1X1_BASE64.length / 2); + const part1 = KITTY_BLACK_1X1_BASE64.substring(0, half); + const part2 = KITTY_BLACK_1X1_BASE64.substring(half); + + // Send chunked data - first chunk (~82 bytes) is under limit + await ctx.proxy.write(`\x1b_Ga=t,f=100,i=777,m=1;${part1}\x1b\\`); + await timeout(50); + + // Second chunk brings total to ~164 bytes, exceeding 100 byte limit + await ctx.proxy.write(`\x1b_Ga=t,f=100,i=777;${part2}\x1b\\`); + await timeout(100); + + // Image should NOT be stored due to size limit + strictEqual(await ctx.page.evaluate(`window.smallLimitAddon._handlers.get('kitty').images.has(777)`), false); + + // Cleanup + await ctx.page.evaluate(() => { + (window as any).smallLimitAddon.dispose(); + }); + }); + test('delete command (a=d) removes specific image by id', async () => { await ctx.proxy.write(`\x1b_Ga=t,f=100,i=10;${KITTY_BLACK_1X1_BASE64}\x1b\\`); await timeout(50);