api-stream.js 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  1. 'use strict'
  2. const assert = require('node:assert')
  3. const { finished } = require('node:stream')
  4. const { AsyncResource } = require('node:async_hooks')
  5. const { InvalidArgumentError, InvalidReturnValueError } = require('../core/errors')
  6. const util = require('../core/util')
  7. const { addSignal, removeSignal } = require('./abort-signal')
  8. function noop () {}
  9. class StreamHandler extends AsyncResource {
  10. constructor (opts, factory, callback) {
  11. if (!opts || typeof opts !== 'object') {
  12. throw new InvalidArgumentError('invalid opts')
  13. }
  14. const { signal, method, opaque, body, onInfo, responseHeaders } = opts
  15. try {
  16. if (typeof callback !== 'function') {
  17. throw new InvalidArgumentError('invalid callback')
  18. }
  19. if (typeof factory !== 'function') {
  20. throw new InvalidArgumentError('invalid factory')
  21. }
  22. if (signal && typeof signal.on !== 'function' && typeof signal.addEventListener !== 'function') {
  23. throw new InvalidArgumentError('signal must be an EventEmitter or EventTarget')
  24. }
  25. if (method === 'CONNECT') {
  26. throw new InvalidArgumentError('invalid method')
  27. }
  28. if (onInfo && typeof onInfo !== 'function') {
  29. throw new InvalidArgumentError('invalid onInfo callback')
  30. }
  31. super('UNDICI_STREAM')
  32. } catch (err) {
  33. if (util.isStream(body)) {
  34. util.destroy(body.on('error', noop), err)
  35. }
  36. throw err
  37. }
  38. this.responseHeaders = responseHeaders || null
  39. this.opaque = opaque || null
  40. this.factory = factory
  41. this.callback = callback
  42. this.res = null
  43. this.abort = null
  44. this.context = null
  45. this.trailers = null
  46. this.body = body
  47. this.onInfo = onInfo || null
  48. if (util.isStream(body)) {
  49. body.on('error', (err) => {
  50. this.onError(err)
  51. })
  52. }
  53. addSignal(this, signal)
  54. }
  55. onConnect (abort, context) {
  56. if (this.reason) {
  57. abort(this.reason)
  58. return
  59. }
  60. assert(this.callback)
  61. this.abort = abort
  62. this.context = context
  63. }
  64. onHeaders (statusCode, rawHeaders, resume, statusMessage) {
  65. const { factory, opaque, context, responseHeaders } = this
  66. const headers = responseHeaders === 'raw' ? util.parseRawHeaders(rawHeaders) : util.parseHeaders(rawHeaders)
  67. if (statusCode < 200) {
  68. if (this.onInfo) {
  69. this.onInfo({ statusCode, headers })
  70. }
  71. return
  72. }
  73. this.factory = null
  74. if (factory === null) {
  75. return
  76. }
  77. const res = this.runInAsyncScope(factory, null, {
  78. statusCode,
  79. headers,
  80. opaque,
  81. context
  82. })
  83. if (
  84. !res ||
  85. typeof res.write !== 'function' ||
  86. typeof res.end !== 'function' ||
  87. typeof res.on !== 'function'
  88. ) {
  89. throw new InvalidReturnValueError('expected Writable')
  90. }
  91. // TODO: Avoid finished. It registers an unnecessary amount of listeners.
  92. finished(res, { readable: false }, (err) => {
  93. const { callback, res, opaque, trailers, abort } = this
  94. this.res = null
  95. if (err || !res?.readable) {
  96. util.destroy(res, err)
  97. }
  98. this.callback = null
  99. this.runInAsyncScope(callback, null, err || null, { opaque, trailers })
  100. if (err) {
  101. abort()
  102. }
  103. })
  104. res.on('drain', resume)
  105. this.res = res
  106. const needDrain = res.writableNeedDrain !== undefined
  107. ? res.writableNeedDrain
  108. : res._writableState?.needDrain
  109. return needDrain !== true
  110. }
  111. onData (chunk) {
  112. const { res } = this
  113. return res ? res.write(chunk) : true
  114. }
  115. onComplete (trailers) {
  116. const { res } = this
  117. removeSignal(this)
  118. if (!res) {
  119. return
  120. }
  121. this.trailers = util.parseHeaders(trailers)
  122. res.end()
  123. }
  124. onError (err) {
  125. const { res, callback, opaque, body } = this
  126. removeSignal(this)
  127. this.factory = null
  128. if (res) {
  129. this.res = null
  130. util.destroy(res, err)
  131. } else if (callback) {
  132. this.callback = null
  133. queueMicrotask(() => {
  134. this.runInAsyncScope(callback, null, err, { opaque })
  135. })
  136. }
  137. if (body) {
  138. this.body = null
  139. util.destroy(body, err)
  140. }
  141. }
  142. }
  143. function stream (opts, factory, callback) {
  144. if (callback === undefined) {
  145. return new Promise((resolve, reject) => {
  146. stream.call(this, opts, factory, (err, data) => {
  147. return err ? reject(err) : resolve(data)
  148. })
  149. })
  150. }
  151. try {
  152. const handler = new StreamHandler(opts, factory, callback)
  153. this.dispatch(opts, handler)
  154. } catch (err) {
  155. if (typeof callback !== 'function') {
  156. throw err
  157. }
  158. const opaque = opts?.opaque
  159. queueMicrotask(() => callback(err, { opaque }))
  160. }
  161. }
  162. module.exports = stream