v7
  1// Copyright 2021 Google LLC
  2//
  3// Licensed under the Apache License, Version 2.0 (the "License");
  4// you may not use this file except in compliance with the License.
  5// You may obtain a copy of the License at
  6//
  7//     https://www.apache.org/licenses/LICENSE-2.0
  8//
  9// Unless required by applicable law or agreed to in writing, software
 10// distributed under the License is distributed on an "AS IS" BASIS,
 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12// See the License for the specific language governing permissions and
 13// limitations under the License.
 14
 15import assert from 'node:assert'
 16import { ChildProcess, spawn, StdioNull, StdioPipe } from 'node:child_process'
 17import { AsyncLocalStorage, createHook } from 'node:async_hooks'
 18import { Readable, Writable } from 'node:stream'
 19import { inspect } from 'node:util'
 20import { RequestInfo, RequestInit } from 'node-fetch'
 21import chalk, { ChalkInstance } from 'chalk'
 22import which from 'which'
 23import {
 24  Duration,
 25  errnoMessage,
 26  exitCodeInfo,
 27  formatCmd,
 28  noop,
 29  parseDuration,
 30  psTree,
 31  quote,
 32  quotePowerShell,
 33} from './util.js'
 34
 35export type Shell = (
 36  pieces: TemplateStringsArray,
 37  ...args: any[]
 38) => ProcessPromise
 39
 40const processCwd = Symbol('processCwd')
 41
 42export type Options = {
 43  [processCwd]: string
 44  cwd?: string
 45  verbose: boolean
 46  env: NodeJS.ProcessEnv
 47  shell: string | boolean
 48  prefix: string
 49  quote: typeof quote
 50  spawn: typeof spawn
 51  log: typeof log
 52}
 53
 54const storage = new AsyncLocalStorage<Options>()
 55const hook = createHook({
 56  init: syncCwd,
 57  before: syncCwd,
 58  promiseResolve: syncCwd,
 59  after: syncCwd,
 60  destroy: syncCwd,
 61})
 62hook.enable()
 63
 64export const defaults: Options = {
 65  [processCwd]: process.cwd(),
 66  verbose: true,
 67  env: process.env,
 68  shell: true,
 69  prefix: '',
 70  quote: () => {
 71    throw new Error('No quote function is defined: https://ï.at/no-quote-func')
 72  },
 73  spawn,
 74  log,
 75}
 76
 77try {
 78  defaults.shell = which.sync('bash')
 79  defaults.prefix = 'set -euo pipefail;'
 80  defaults.quote = quote
 81} catch (err) {
 82  if (process.platform == 'win32') {
 83    defaults.shell = which.sync('powershell.exe')
 84    defaults.quote = quotePowerShell
 85  }
 86}
 87
 88function getStore() {
 89  return storage.getStore() || defaults
 90}
 91
 92export const $ = new Proxy<Shell & Options>(
 93  function (pieces, ...args) {
 94    const from = new Error().stack!.split(/^\s*at\s/m)[2].trim()
 95    if (pieces.some((p) => p == undefined)) {
 96      throw new Error(`Malformed command at ${from}`)
 97    }
 98    let resolve: Resolve, reject: Resolve
 99    const promise = new ProcessPromise((...args) => ([resolve, reject] = args))
100    let cmd = pieces[0],
101      i = 0
102    while (i < args.length) {
103      let s
104      if (Array.isArray(args[i])) {
105        s = args[i].map((x: any) => $.quote(substitute(x))).join(' ')
106      } else {
107        s = $.quote(substitute(args[i]))
108      }
109      cmd += s + pieces[++i]
110    }
111    promise._bind(cmd, from, resolve!, reject!, getStore())
112    // Postpone run to allow promise configuration.
113    setImmediate(() => promise.isHalted || promise.run())
114    return promise
115  } as Shell & Options,
116  {
117    set(_, key, value) {
118      const target = key in Function.prototype ? _ : getStore()
119      Reflect.set(target, key, value)
120      return true
121    },
122    get(_, key) {
123      const target = key in Function.prototype ? _ : getStore()
124      return Reflect.get(target, key)
125    },
126  }
127)
128
129function substitute(arg: ProcessPromise | any) {
130  if (arg?.stdout) {
131    return arg.stdout.replace(/\n$/, '')
132  }
133  return `${arg}`
134}
135
136type Resolve = (out: ProcessOutput) => void
137type IO = StdioPipe | StdioNull
138
139export class ProcessPromise extends Promise<ProcessOutput> {
140  child?: ChildProcess
141  private _command = ''
142  private _from = ''
143  private _resolve: Resolve = noop
144  private _reject: Resolve = noop
145  private _snapshot = getStore()
146  private _stdio: [IO, IO, IO] = ['inherit', 'pipe', 'pipe']
147  private _nothrow = false
148  private _quiet = false
149  private _timeout?: number
150  private _timeoutSignal?: string
151  private _resolved = false
152  private _halted = false
153  private _piped = false
154  _prerun = noop
155  _postrun = noop
156
157  _bind(
158    cmd: string,
159    from: string,
160    resolve: Resolve,
161    reject: Resolve,
162    options: Options
163  ) {
164    this._command = cmd
165    this._from = from
166    this._resolve = resolve
167    this._reject = reject
168    this._snapshot = { ...options }
169  }
170
171  run(): ProcessPromise {
172    const $ = this._snapshot
173    if (this.child) return this // The _run() can be called from a few places.
174    this._prerun() // In case $1.pipe($2), the $2 returned, and on $2._run() invoke $1._run().
175    $.log({
176      kind: 'cmd',
177      cmd: this._command,
178      verbose: $.verbose && !this._quiet,
179    })
180    this.child = $.spawn($.prefix + this._command, {
181      cwd: $.cwd ?? $[processCwd],
182      shell: typeof $.shell === 'string' ? $.shell : true,
183      stdio: this._stdio,
184      windowsHide: true,
185      env: $.env,
186    })
187    this.child.on('close', (code, signal) => {
188      let message = `exit code: ${code}`
189      if (code != 0 || signal != null) {
190        message = `${stderr || '\n'}    at ${this._from}`
191        message += `\n    exit code: ${code}${
192          exitCodeInfo(code) ? ' (' + exitCodeInfo(code) + ')' : ''
193        }`
194        if (signal != null) {
195          message += `\n    signal: ${signal}`
196        }
197      }
198      let output = new ProcessOutput(
199        code,
200        signal,
201        stdout,
202        stderr,
203        combined,
204        message
205      )
206      if (code === 0 || this._nothrow) {
207        this._resolve(output)
208      } else {
209        this._reject(output)
210      }
211      this._resolved = true
212    })
213    this.child.on('error', (err: NodeJS.ErrnoException) => {
214      const message =
215        `${err.message}\n` +
216        `    errno: ${err.errno} (${errnoMessage(err.errno)})\n` +
217        `    code: ${err.code}\n` +
218        `    at ${this._from}`
219      this._reject(
220        new ProcessOutput(null, null, stdout, stderr, combined, message)
221      )
222      this._resolved = true
223    })
224    let stdout = '',
225      stderr = '',
226      combined = ''
227    let onStdout = (data: any) => {
228      $.log({ kind: 'stdout', data, verbose: $.verbose && !this._quiet })
229      stdout += data
230      combined += data
231    }
232    let onStderr = (data: any) => {
233      $.log({ kind: 'stderr', data, verbose: $.verbose && !this._quiet })
234      stderr += data
235      combined += data
236    }
237    if (!this._piped) this.child.stdout?.on('data', onStdout) // If process is piped, don't collect or print output.
238    this.child.stderr?.on('data', onStderr) // Stderr should be printed regardless of piping.
239    this._postrun() // In case $1.pipe($2), after both subprocesses are running, we can pipe $1.stdout to $2.stdin.
240    if (this._timeout && this._timeoutSignal) {
241      const t = setTimeout(() => this.kill(this._timeoutSignal), this._timeout)
242      this.finally(() => clearTimeout(t)).catch(noop)
243    }
244    return this
245  }
246
247  get stdin(): Writable {
248    this.stdio('pipe')
249    this.run()
250    assert(this.child)
251    if (this.child.stdin == null)
252      throw new Error('The stdin of subprocess is null.')
253    return this.child.stdin
254  }
255
256  get stdout(): Readable {
257    this.run()
258    assert(this.child)
259    if (this.child.stdout == null)
260      throw new Error('The stdout of subprocess is null.')
261    return this.child.stdout
262  }
263
264  get stderr(): Readable {
265    this.run()
266    assert(this.child)
267    if (this.child.stderr == null)
268      throw new Error('The stderr of subprocess is null.')
269    return this.child.stderr
270  }
271
272  get exitCode(): Promise<number | null> {
273    return this.then(
274      (p) => p.exitCode,
275      (p) => p.exitCode
276    )
277  }
278
279  then<R = ProcessOutput, E = ProcessOutput>(
280    onfulfilled?:
281      | ((value: ProcessOutput) => PromiseLike<R> | R)
282      | undefined
283      | null,
284    onrejected?:
285      | ((reason: ProcessOutput) => PromiseLike<E> | E)
286      | undefined
287      | null
288  ): Promise<R | E> {
289    if (this.isHalted && !this.child) {
290      throw new Error('The process is halted!')
291    }
292    return super.then(onfulfilled, onrejected)
293  }
294
295  catch<T = ProcessOutput>(
296    onrejected?:
297      | ((reason: ProcessOutput) => PromiseLike<T> | T)
298      | undefined
299      | null
300  ): Promise<ProcessOutput | T> {
301    return super.catch(onrejected)
302  }
303
304  pipe(dest: Writable | ProcessPromise): ProcessPromise {
305    if (typeof dest == 'string')
306      throw new Error('The pipe() method does not take strings. Forgot $?')
307    if (this._resolved) {
308      if (dest instanceof ProcessPromise) dest.stdin.end() // In case of piped stdin, we may want to close stdin of dest as well.
309      throw new Error(
310        "The pipe() method shouldn't be called after promise is already resolved!"
311      )
312    }
313    this._piped = true
314    if (dest instanceof ProcessPromise) {
315      dest.stdio('pipe')
316      dest._prerun = this.run.bind(this)
317      dest._postrun = () => {
318        if (!dest.child)
319          throw new Error(
320            'Access to stdin of pipe destination without creation a subprocess.'
321          )
322        this.stdout.pipe(dest.stdin)
323      }
324      return dest
325    } else {
326      this._postrun = () => this.stdout.pipe(dest)
327      return this
328    }
329  }
330
331  async kill(signal = 'SIGTERM'): Promise<void> {
332    if (!this.child)
333      throw new Error('Trying to kill a process without creating one.')
334    if (!this.child.pid) throw new Error('The process pid is undefined.')
335    let children = await psTree(this.child.pid)
336    for (const p of children) {
337      try {
338        process.kill(+p.PID, signal)
339      } catch (e) {}
340    }
341    try {
342      process.kill(this.child.pid, signal)
343    } catch (e) {}
344  }
345
346  stdio(stdin: IO, stdout: IO = 'pipe', stderr: IO = 'pipe'): ProcessPromise {
347    this._stdio = [stdin, stdout, stderr]
348    return this
349  }
350
351  nothrow(): ProcessPromise {
352    this._nothrow = true
353    return this
354  }
355
356  quiet(): ProcessPromise {
357    this._quiet = true
358    return this
359  }
360
361  timeout(d: Duration, signal = 'SIGTERM'): ProcessPromise {
362    this._timeout = parseDuration(d)
363    this._timeoutSignal = signal
364    return this
365  }
366
367  halt(): ProcessPromise {
368    this._halted = true
369    return this
370  }
371
372  get isHalted(): boolean {
373    return this._halted
374  }
375}
376
377export class ProcessOutput extends Error {
378  private readonly _code: number | null
379  private readonly _signal: NodeJS.Signals | null
380  private readonly _stdout: string
381  private readonly _stderr: string
382  private readonly _combined: string
383
384  constructor(
385    code: number | null,
386    signal: NodeJS.Signals | null,
387    stdout: string,
388    stderr: string,
389    combined: string,
390    message: string
391  ) {
392    super(message)
393    this._code = code
394    this._signal = signal
395    this._stdout = stdout
396    this._stderr = stderr
397    this._combined = combined
398  }
399
400  toString() {
401    return this._combined
402  }
403
404  get stdout() {
405    return this._stdout
406  }
407
408  get stderr() {
409    return this._stderr
410  }
411
412  get exitCode() {
413    return this._code
414  }
415
416  get signal() {
417    return this._signal
418  }
419
420  [inspect.custom]() {
421    let stringify = (s: string, c: ChalkInstance) =>
422      s.length === 0 ? "''" : c(inspect(s))
423    return `ProcessOutput {
424  stdout: ${stringify(this.stdout, chalk.green)},
425  stderr: ${stringify(this.stderr, chalk.red)},
426  signal: ${inspect(this.signal)},
427  exitCode: ${(this.exitCode === 0 ? chalk.green : chalk.red)(this.exitCode)}${
428    exitCodeInfo(this.exitCode)
429      ? chalk.grey(' (' + exitCodeInfo(this.exitCode) + ')')
430      : ''
431  }
432}`
433  }
434}
435
436export function within<R>(callback: () => R): R {
437  return storage.run({ ...getStore() }, callback)
438}
439
440function syncCwd() {
441  if ($[processCwd] != process.cwd()) process.chdir($[processCwd])
442}
443
444export function cd(dir: string | ProcessOutput) {
445  if (dir instanceof ProcessOutput) {
446    dir = dir.toString().replace(/\n+$/, '')
447  }
448
449  $.log({ kind: 'cd', dir })
450  process.chdir(dir)
451  $[processCwd] = process.cwd()
452}
453
454export type LogEntry =
455  | {
456      kind: 'cmd'
457      verbose: boolean
458      cmd: string
459    }
460  | {
461      kind: 'stdout' | 'stderr'
462      verbose: boolean
463      data: Buffer
464    }
465  | {
466      kind: 'cd'
467      dir: string
468    }
469  | {
470      kind: 'fetch'
471      url: RequestInfo
472      init?: RequestInit
473    }
474  | {
475      kind: 'retry'
476      error: string
477    }
478  | {
479      kind: 'custom'
480      data: any
481    }
482
483export function log(entry: LogEntry) {
484  switch (entry.kind) {
485    case 'cmd':
486      if (!entry.verbose) return
487      process.stderr.write(formatCmd(entry.cmd))
488      break
489    case 'stdout':
490    case 'stderr':
491      if (!entry.verbose) return
492      process.stderr.write(entry.data)
493      break
494    case 'cd':
495      if (!$.verbose) return
496      process.stderr.write('$ ' + chalk.greenBright('cd') + ` ${entry.dir}\n`)
497      break
498    case 'fetch':
499      if (!$.verbose) return
500      const init = entry.init ? ' ' + inspect(entry.init) : ''
501      process.stderr.write(
502        '$ ' + chalk.greenBright('fetch') + ` ${entry.url}${init}\n`
503      )
504      break
505    case 'retry':
506      if (!$.verbose) return
507      process.stderr.write(entry.error + '\n')
508  }
509}