doc/ws.md

@@ -84,6 +84,8 @@ This class represents a WebSocket server. It extends the `EventEmitter`.
   - `verifyClient` {Function} A function which can be used to validate incoming
     connections. See description below. (Usage is discouraged: see
     [Issue #337](https://github.com/websockets/ws/issues/377#issuecomment-462152231))
+  - `WebSocket` {Function} Specifies the `WebSocket` class to be used. It must
+    be extended from the original `WebSocket`. Defaults to `WebSocket`.
 - `callback` {Function}
 
 Create a new server instance. One and only one of `port`, `server` or `noServer`

lib/websocket-server.js

@@ -49,6 +49,8 @@ class WebSocketServer extends EventEmitter {
    * @param {Boolean} [options.skipUTF8Validation=false] Specifies whether or
    *     not to skip UTF-8 validation for text and close messages
    * @param {Function} [options.verifyClient] A hook to reject connections
+   * @param {Function} [options.WebSocket=WebSocket] Specifies the `WebSocket`
+   *     class to use. It must be the `WebSocket` class or class that extends it
    * @param {Function} [callback] A listener for the `listening` event
    */
   constructor(options, callback) {
@@ -67,6 +69,7 @@ class WebSocketServer extends EventEmitter {
       host: null,
       path: null,
       port: null,
+      WebSocket,
       ...options
     };
 
@@ -356,7 +359,7 @@ class WebSocketServer extends EventEmitter {
       `Sec-WebSocket-Accept: ${digest}`
     ];
 
-    const ws = new WebSocket(null);
+    const ws = new this.options.WebSocket(null);
 
     if (protocols.size) {
       //

lib/websocket.js

@@ -766,6 +766,45 @@ function initAsClient(websocket, address, protocols, options) {
     opts.path = parts[1];
   }
 
+  if (opts.followRedirects) {
+    if (websocket._redirects === 0) {
+      websocket._originalHost = parsedUrl.host;
+
+      const headers = options && options.headers;
+
+      //
+      // Shallow copy the user provided options so that headers can be changed
+      // without mutating the original object.
+      //
+      options = { ...options, headers: {} };
+
+      if (headers) {
+        for (const [key, value] of Object.entries(headers)) {
+          options.headers[key.toLowerCase()] = value;
+        }
+      }
+    } else if (parsedUrl.host !== websocket._originalHost) {
+      //
+      // Match curl 7.77.0 behavior and drop the following headers. These
+      // headers are also dropped when following a redirect to a subdomain.
+      //
+      delete opts.headers.authorization;
+      delete opts.headers.cookie;
+      delete opts.headers.host;
+      opts.auth = undefined;
+    }
+
+    //
+    // Match curl 7.77.0 behavior and make the first `Authorization` header win.
+    // If the `Authorization` header is set, then there is nothing to do as it
+    // will take precedence.
+    //
+    if (opts.auth && !options.headers.authorization) {
+      options.headers.authorization =
+        'Basic ' + Buffer.from(opts.auth).toString('base64');
+    }
+  }
+
   let req = (websocket._req = get(opts));
 
   if (opts.timeout) {

package.json

@@ -1,6 +1,6 @@
 {
   "name": "ws",
-  "version": "8.4.2",
+  "version": "8.5.0",
   "description": "Simple to use, blazing fast and thoroughly tested websocket client and server for Node.js",
   "keywords": [
     "HyBi",

test/websocket-server.test.js

@@ -89,6 +89,32 @@ describe('WebSocketServer', () => {
           wss.close(done);
         });
       });
+
+      it('honors the `WebSocket` option', (done) => {
+        class CustomWebSocket extends WebSocket.WebSocket {
+          get foo() {
+            return 'foo';
+          }
+        }
+
+        const wss = new WebSocket.Server(
+          {
+            port: 0,
+            WebSocket: CustomWebSocket
+          },
+          () => {
+            const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
+
+            ws.on('open', ws.close);
+          }
+        );
+
+        wss.on('connection', (ws) => {
+          assert.ok(ws instanceof CustomWebSocket);
+          assert.strictEqual(ws.foo, 'foo');
+          wss.close(done);
+        });
+      });
     });
 
     it('emits an error if http server bind fails', (done) => {

test/websocket.test.js

@@ -1101,11 +1101,9 @@ describe('WebSocket', () => {
     });
 
     it('emits an error if the redirect URL is invalid (1/2)', (done) => {
-      const onUpgrade = (req, socket) => {
+      server.once('upgrade', (req, socket) => {
         socket.end('HTTP/1.1 302 Found\r\nLocation: ws://\r\n\r\n');
-      };
-
-      server.on('upgrade', onUpgrade);
+      });
 
       const ws = new WebSocket(`ws://localhost:${server.address().port}`, {
         followRedirects: true
@@ -1117,17 +1115,14 @@ describe('WebSocket', () => {
         assert.strictEqual(err.message, 'Invalid URL: ws://');
         assert.strictEqual(ws._redirects, 1);
 
-        server.removeListener('upgrade', onUpgrade);
         ws.on('close', () => done());
       });
     });
 
     it('emits an error if the redirect URL is invalid (2/2)', (done) => {
-      const onUpgrade = (req, socket) => {
+      server.once('upgrade', (req, socket) => {
         socket.end('HTTP/1.1 302 Found\r\nLocation: http://localhost\r\n\r\n');
-      };
-
-      server.on('upgrade', onUpgrade);
+      });
 
       const ws = new WebSocket(`ws://localhost:${server.address().port}`, {
         followRedirects: true
@@ -1142,10 +1137,122 @@ describe('WebSocket', () => {
         );
         assert.strictEqual(ws._redirects, 1);
 
-        server.removeListener('upgrade', onUpgrade);
         ws.on('close', () => done());
       });
     });
+
+    it('uses the first url userinfo when following redirects', (done) => {
+      const wss = new WebSocket.Server({ noServer: true, path: '/foo' });
+      const authorization = 'Basic Zm9vOmJhcg==';
+
+      server.once('upgrade', (req, socket) => {
+        socket.end('HTTP/1.1 302 Found\r\nLocation: /foo\r\n\r\n');
+        server.once('upgrade', (req, socket, head) => {
+          wss.handleUpgrade(req, socket, head, (ws, req) => {
+            assert.strictEqual(req.headers.authorization, authorization);
+            ws.close();
+          });
+        });
+      });
+
+      const port = server.address().port;
+      const ws = new WebSocket(`ws://foo:bar@localhost:${port}`, {
+        followRedirects: true
+      });
+
+      assert.strictEqual(ws._req.getHeader('Authorization'), authorization);
+
+      ws.on('close', (code) => {
+        assert.strictEqual(code, 1005);
+        assert.strictEqual(ws.url, `ws://foo:bar@localhost:${port}/foo`);
+        assert.strictEqual(ws._redirects, 1);
+
+        wss.close(done);
+      });
+    });
+
+    describe('When the redirect host is different', () => {
+      it('drops the `auth` option', (done) => {
+        const wss = new WebSocket.Server({ port: 0 }, () => {
+          const port = wss.address().port;
+
+          server.once('upgrade', (req, socket) => {
+            socket.end(
+              `HTTP/1.1 302 Found\r\nLocation: ws://localhost:${port}/\r\n\r\n`
+            );
+          });
+
+          const ws = new WebSocket(`ws://localhost:${server.address().port}`, {
+            auth: 'foo:bar',
+            followRedirects: true
+          });
+
+          assert.strictEqual(
+            ws._req.getHeader('Authorization'),
+            'Basic Zm9vOmJhcg=='
+          );
+
+          ws.on('close', (code) => {
+            assert.strictEqual(code, 1005);
+            assert.strictEqual(ws.url, `ws://localhost:${port}/`);
+            assert.strictEqual(ws._redirects, 1);
+
+            wss.close(done);
+          });
+        });
+
+        wss.on('connection', (ws, req) => {
+          assert.strictEqual(req.headers.authorization, undefined);
+          ws.close();
+        });
+      });
+
+      it('drops the Authorization, Cookie, and Host headers', (done) => {
+        const wss = new WebSocket.Server({ port: 0 }, () => {
+          const port = wss.address().port;
+
+          server.once('upgrade', (req, socket) => {
+            socket.end(
+              `HTTP/1.1 302 Found\r\nLocation: ws://localhost:${port}/\r\n\r\n`
+            );
+          });
+
+          const ws = new WebSocket(`ws://localhost:${server.address().port}`, {
+            headers: {
+              Authorization: 'Basic Zm9vOmJhcg==',
+              Cookie: 'foo=bar',
+              Host: 'foo'
+            },
+            followRedirects: true
+          });
+
+          assert.strictEqual(
+            ws._req.getHeader('Authorization'),
+            'Basic Zm9vOmJhcg=='
+          );
+          assert.strictEqual(ws._req.getHeader('Cookie'), 'foo=bar');
+          assert.strictEqual(ws._req.getHeader('Host'), 'foo');
+
+          ws.on('close', (code) => {
+            assert.strictEqual(code, 1005);
+            assert.strictEqual(ws.url, `ws://localhost:${port}/`);
+            assert.strictEqual(ws._redirects, 1);
+
+            wss.close(done);
+          });
+        });
+
+        wss.on('connection', (ws, req) => {
+          assert.strictEqual(req.headers.authorization, undefined);
+          assert.strictEqual(req.headers.cookie, undefined);
+          assert.strictEqual(
+            req.headers.host,
+            `localhost:${wss.address().port}`
+          );
+          ws.close();
+        });
+      });
+    });
   });
 
   describe('Connection with query string', () => {
@@ -2757,18 +2864,17 @@ describe('WebSocket', () => {
         requestCert: true
       });
 
-      let success = false;
-      const wss = new WebSocket.Server({
-        verifyClient: (info) => {
-          success = !!info.req.client.authorized;
-          return true;
-        },
-        server
-      });
+      const wss = new WebSocket.Server({ noServer: true });
 
-      wss.on('connection', () => {
-        assert.ok(success);
-        server.close(done);
+      server.on('upgrade', (request, socket, head) => {
+        assert.ok(socket.authorized);
+
+        wss.handleUpgrade(request, socket, head, (ws) => {
+          ws.on('close', (code) => {
+            assert.strictEqual(code, 1005);
+            server.close(done);
+          });
+        });
       });
 
       server.listen(0, () => {