handling upgrades

README.md

@@ -118,6 +118,71 @@ const wsServer = new WebSocketServer({
 });
 ```
 
+## Handling requests before connection
+
+Instead of this:
+
+```js
+const http = require("http");
+const express = require("express");
+const ws = require("ws");
+
+const app = express();
+const wsServer = new ws.WebSocketServer({ noServer: true, path: "/wspath" }); // path is optional
+
+app.get("/", (_, res) => res.send("Hello, world!"));
+
+const server = http.createServer(app);
+server.on("upgrade", async (request, socket, head) => {
+    const user = await getUserFromDatabase(request.headers['authorization']); // your auth logic
+    if(!user) return socket.destroy();
+
+    wsServer.handleUpgrade(request, socket, head, (ws) => {
+        ws.user = user;
+        wsServer.emit("connection", ws, request);
+    });
+});
+
+server.listen(3000);
+```
+
+You can do this:
+```js
+const { WebSocketServer } = require("ultimate-ws");
+const express = require("ultimate-express");
+
+const app = express();
+
+const wsServer = new WebSocketServer({
+    server: app,
+    handleUpgrade: async (request) => {
+        const user = await getUserFromDatabase(request.headers['authorization']); // your auth logic
+        if(!user) {
+            // request has `req` and `res` properties, which are instances of `uws.HttpRequest` and `uws.HttpResponse`
+            request.res.cork(() => {
+                request.res.writeStatus("401 Unauthorized");
+                request.res.end();
+            });
+            return false;
+        }
+
+        return (ws, request) => {
+            ws.user = user;
+            wsServer.emit("connection", ws, request);
+        }
+    }
+});
+
+app.get("/", (_, res) => res.send("Hello, world!"));
+
+app.listen(3000);
+```
+
+- if `handleUpgrade` returns a function, it will be called with the new `WebSocket` instance and original request. "connection" will not be emitted automatically.
+- if it returns `false`, the connection will not be upgraded. It's your responsibility to destroy the socket.
+- if it returns nothing or anything else, the connection will be handled as usual, and "connection" event will be emitted.
+- By default (`handleUpgrade: undefined`), the connection will be handled as usual, and "connection" event will be emitted.
+
 ## Compatibility
 
 All commonly used `ws` features are supported. Almost all ws servers should work, as it's built with maximum compatibility in mind.

package.json

@@ -1,6 +1,6 @@
 {
   "name": "ultimate-ws",
-  "version": "1.0.5",
+  "version": "1.0.6",
   "description": "The Ultimate WebSocket server. ws-compatible server, based on uWebSockets.",
   "main": "src/index.js",
   "directories": {

src/request.js

@@ -70,6 +70,11 @@ module.exports = class IncomingMessage {
             localPort: this.app.port,
             remotePort: this.app.port,
             encrypted: this.app.ssl,
+            destroy: () => this.res.close(),
+            end: (data, encoding, callback) => {
+                this.res.end(data, true);
+                if(callback) process.nextTick(callback);
+            },
         };
     }
 

src/server.js

@@ -95,6 +95,9 @@ module.exports = class WebSocketServer extends EventEmitter {
                 }
                 const headers = [];
                 const msg = new IncomingMessage(this, req, res);
+                res.onAborted(() => {
+                    msg.finished = true;
+                });
 
                 if(this.options.verifyClient) {
                     if(this.options.verifyClient.length === 1) {
@@ -113,6 +116,9 @@ module.exports = class WebSocketServer extends EventEmitter {
                                 req: msg,
                                 secure: this.ssl,
                             }, (result, code, name, headers = {}) => {
+                                if(msg.finished) {
+                                    return resolve(false);
+                                }
                                 if(!result) {
                                     res.writeStatus(`${code} ${name}`);
                                     for(const header in headers) {
@@ -130,6 +136,17 @@ module.exports = class WebSocketServer extends EventEmitter {
                     }
                 }
 
+                let onOpen;
+                if(this.options.handleUpgrade) {
+                    const result = await this.options.handleUpgrade(msg);
+                    if(result === false) {
+                        return;
+                    }
+                    if(typeof result === 'function') {
+                        onOpen = result;
+                    }
+                }
+
                 this.emit("headers", headers, msg);
                 res.cork(() => {
                     if(headers.length || this.options.handleProtocols) {
@@ -147,7 +164,7 @@ module.exports = class WebSocketServer extends EventEmitter {
                         protocol = this.options.handleProtocols(protocols, msg);
                     }
                     res.upgrade(
-                        { req: msg },
+                        { req: msg, onOpen },
                         req.getHeader('sec-websocket-key'),
                         protocol ?? req.getHeader('sec-websocket-protocol'),
                         req.getHeader('sec-websocket-extensions'),
@@ -158,7 +175,11 @@ module.exports = class WebSocketServer extends EventEmitter {
             open: (ws) => {
                 ws.client = new this.options.WebSocket(ws, ws.req, this);
                 if(this.clients) this.clients.add(ws.client);
-                this.emit("connection", ws.client, ws.req);
+                if(onOpen) {
+                    onOpen(ws.client, ws.req);
+                } else {
+                    this.emit("connection", ws.client, ws.req);
+                }
             },
             close: (ws) => {
                 ws.client.readyState = WS.CLOSED;