stream.js 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  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. * @private
  26. */
  27. function duplexOnError(err) {
  28. this.removeListener('error', duplexOnError);
  29. this.destroy();
  30. if (this.listenerCount('error') === 0) {
  31. // Do not suppress the throwing behavior.
  32. this.emit('error', err);
  33. }
  34. }
  35. /**
  36. * Wraps a `WebSocket` in a duplex stream.
  37. *
  38. * @param {WebSocket} ws The `WebSocket` to wrap
  39. * @param {Object} options The options for the `Duplex` constructor
  40. * @return {stream.Duplex} The duplex stream
  41. * @public
  42. */
  43. function createWebSocketStream(ws, options) {
  44. let resumeOnReceiverDrain = true;
  45. function receiverOnDrain() {
  46. if (resumeOnReceiverDrain) ws._socket.resume();
  47. }
  48. if (ws.readyState === ws.CONNECTING) {
  49. ws.once('open', function open() {
  50. ws._receiver.removeAllListeners('drain');
  51. ws._receiver.on('drain', receiverOnDrain);
  52. });
  53. } else {
  54. ws._receiver.removeAllListeners('drain');
  55. ws._receiver.on('drain', receiverOnDrain);
  56. }
  57. const duplex = new Duplex({
  58. ...options,
  59. autoDestroy: false,
  60. emitClose: false,
  61. objectMode: false,
  62. readableObjectMode: 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. duplex.destroy(err);
  73. });
  74. ws.once('close', function close() {
  75. if (duplex.destroyed) return;
  76. duplex.push(null);
  77. });
  78. duplex._destroy = function(err, callback) {
  79. if (ws.readyState === ws.CLOSED) {
  80. callback(err);
  81. process.nextTick(emitClose, duplex);
  82. return;
  83. }
  84. ws.once('close', function close() {
  85. callback(err);
  86. process.nextTick(emitClose, duplex);
  87. });
  88. ws.terminate();
  89. };
  90. duplex._final = function(callback) {
  91. if (ws.readyState === ws.CONNECTING) {
  92. ws.once('open', function open() {
  93. duplex._final(callback);
  94. });
  95. return;
  96. }
  97. if (ws._socket._writableState.finished) {
  98. if (duplex._readableState.endEmitted) duplex.destroy();
  99. callback();
  100. } else {
  101. ws._socket.once('finish', function finish() {
  102. // `duplex` is not destroyed here because the `'end'` event will be
  103. // emitted on `duplex` after this `'finish'` event. The EOF signaling
  104. // `null` chunk is, in fact, pushed when the WebSocket emits `'close'`.
  105. callback();
  106. });
  107. ws.close();
  108. }
  109. };
  110. duplex._read = function() {
  111. if (ws.readyState === ws.OPEN && !resumeOnReceiverDrain) {
  112. resumeOnReceiverDrain = true;
  113. if (!ws._receiver._writableState.needDrain) ws._socket.resume();
  114. }
  115. };
  116. duplex._write = function(chunk, encoding, callback) {
  117. if (ws.readyState === ws.CONNECTING) {
  118. ws.once('open', function open() {
  119. duplex._write(chunk, encoding, callback);
  120. });
  121. return;
  122. }
  123. ws.send(chunk, callback);
  124. };
  125. duplex.on('end', duplexOnEnd);
  126. duplex.on('error', duplexOnError);
  127. return duplex;
  128. }
  129. module.exports = createWebSocketStream;