diff --git a/src/commands/filesystem.test.ts b/src/commands/filesystem.test.ts new file mode 100644 index 0000000..6cd6607 --- /dev/null +++ b/src/commands/filesystem.test.ts @@ -0,0 +1,147 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { Context } from '../context' +import { Jishaku } from '../jishaku' +import { filesystemCommands } from './filesystem' + +const catCommand = filesystemCommands[0] +const curlCommand = filesystemCommands[1] + +function makeJsk(): Jishaku { + // biome-ignore lint/suspicious/noExplicitAny: minimal fake client for tests. + return new Jishaku({ token: 't0ken-fake' } as any, { consoleLog: false }) +} + +function makeContext(command: string, args: string) { + const send = vi.fn(async (payload: unknown) => ({ payload })) + // biome-ignore lint/suspicious/noExplicitAny: minimal fake message for tests. + const message = { channel: { send }, author: {} } as any + const source = { kind: 'message' as const, message } + const ctx = new Context(makeJsk(), source, command, args) + return { ctx, send } +} + +describe('jsk cat — line span validation', () => { + let dir: string + let filePath: string + + beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'djsk-cat-')) + filePath = join(dir, 'file.txt') + await writeFile(filePath, 'one\ntwo\nthree\n', 'utf-8') + }) + + afterEach(async () => { + await rm(dir, { recursive: true, force: true }) + }) + + it('reads a valid line span', async () => { + const { ctx, send } = makeContext('cat', `${filePath}#L2-3`) + + await catCommand.handler(ctx) + + const [payload] = send.mock.calls[0] as [{ content: string }] + expect(payload.content).toContain('two') + expect(payload.content).toContain('three') + expect(payload.content).not.toContain('one') + }) + + it('rejects #L0 instead of silently returning the last line', async () => { + const { ctx, send } = makeContext('cat', `${filePath}#L0`) + + await catCommand.handler(ctx) + + expect(send).toHaveBeenCalledTimes(1) + const [payload] = send.mock.calls[0] as [string] + expect(payload).toContain('Line numbers must start at 1') + }) + + it('rejects a range end before its start', async () => { + const { ctx, send } = makeContext('cat', `${filePath}#L5-2`) + + await catCommand.handler(ctx) + + const [payload] = send.mock.calls[0] as [string] + expect(payload).toContain('Line numbers must start at 1') + }) +}) + +describe('jsk curl — timeout and size guarding', () => { + const originalFetch = global.fetch + + afterEach(() => { + global.fetch = originalFetch + vi.unstubAllGlobals() + }) + + it('reports a timeout instead of hanging when the request is aborted', async () => { + global.fetch = vi.fn( + () => + new Promise((_resolve, reject) => { + const error = new Error('The operation was aborted due to timeout') + error.name = 'TimeoutError' + reject(error) + }), + ) as unknown as typeof fetch + + const { ctx, send } = makeContext('curl', 'https://example.com/slow') + + await curlCommand.handler(ctx) + + const [payload] = send.mock.calls[0] as [string] + expect(payload).toContain('timed out') + }) + + it('refuses a response whose declared Content-Length exceeds the cap without downloading it', async () => { + const fetchMock = vi.fn( + async () => + new Response('', { + status: 200, + headers: { 'content-length': String(50 * 1024 * 1024) }, + }), + ) + global.fetch = fetchMock as unknown as typeof fetch + + const { ctx, send } = makeContext('curl', 'https://example.com/huge') + + await curlCommand.handler(ctx) + + const [payload] = send.mock.calls[0] as [string] + expect(payload).toContain('Refusing to download') + }) + + it('cuts off a response that exceeds the cap even without an accurate Content-Length', async () => { + // No content-length header at all — the streamed-byte cap must still catch this. + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(9 * 1024 * 1024).fill(97)) // 9 MiB, over the 8 MiB cap + controller.close() + }, + }) + global.fetch = vi.fn( + async () => new Response(stream, { status: 200 }), + ) as unknown as typeof fetch + + const { ctx, send } = makeContext('curl', 'https://example.com/unbounded') + + await curlCommand.handler(ctx) + + const [payload] = send.mock.calls[0] as [string] + expect(payload).toContain('Refusing to download') + }) + + it('downloads and displays a normal response', async () => { + global.fetch = vi.fn( + async () => new Response('hello from the internet', { status: 200 }), + ) as unknown as typeof fetch + + const { ctx, send } = makeContext('curl', 'https://example.com/ok') + + await curlCommand.handler(ctx) + + const [payload] = send.mock.calls[0] as [{ content: string }] + expect(payload.content).toContain('hello from the internet') + }) +}) diff --git a/src/commands/filesystem.ts b/src/commands/filesystem.ts index 98f99d8..d8dea41 100644 --- a/src/commands/filesystem.ts +++ b/src/commands/filesystem.ts @@ -1,7 +1,7 @@ import { readFile, stat } from 'node:fs/promises' import { basename } from 'node:path' import type { Context } from '../context' -import { toFile } from '../util/format' +import { naturalSize, toFile } from '../util/format' import type { Command } from './registry' const MAX_FILE_SIZE = 8 * 1024 * 1024 // 8 MiB @@ -9,6 +9,44 @@ const CODEBLOCK_BYTE_LIMIT = 20_000 const CAT_ARG = /^(?:\.\/+)?(.+?)(?:#L?(\d+)(?:-L?(\d+))?)?$/ +// Mirrors MAX_FILE_SIZE: `curl` has no equivalent of cat's on-disk size check (there's no +// `stat()` to consult before downloading), so the cap is enforced by capping the fetch itself. +const MAX_RESPONSE_SIZE = 8 * 1024 * 1024 // 8 MiB +const FETCH_TIMEOUT_MS = 15_000 + +/** Thrown by {@link readTextLimited} once the response body exceeds `maxBytes`. */ +class ResponseTooLargeError extends Error {} + +/** + * Reads `response`'s body as UTF-8 text, aborting once it exceeds `maxBytes`. + * + * Unlike `response.text()`, this doesn't trust `Content-Length` (absent or inaccurate for + * chunked/compressed responses) — it enforces the cap against the bytes actually received, + * streaming and cancelling as soon as the limit is crossed instead of buffering an + * unboundedly large body into memory first. + */ +async function readTextLimited(response: Response, maxBytes: number): Promise { + if (!response.body) return response.text() + + const reader = response.body.getReader() + const chunks: Uint8Array[] = [] + let total = 0 + + while (true) { + const { done, value } = await reader.read() + if (done) break + + total += value.byteLength + if (total > maxBytes) { + await reader.cancel().catch(() => {}) + throw new ResponseTooLargeError() + } + chunks.push(value) + } + + return Buffer.concat(chunks).toString('utf-8') +} + /** Sends text as a plain codeblock, falling back to a file attachment when it is large. */ async function sendText(ctx: Context, content: string, filename: string): Promise { const scrubbed = ctx.jsk.scrub(content) @@ -40,6 +78,11 @@ const catCommand: Command = { ? ([Number.parseInt(match[2], 10), Number.parseInt(match[3] ?? match[2], 10)] as const) : null + if (lineSpan && (lineSpan[0] < 1 || lineSpan[1] < lineSpan[0])) { + await ctx.send('Line numbers must start at 1, and the range end must not precede its start.') + return + } + let info: Awaited> try { info = await stat(path) @@ -89,13 +132,36 @@ const curlCommand: Command = { let response: Response try { - response = await fetch(url) + response = await fetch(url, { signal: AbortSignal.timeout(FETCH_TIMEOUT_MS) }) + } catch (error) { + const message = + error instanceof Error && error.name === 'TimeoutError' + ? `Request timed out after ${FETCH_TIMEOUT_MS}ms.` + : `Request failed: ${error instanceof Error ? error.message : String(error)}` + await ctx.send(message) + return + } + + const declaredSize = Number(response.headers.get('content-length')) + if (Number.isFinite(declaredSize) && declaredSize > MAX_RESPONSE_SIZE) { + await ctx.send( + `Refusing to download a response larger than ${naturalSize(MAX_RESPONSE_SIZE)} ` + + `(reported ${naturalSize(declaredSize)}).`, + ) + return + } + + let data: string + try { + data = await readTextLimited(response, MAX_RESPONSE_SIZE) } catch (error) { - await ctx.send(`Request failed: ${error instanceof Error ? error.message : String(error)}`) + if (!(error instanceof ResponseTooLargeError)) throw error + await ctx.send( + `Refusing to download a response larger than ${naturalSize(MAX_RESPONSE_SIZE)}.`, + ) return } - const data = await response.text() if (!data) { await ctx.send(`HTTP response was empty (status code ${response.status}).`) return diff --git a/src/commands/js.test.ts b/src/commands/js.test.ts index 65cc4db..556856b 100644 --- a/src/commands/js.test.ts +++ b/src/commands/js.test.ts @@ -118,6 +118,23 @@ describe('jsk js — terminal output capture', () => { expect(payload.content).not.toContain('[33m') }) + it('keeps codeblock fences on every page when terminal output needs pagination', async () => { + const { ctx, send } = makeContext("process.stdout.write('x'.repeat(3000))") + + await jsCommand.handler(ctx) + + // Terminal output alone already exceeds the message limit, so it must be sent through the + // codeblock-aware pagination (one message per fenced page) rather than combined with the + // result into one oversized string and split at fixed offsets — which would only leave the + // fences intact at the very start of page 1 and the very end of the last page. + expect(send).toHaveBeenCalled() + const [payload] = send.mock.calls[0] as [{ content: string }] + expect(payload.content.startsWith('```\n')).toBe(true) + const closeIndex = payload.content.indexOf('```', 4) + expect(closeIndex).toBeGreaterThan(-1) + expect(payload.content.slice(closeIndex)).toMatch(/^```\n-- Page 1\/\d+ --$/) + }) + it('does not scrub non-token secrets when security mode is off', async () => { process.env.JS_TEST_SECRET_KEY_2 = 'another-secret-value-654321' try { diff --git a/src/commands/js.ts b/src/commands/js.ts index 554173f..e463d2b 100644 --- a/src/commands/js.ts +++ b/src/commands/js.ts @@ -2,7 +2,7 @@ import vm from 'node:vm' import type { Context } from '../context' import { installPrototypeGuards, installRestGuard } from '../prototype-guard' import { guardOutbound } from '../security' -import { inspectResult, stripAnsi } from '../util/format' +import { inspectResult, MESSAGE_LIMIT, stripAnsi } from '../util/format' import { loadLibraryModule } from '../util/meta' import type { Command } from './registry' @@ -206,6 +206,45 @@ function captureTerminalOutput(scrub: ((text: string) => string) | null): { } } +/** + * Sends `terminalOutput` (wrapped in a codeblock) followed by `text` (plain), preferring a + * single combined message when it fits. + * + * When it doesn't fit, the two are sent as separate, independently-paginated messages instead + * of combining them into one oversized string and handing that to {@link Context.sendResult}: + * that pagination is codeblock-*unaware* (by design — it's also used for plain results with no + * codeblock at all), so it slices the combined text at fixed byte offsets with no regard for + * where the codeblock's fences landed. The fences only happen to survive at the very start of + * page 1 and the very end of the last page; every page in between is missing both, since + * nothing re-opens/re-closes the codeblock at the split points. Sending `terminalOutput` + * through {@link Context.sendCodeblock} instead re-wraps every page in its own fences. + */ +async function sendTerminalAndText( + ctx: Context, + terminalOutput: string, + text: string | null, + filename: string, +): Promise { + if (!terminalOutput) { + if (text !== null) await ctx.sendResult(text, filename) + return + } + + const codeblock = `\`\`\`\n${terminalOutput}\n\`\`\`` + const combined = text !== null ? `${codeblock}\n${text}` : codeblock + + // Matches Context.sendResult's own scrubbed-length check, so this decides on the same basis + // it will (scrubbing doesn't reliably preserve length, and ctx.sendResult scrubs again below + // regardless — redaction is idempotent, so double-scrubbing is harmless). + if (ctx.jsk.scrub(combined).length <= MESSAGE_LIMIT) { + await ctx.sendResult(combined, filename) + return + } + + await ctx.sendCodeblock(terminalOutput, '', filename) + if (text !== null) await ctx.sendResult(text, filename) +} + async function sendResult(ctx: Context, result: unknown, terminalOutput: string): Promise { const resultText = isMessage(result) ? // biome-ignore lint/suspicious/noExplicitAny: verified Message-like above. @@ -214,14 +253,12 @@ async function sendResult(ctx: Context, result: unknown, terminalOutput: string) ? null : inspectResult(result) - const parts: string[] = [] - if (terminalOutput) parts.push(`\`\`\`\n${terminalOutput}\n\`\`\``) - if (resultText !== null) parts.push(resultText) - if (parts.length === 0) return + if (!terminalOutput && resultText === null) return - // Routed through ctx.sendResult (not a raw send) so the captured terminal output gets the - // same token redaction / security-mode secret scrubbing as everything else djsk sends. - await ctx.sendResult(parts.join('\n'), 'output.js') + // Routed through sendTerminalAndText (which itself routes through ctx.sendResult/ + // sendCodeblock, not a raw send) so the captured terminal output gets the same token + // redaction / security-mode secret scrubbing as everything else djsk sends. + await sendTerminalAndText(ctx, terminalOutput, resultText, 'output.js') } const jsCommand: Command = { @@ -359,10 +396,7 @@ const jsCommand: Command = { if (error instanceof EvalTimedOutError) { await ctx.react('⏱️') - const parts = terminalOutput - ? [`\`\`\`\n${terminalOutput}\n\`\`\``, error.message] - : [error.message] - await ctx.sendResult(parts.join('\n'), 'output.js') + await sendTerminalAndText(ctx, terminalOutput, error.message, 'output.js') return } diff --git a/src/commands/shell.test.ts b/src/commands/shell.test.ts index 6afbe3b..e427133 100644 --- a/src/commands/shell.test.ts +++ b/src/commands/shell.test.ts @@ -73,6 +73,39 @@ describe('jsk sh — final output rendering', () => { expect(sentMessage.react).toHaveBeenCalledWith('➡️') }) + it('keeps codeblock fences on every page, both the initial (tail) page and after paging', async () => { + stubReader(['x'.repeat(3000)]) + const { ctx, sentMessage } = makeContext('echo big') + + await shellCommand.handler(ctx) + + // The initial render (last page, tail-first) is set via a plain-string edit, unlike + // `jsk js`'s old bug — `jsk sh` has always built every page through `wrapPages`' + // prefix/suffix (see format.ts), which re-wraps each page's own fences, rather than + // wrapping the whole output once and paginating that as one codeblock-unaware blob. + const [initialContent] = sentMessage.edit.mock.calls[0] as [string] + expect(initialContent.startsWith('```powershell\n')).toBe(true) + const closeIndex = initialContent.indexOf('```', 4) + expect(closeIndex).toBeGreaterThan(-1) + expect(initialContent.slice(closeIndex)).toMatch(/^```\n-- Page \d+\/\d+ --$/) + + // Page backward via ⬅️ and check the newly-rendered page also has both fences. + const collector = sentMessage.createReactionCollector.mock.results[0].value as EventEmitter + collector.emit( + 'collect', + { emoji: { name: '⬅️' }, users: { remove: vi.fn(async () => {}) } }, + { id: 'owner-1', bot: false }, + ) + await new Promise((resolve) => setTimeout(resolve, 0)) + + const lastEditCall = sentMessage.edit.mock.calls.at(-1) as [{ content: string }] + const pagedContent = lastEditCall[0].content + expect(pagedContent.startsWith('```powershell\n')).toBe(true) + const pagedCloseIndex = pagedContent.indexOf('```', 4) + expect(pagedCloseIndex).toBeGreaterThan(-1) + expect(pagedContent.slice(pagedCloseIndex)).toMatch(/^```\n-- Page \d+\/\d+ --$/) + }) + it('falls back to a file attachment (in addition to the tail) for very large output', async () => { stubReader(['x'.repeat(25000)]) const { ctx, send, sentMessage } = makeContext('echo huge') @@ -118,4 +151,56 @@ describe('jsk sh — final output rendering', () => { const [, options] = vi.mocked(ShellReader).mock.calls.at(-1) as [string, { shell?: unknown }] expect(options.shell).toBeNull() }) + + it('does not send a duplicate message when a periodic flush is still in flight when the command finishes', async () => { + vi.useFakeTimers() + try { + let resolveDone: (code: number) => void = () => {} + const done = new Promise((resolve) => { + resolveDone = resolve + }) + + let resolveFirstSend: (value: unknown) => void = () => {} + const firstSend = new Promise((resolve) => { + resolveFirstSend = resolve + }) + + fakeReaderImpl = (_code, options) => { + options.onLine('some output') + return { ps1: 'PS >', highlight: 'powershell', done, kill: vi.fn() } + } + + const sentMessage = { + react: vi.fn(async () => {}), + edit: vi.fn(async (payload: unknown) => ({ ...sentMessage, payload })), + createReactionCollector: vi.fn(() => new EventEmitter()), + } + // biome-ignore lint/suspicious/noExplicitAny: minimal test double. + const send = vi.fn((): any => (send.mock.calls.length === 1 ? firstSend : sentMessage)) + // biome-ignore lint/suspicious/noExplicitAny: minimal fake message for tests. + const message = { channel: { send }, author: { id: 'owner-1' } } as any + const ctx = new Context(makeJsk(), { kind: 'message', message }, 'sh', 'echo hi') + + const handlerPromise = shellCommand.handler(ctx) + + // Let the periodic flush fire once, kicking off an in-flight ctx.send() that won't + // resolve yet — simulating a slow Discord API response. + await vi.advanceTimersByTimeAsync(1500) + expect(send).toHaveBeenCalledTimes(1) + + // The command "finishes" while that send is still pending. + resolveDone(0) + await vi.advanceTimersByTimeAsync(0) + + // Now let the in-flight send resolve. + resolveFirstSend(sentMessage) + await handlerPromise + + // Only one message for this single, short output — not a second, independent one racing + // the in-flight periodic flush. + expect(send).toHaveBeenCalledTimes(1) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/src/commands/shell.ts b/src/commands/shell.ts index 113f2a1..9d71067 100644 --- a/src/commands/shell.ts +++ b/src/commands/shell.ts @@ -56,8 +56,20 @@ const shellCommand: Command = { } } + // Serializes flush() calls through a single chain (never more than one send/edit request in + // flight at once), so awaiting it always means "whatever was pending has now fully settled + // `message`" — including a call the periodic interval already kicked off. Without this, the + // final render below could still see `message` as `null` while an interval-triggered flush + // is mid-`ctx.send()`, and would send a second, independent message for the same output + // instead of reusing/editing the one already in flight. + let flushChain: Promise = Promise.resolve() + const runFlush = (): Promise => { + flushChain = flushChain.then(flush) + return flushChain + } + const task = ctx.jsk.submitTask('jsk sh', () => reader.kill()) - const interval = setInterval(() => void flush(), EDIT_INTERVAL) + const interval = setInterval(() => void runFlush(), EDIT_INTERVAL) try { const exitCode = await reader.done @@ -68,6 +80,11 @@ const shellCommand: Command = { ctx.jsk.removeTask(task) } + // Drain the flush chain (including anything still in flight from the interval) before + // touching `message` below — see runFlush's doc comment. This also performs the final + // tail-render (with the exit code appended above) for the common case. + await runFlush() + // Final render: like `jsk js`, output that doesn't fit in one message gets ⬅️/➡️ // pagination over the FULL output (not just the live tail) instead of staying // tail-truncated, falling back to a file attachment (alongside the tail) only when it's @@ -78,12 +95,10 @@ const shellCommand: Command = { const pages = wrapPages(output, { prefix, suffix: '```', maxSize: 1940 }) if (pages.length <= 1) { - await flush() return } if (pages.length > MAX_PAGINATED_PAGES) { - await flush() await ctx.send({ files: [toFile('output.txt', output)] }) return }