diff --git a/index.d.ts b/index.d.ts index 9c18dfb..bc281ff 100644 --- a/index.d.ts +++ b/index.d.ts @@ -17,6 +17,16 @@ export interface Size { cols: number rows: number } +/** + * Set the close-on-exec flag on a file descriptor. This is `fcntl(fd, F_SETFD, FD_CLOEXEC)` under + * the covers. + */ +export function setCloseOnExec(fd: number, closeOnExec: boolean): void +/** + * Get the close-on-exec flag on a file descriptor. This is `fcntl(fd, F_GETFD) & FD_CLOEXEC == + *_CLOEXEC` under the covers. + */ +export function getCloseOnExec(fd: number): boolean export class Pty { /** The pid of the forked process. */ pid: number diff --git a/index.js b/index.js index 1a2a909..251c62d 100644 --- a/index.js +++ b/index.js @@ -310,6 +310,8 @@ if (!nativeBinding) { throw new Error(`Failed to load native binding`) } -const { Pty } = nativeBinding +const { Pty, setCloseOnExec, getCloseOnExec } = nativeBinding module.exports.Pty = Pty +module.exports.setCloseOnExec = setCloseOnExec +module.exports.getCloseOnExec = getCloseOnExec diff --git a/package-lock.json b/package-lock.json index 9c50be2..f95d728 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@replit/ruspty", - "version": "3.0.5", + "version": "3.1.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@replit/ruspty", - "version": "3.0.5", + "version": "3.1.0", "license": "MIT", "devDependencies": { "@napi-rs/cli": "^2.18.2", diff --git a/package.json b/package.json index f9620c7..821c16c 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@replit/ruspty", - "version": "3.0.5", + "version": "3.1.0", "main": "dist/wrapper.js", "types": "dist/wrapper.d.ts", "author": "Szymon Kaliski ", diff --git a/src/lib.rs b/src/lib.rs index 8242be8..f87eb52 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFun use napi::Status::GenericFailure; use napi::{self, Env}; use nix::errno::Errno; +use nix::fcntl::{fcntl, FcntlArg, FdFlag}; use nix::poll::{poll, PollFd, PollFlags, PollTimeout}; use nix::pty::{openpty, Winsize}; use nix::sys::termios::{self, SetArg}; @@ -200,7 +201,7 @@ impl Pty { thread::spawn(move || { let wait_result = child.wait(); - // try to wait for the controller fd to be fully read + // try to wait for the controller fd to be fully read poll_controller_fd_until_read(raw_controller_fd); // we don't drop the controller fd immediately @@ -302,3 +303,55 @@ impl Pty { Ok(()) } } + +/// Set the close-on-exec flag on a file descriptor. This is `fcntl(fd, F_SETFD, FD_CLOEXEC)` under +/// the covers. +#[napi] +#[allow(dead_code)] +fn set_close_on_exec(fd: i32, close_on_exec: bool) -> Result<(), napi::Error> { + let mut flags = match fcntl(fd as RawFd, FcntlArg::F_GETFD) { + Ok(flags) => FdFlag::from_bits_truncate(flags), + Err(err) => { + return Err(napi::Error::new( + GenericFailure, + format!("fcntl F_GETFD: {}", err,), + )); + } + }; + if close_on_exec { + if flags.contains(FdFlag::FD_CLOEXEC) { + // It's already there! + return Ok(()); + } + flags.insert(FdFlag::FD_CLOEXEC); + } else { + if !flags.contains(FdFlag::FD_CLOEXEC) { + // It's already removed! + return Ok(()); + } + flags.remove(FdFlag::FD_CLOEXEC); + } + + if let Err(err) = fcntl(fd as RawFd, FcntlArg::F_SETFD(flags)) { + return Err(napi::Error::new( + GenericFailure, + format!("fcntl F_SETFD: {}", err,), + )); + }; + + Ok(()) +} + +/// Get the close-on-exec flag on a file descriptor. This is `fcntl(fd, F_GETFD) & FD_CLOEXEC == +///_CLOEXEC` under the covers. +#[napi] +#[allow(dead_code)] +fn get_close_on_exec(fd: i32) -> Result { + match fcntl(fd as RawFd, FcntlArg::F_GETFD) { + Ok(flags) => Ok(FdFlag::from_bits_truncate(flags).contains(FdFlag::FD_CLOEXEC)), + Err(err) => Err(napi::Error::new( + GenericFailure, + format!("fcntl F_GETFD: {}", err,), + )), + } +} diff --git a/tests/index.test.ts b/tests/index.test.ts index cdf5d94..6b6b856 100644 --- a/tests/index.test.ts +++ b/tests/index.test.ts @@ -1,4 +1,4 @@ -import { Pty } from '../wrapper'; +import { Pty, getCloseOnExec, setCloseOnExec } from '../wrapper'; import { readdirSync, readlinkSync } from 'fs'; import { describe, test, expect } from 'vitest'; @@ -15,7 +15,10 @@ function getOpenFds(): FdRecord { for (const filename of readdirSync(procSelfFd)) { try { const linkTarget = readlinkSync(procSelfFd + filename); - if (linkTarget === 'anon_inode:[timerfd]' || linkTarget.startsWith("socket:[")) { + if ( + linkTarget === 'anon_inode:[timerfd]' || + linkTarget.startsWith('socket:[') + ) { continue; } @@ -31,212 +34,244 @@ function getOpenFds(): FdRecord { return fds; } -describe('PTY', () => { - - test('spawns and exits', () => new Promise(done => { - const oldFds = getOpenFds(); - const message = 'hello from a pty'; - let buffer = ''; - - const pty = new Pty({ - command: '/bin/echo', - args: [message], - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(0); - expect(buffer.trim()).toBe(message); - expect(getOpenFds()).toStrictEqual(oldFds); - done() - }, - }); - - const readStream = pty.read; - readStream.on('data', (chunk) => { - buffer = chunk.toString(); - }); - })); - - test('captures an exit code', () => new Promise(done => { - const oldFds = getOpenFds(); - new Pty({ - command: '/bin/sh', - args: ['-c', 'exit 17'], - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(17); - expect(getOpenFds()).toStrictEqual(oldFds); - done(); - }, - }); - })); - - test('can be written to', () => new Promise(done => { - const oldFds = getOpenFds(); - - // The message should end in newline so that the EOT can signal that the input has ended and not - // just the line. - const message = 'hello cat\n'; - let buffer = ''; - - // We have local echo enabled, so we'll read the message twice. - const result = process.platform === "darwin" - ? 'hello cat\r\n^D\b\bhello cat\r\n' - : 'hello cat\r\nhello cat\r\n'; - - const pty = new Pty({ - command: '/bin/cat', - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(0); - expect(buffer.trim()).toBe(result.trim()); - expect(getOpenFds()).toStrictEqual(oldFds); - done(); - }, - }); - - const writeStream = pty.write; - const readStream = pty.read; - - readStream.on('data', (data) => { - buffer += data.toString(); - }); - writeStream.write(message); - writeStream.end(EOT); - })); - - test('can be resized', () => new Promise(done => { - const oldFds = getOpenFds(); - let buffer = ''; - let state: 'expectPrompt' | 'expectDone1' | 'expectDone2' | 'done' = 'expectPrompt'; - const pty = new Pty({ - command: '/bin/sh', - size: { rows: 24, cols: 80 }, - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(0); - - expect(state).toBe('done'); - expect(getOpenFds()).toStrictEqual(oldFds); - done() - }, - }); - - const writeStream = pty.write; - const readStream = pty.read; - - readStream.on('data', (data) => { - buffer += data.toString(); - - if (state === 'expectPrompt' && buffer.endsWith('$ ')) { - writeStream.write("stty size; echo 'done1'\n"); - state = 'expectDone1'; - return; - } - - if (state === 'expectDone1' && buffer.includes('done1\r\n')) { - state = 'expectDone2'; - expect(buffer).toContain('24 80'); - pty.resize({ rows: 60, cols: 100 }); - - writeStream.write("stty size; echo 'done2'\n"); - return; - } - - if (state === 'expectDone2' && buffer.includes('done2\r\n')) { - expect(buffer).toContain('60 100'); - state = 'done'; - - writeStream.write(EOT); - return; - } - }); - - })); - - test('respects working directory', () => new Promise(done => { - const oldFds = getOpenFds(); - const cwd = process.cwd(); - let buffer = ''; - - const pty = new Pty({ - command: '/bin/pwd', - dir: cwd, - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(0); - expect(buffer.trim()).toBe(cwd); - expect(getOpenFds()).toStrictEqual(oldFds); - done(); - }, - }); - - const readStream = pty.read; - readStream.on('data', (data) => { - buffer += data.toString(); - }); - })); - - test('respects env', () => new Promise(done => { - const oldFds = getOpenFds(); - const message = 'hello from env'; - let buffer = ''; - - const pty = new Pty({ - command: '/bin/sh', - args: ['-c', 'echo $ENV_VARIABLE && exit'], - envs: { - ENV_VARIABLE: message, - }, - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(0); - expect(buffer.trim()).toBe(message); - expect(getOpenFds()).toStrictEqual(oldFds); - done(); - }, - }); - - const readStream = pty.read; - readStream.on('data', (data) => { - buffer += data.toString(); - }); - })); - - test('ordering is correct', () => new Promise(done => { - const oldFds = getOpenFds(); - let buffer = Buffer.from(''); - const n = 1024; - const pty = new Pty({ - command: '/bin/sh', - args: ['-c', `for i in $(seq 0 ${n}); do /bin/echo $i; done && exit`], - onExit: (err, exitCode) => { - expect(err).toBeNull(); - expect(exitCode).toBe(0); - expect(buffer.toString().trim()).toBe([...Array(n + 1).keys()].join('\r\n')); - expect(getOpenFds()).toStrictEqual(oldFds); - done(); - }, - }); - - const readStream = pty.read; - readStream.on('data', (data) => { - buffer = Buffer.concat([buffer, data]); - }); - })); - - test("doesn't break when executing non-existing binary", () => new Promise((done) => { - const oldFds = getOpenFds(); - try { - new Pty({ - command: '/bin/this-does-not-exist', - onExit: () => { - expect(getOpenFds()).toStrictEqual(oldFds); - }, - }); - } catch (e: any) { - expect(e.message).toContain('No such file or directory'); - - done(); +describe( + 'PTY', + () => { + test('spawns and exits', () => + new Promise((done) => { + const oldFds = getOpenFds(); + const message = 'hello from a pty'; + let buffer = ''; + + const pty = new Pty({ + command: '/bin/echo', + args: [message], + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(0); + expect(buffer.trim()).toBe(message); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + + const readStream = pty.read; + readStream.on('data', (chunk) => { + buffer = chunk.toString(); + }); + })); + + test('captures an exit code', () => + new Promise((done) => { + const oldFds = getOpenFds(); + new Pty({ + command: '/bin/sh', + args: ['-c', 'exit 17'], + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(17); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + })); + + test('can be written to', () => + new Promise((done) => { + const oldFds = getOpenFds(); + + // The message should end in newline so that the EOT can signal that the input has ended and not + // just the line. + const message = 'hello cat\n'; + let buffer = ''; + + // We have local echo enabled, so we'll read the message twice. + const result = + process.platform === 'darwin' + ? 'hello cat\r\n^D\b\bhello cat\r\n' + : 'hello cat\r\nhello cat\r\n'; + + const pty = new Pty({ + command: '/bin/cat', + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(0); + expect(buffer.trim()).toBe(result.trim()); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + + const writeStream = pty.write; + const readStream = pty.read; + + readStream.on('data', (data) => { + buffer += data.toString(); + }); + writeStream.write(message); + writeStream.end(EOT); + })); + + test('can be resized', () => + new Promise((done) => { + const oldFds = getOpenFds(); + let buffer = ''; + let state: 'expectPrompt' | 'expectDone1' | 'expectDone2' | 'done' = + 'expectPrompt'; + const pty = new Pty({ + command: '/bin/sh', + size: { rows: 24, cols: 80 }, + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(0); + + expect(state).toBe('done'); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + + const writeStream = pty.write; + const readStream = pty.read; + + readStream.on('data', (data) => { + buffer += data.toString(); + + if (state === 'expectPrompt' && buffer.endsWith('$ ')) { + writeStream.write("stty size; echo 'done1'\n"); + state = 'expectDone1'; + return; + } + + if (state === 'expectDone1' && buffer.includes('done1\r\n')) { + state = 'expectDone2'; + expect(buffer).toContain('24 80'); + pty.resize({ rows: 60, cols: 100 }); + + writeStream.write("stty size; echo 'done2'\n"); + return; + } + + if (state === 'expectDone2' && buffer.includes('done2\r\n')) { + expect(buffer).toContain('60 100'); + state = 'done'; + + writeStream.write(EOT); + return; + } + }); + })); + + test('respects working directory', () => + new Promise((done) => { + const oldFds = getOpenFds(); + const cwd = process.cwd(); + let buffer = ''; + + const pty = new Pty({ + command: '/bin/pwd', + dir: cwd, + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(0); + expect(buffer.trim()).toBe(cwd); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + + const readStream = pty.read; + readStream.on('data', (data) => { + buffer += data.toString(); + }); + })); + + test('respects env', () => + new Promise((done) => { + const oldFds = getOpenFds(); + const message = 'hello from env'; + let buffer = ''; + + const pty = new Pty({ + command: '/bin/sh', + args: ['-c', 'echo $ENV_VARIABLE && exit'], + envs: { + ENV_VARIABLE: message, + }, + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(0); + expect(buffer.trim()).toBe(message); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + + const readStream = pty.read; + readStream.on('data', (data) => { + buffer += data.toString(); + }); + })); + + test('ordering is correct', () => + new Promise((done) => { + const oldFds = getOpenFds(); + let buffer = Buffer.from(''); + const n = 1024; + const pty = new Pty({ + command: '/bin/sh', + args: ['-c', `for i in $(seq 0 ${n}); do /bin/echo $i; done && exit`], + onExit: (err, exitCode) => { + expect(err).toBeNull(); + expect(exitCode).toBe(0); + expect(buffer.toString().trim()).toBe( + [...Array(n + 1).keys()].join('\r\n'), + ); + expect(getOpenFds()).toStrictEqual(oldFds); + done(); + }, + }); + + const readStream = pty.read; + readStream.on('data', (data) => { + buffer = Buffer.concat([buffer, data]); + }); + })); + + test("doesn't break when executing non-existing binary", () => + new Promise((done) => { + const oldFds = getOpenFds(); + try { + new Pty({ + command: '/bin/this-does-not-exist', + onExit: () => { + expect(getOpenFds()).toStrictEqual(oldFds); + }, + }); + } catch (e: any) { + expect(e.message).toContain('No such file or directory'); + + done(); + } + })); + }, + { repeats: 50 }, +); + +describe('setCloseOnExec', () => { + test('setCloseOnExec', () => { + // stdio typically never has the close-on-exec flag since it's always expected to be + // inheritable. But just to be safe, we'll keep it as it was when started. + const originalFlag = getCloseOnExec(0); + + setCloseOnExec(0, true); + expect(getCloseOnExec(0)).toBe(true); + + setCloseOnExec(0, false); + expect(getCloseOnExec(0)).toBe(false); + + if (originalFlag) { + setCloseOnExec(0, originalFlag); } - })); -}, { repeats: 50 }); + }); +}); diff --git a/wrapper.ts b/wrapper.ts index db82a9b..44bbecc 100644 --- a/wrapper.ts +++ b/wrapper.ts @@ -1,5 +1,10 @@ import { PassThrough, Readable, Writable } from 'stream'; -import { Pty as RawPty, type Size } from './index.js'; +import { + Pty as RawPty, + type Size, + setCloseOnExec as rawSetCloseOnExec, + getCloseOnExec as rawGetCloseOnExec, +} from './index.js'; import { type PtyOptions as RawOptions } from './index.js'; import fs from 'fs'; @@ -8,7 +13,7 @@ export type PtyOptions = RawOptions; type ExitResult = { error: NodeJS.ErrnoException | null; code: number; -} +}; /** * A very thin wrapper around PTYs and processes. @@ -48,14 +53,14 @@ export class Pty { let resolve: (value: ExitResult) => void; let exitResult: Promise = new Promise((res) => { resolve = res; - }) + }); const mockedExit = (error: NodeJS.ErrnoException | null, code: number) => { resolve({ error, code }); - } + }; // when pty exits, we should wait until the fd actually ends (end OR error) // before closing the pty - // we use a mocked exit function to capture the exit result + // we use a mocked exit function to capture the exit result // and then call the real exit function after the fd is fully read this.#pty = new RawPty({ ...options, onExit: mockedExit }); const fd = this.#pty.fd(); @@ -79,10 +84,10 @@ export class Pty { exitResult.then((result) => realExit(result.error, result.code)); this.#pty.close(); userFacingRead.end(); - } + }; // catch end events - read.on('end', eof) + read.on('end', eof); // strip out EIO errors read.on('error', (err: NodeJS.ErrnoException) => { @@ -116,3 +121,15 @@ export class Pty { return this.#pty.pid; } } + +/** + * Set the close-on-exec flag on a file descriptor. This is `fcntl(fd, F_SETFD, FD_CLOEXEC)` under + * the covers. + */ +export const setCloseOnExec = rawGetCloseOnExec; + +/** + * Get the close-on-exec flag on a file descriptor. This is `fcntl(fd, F_GETFD) & FD_CLOEXEC == + * FD_CLOEXEC` under the covers. + */ +export const getCloseOnExec = rawGetCloseOnExec;