You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

164 lines
3.8 KiB

4 years ago
  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. writableObjectMode: false
  63. });
  64. ws.on('message', function message(msg) {
  65. if (!duplex.push(msg)) {
  66. resumeOnReceiverDrain = false;
  67. ws._socket.pause();
  68. }
  69. });
  70. ws.once('error', function error(err) {
  71. if (duplex.destroyed) return;
  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. let called = false;
  85. ws.once('error', function error(err) {
  86. called = true;
  87. callback(err);
  88. });
  89. ws.once('close', function close() {
  90. if (!called) callback(err);
  91. process.nextTick(emitClose, duplex);
  92. });
  93. ws.terminate();
  94. };
  95. duplex._final = function (callback) {
  96. if (ws.readyState === ws.CONNECTING) {
  97. ws.once('open', function open() {
  98. duplex._final(callback);
  99. });
  100. return;
  101. }
  102. // If the value of the `_socket` property is `null` it means that `ws` is a
  103. // client websocket and the handshake failed. In fact, when this happens, a
  104. // socket is never assigned to the websocket. Wait for the `'error'` event
  105. // that will be emitted by the websocket.
  106. if (ws._socket === null) return;
  107. if (ws._socket._writableState.finished) {
  108. callback();
  109. if (duplex._readableState.endEmitted) duplex.destroy();
  110. } else {
  111. ws._socket.once('finish', function finish() {
  112. // `duplex` is not destroyed here because the `'end'` event will be
  113. // emitted on `duplex` after this `'finish'` event. The EOF signaling
  114. // `null` chunk is, in fact, pushed when the websocket emits `'close'`.
  115. callback();
  116. });
  117. ws.close();
  118. }
  119. };
  120. duplex._read = function () {
  121. if (ws.readyState === ws.OPEN && !resumeOnReceiverDrain) {
  122. resumeOnReceiverDrain = true;
  123. if (!ws._receiver._writableState.needDrain) ws._socket.resume();
  124. }
  125. };
  126. duplex._write = function (chunk, encoding, callback) {
  127. if (ws.readyState === ws.CONNECTING) {
  128. ws.once('open', function open() {
  129. duplex._write(chunk, encoding, callback);
  130. });
  131. return;
  132. }
  133. ws.send(chunk, callback);
  134. };
  135. duplex.on('end', duplexOnEnd);
  136. duplex.on('error', duplexOnError);
  137. return duplex;
  138. }
  139. module.exports = createWebSocketStream;