stream.js 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165
  1. 'use strict';
  2. const { Duplex } = require('stream');
  3. /**
  4. * Emits the `'close'` event on a stream.
  5. *
  6. * @param {stream.Duplex} The stream.
  7. * @private
  8. */
  9. function emitClose(stream) {
  10. stream.emit('close');
  11. }
  12. /**
  13. * The listener of the `'end'` event.
  14. *
  15. * @private
  16. */
  17. function duplexOnEnd() {
  18. if (!this.destroyed && this._writableState.finished) {
  19. this.destroy();
  20. }
  21. }
  22. /**
  23. * The listener of the `'error'` event.
  24. *
  25. * @param {Error} err The error
  26. * @private
  27. */
  28. function duplexOnError(err) {
  29. this.removeListener('error', duplexOnError);
  30. this.destroy();
  31. if (this.listenerCount('error') === 0) {
  32. // Do not suppress the throwing behavior.
  33. this.emit('error', err);
  34. }
  35. }
  36. /**
  37. * Wraps a `WebSocket` in a duplex stream.
  38. *
  39. * @param {WebSocket} ws The `WebSocket` to wrap
  40. * @param {Object} [options] The options for the `Duplex` constructor
  41. * @return {stream.Duplex} The duplex stream
  42. * @public
  43. */
  44. function createWebSocketStream(ws, options) {
  45. let resumeOnReceiverDrain = true;
  46. function receiverOnDrain() {
  47. if (resumeOnReceiverDrain) ws._socket.resume();
  48. }
  49. if (ws.readyState === ws.CONNECTING) {
  50. ws.once('open', function open() {
  51. ws._receiver.removeAllListeners('drain');
  52. ws._receiver.on('drain', receiverOnDrain);
  53. });
  54. } else {
  55. ws._receiver.removeAllListeners('drain');
  56. ws._receiver.on('drain', receiverOnDrain);
  57. }
  58. const duplex = new Duplex({
  59. ...options,
  60. autoDestroy: false,
  61. emitClose: false,
  62. objectMode: false,
  63. writableObjectMode: false
  64. });
  65. ws.on('message', function message(msg) {
  66. if (!duplex.push(msg)) {
  67. resumeOnReceiverDrain = false;
  68. ws._socket.pause();
  69. }
  70. });
  71. ws.once('error', function error(err) {
  72. if (duplex.destroyed) return;
  73. duplex.destroy(err);
  74. });
  75. ws.once('close', function close() {
  76. if (duplex.destroyed) return;
  77. duplex.push(null);
  78. });
  79. duplex._destroy = function (err, callback) {
  80. if (ws.readyState === ws.CLOSED) {
  81. callback(err);
  82. process.nextTick(emitClose, duplex);
  83. return;
  84. }
  85. let called = false;
  86. ws.once('error', function error(err) {
  87. called = true;
  88. callback(err);
  89. });
  90. ws.once('close', function close() {
  91. if (!called) callback(err);
  92. process.nextTick(emitClose, duplex);
  93. });
  94. ws.terminate();
  95. };
  96. duplex._final = function (callback) {
  97. if (ws.readyState === ws.CONNECTING) {
  98. ws.once('open', function open() {
  99. duplex._final(callback);
  100. });
  101. return;
  102. }
  103. // If the value of the `_socket` property is `null` it means that `ws` is a
  104. // client websocket and the handshake failed. In fact, when this happens, a
  105. // socket is never assigned to the websocket. Wait for the `'error'` event
  106. // that will be emitted by the websocket.
  107. if (ws._socket === null) return;
  108. if (ws._socket._writableState.finished) {
  109. callback();
  110. if (duplex._readableState.endEmitted) duplex.destroy();
  111. } else {
  112. ws._socket.once('finish', function finish() {
  113. // `duplex` is not destroyed here because the `'end'` event will be
  114. // emitted on `duplex` after this `'finish'` event. The EOF signaling
  115. // `null` chunk is, in fact, pushed when the websocket emits `'close'`.
  116. callback();
  117. });
  118. ws.close();
  119. }
  120. };
  121. duplex._read = function () {
  122. if (ws.readyState === ws.OPEN && !resumeOnReceiverDrain) {
  123. resumeOnReceiverDrain = true;
  124. if (!ws._receiver._writableState.needDrain) ws._socket.resume();
  125. }
  126. };
  127. duplex._write = function (chunk, encoding, callback) {
  128. if (ws.readyState === ws.CONNECTING) {
  129. ws.once('open', function open() {
  130. duplex._write(chunk, encoding, callback);
  131. });
  132. return;
  133. }
  134. ws.send(chunk, callback);
  135. };
  136. duplex.on('end', duplexOnEnd);
  137. duplex.on('error', duplexOnError);
  138. return duplex;
  139. }
  140. module.exports = createWebSocketStream;