From 9459818906634084a3cf543436a2a6788b12a2a1 Mon Sep 17 00:00:00 2001
From: Josh Kelley <joshkel@gmail.com>
Date: Tue, 27 Jun 2023 14:41:47 -0400
Subject: [PATCH] Improve types

Events such as connect/disconnect/error were typed as accepting full PluginState objects, but the `mode` property was missing.

The `Request.websocket()` method was typed as always returning a `'websocket'` mode object, but it may return an `'http'` mode as well.
---
 hapi-plugin-websocket.d.ts | 13 ++++++++++++-
 hapi-plugin-websocket.js   | 10 +++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/hapi-plugin-websocket.d.ts b/hapi-plugin-websocket.d.ts
index ce49b8e..afb32d9 100644
--- a/hapi-plugin-websocket.d.ts
+++ b/hapi-plugin-websocket.d.ts
@@ -39,6 +39,17 @@ declare namespace HAPIPluginWebsocket {
         initially: boolean
     }
 
+    interface PluginStateHttp {
+        mode: "http"
+        ctx: null
+        wss: null
+        ws: null
+        wsf: null
+        req: null
+        peers: null
+        initially?: null
+    }
+
     interface PluginSpecificConfiguration {
         only?: boolean
         subprotocol?: string
@@ -88,7 +99,7 @@ export = HAPIPluginWebsocket
 
 declare module "@hapi/hapi" {
     export interface Request<Refs extends ReqRef = ReqRefDefaults> extends Podium {
-        websocket(): HAPIPluginWebsocket.PluginState
+        websocket(): HAPIPluginWebsocket.PluginState | HAPIPluginWebsocket.PluginStateHttp
     }
     export interface PluginsStates {
         websocket: HAPIPluginWebsocket.PluginState
diff --git a/hapi-plugin-websocket.js b/hapi-plugin-websocket.js
index dc4a399..95d4237 100644
--- a/hapi-plugin-websocket.js
+++ b/hapi-plugin-websocket.js
@@ -206,7 +206,7 @@ const register = async (server, pluginOptions) => {
             const ctx = {}
 
             /*  allow application to hook into WebSocket connection  */
-            routeOptions.connect.call(ctx, { ctx, wss, ws, wsf, req, peers })
+            routeOptions.connect.call(ctx, { mode: "websocket", ctx, wss, ws, wsf, req, peers })
 
             /*  determine HTTP headers for simulated HTTP request:
                 take headers of initial HTTP upgrade request, but explicitly remove Accept-Encoding,
@@ -253,7 +253,7 @@ const register = async (server, pluginOptions) => {
                 /*  framed WebSocket communication (correlated request/reply)  */
                 wsf.on("message", async (ev) => {
                     /*  allow application to hook into raw WebSocket frame processing  */
-                    routeOptions.frameMessage.call(ctx, { ctx, wss, ws, wsf, req, peers }, ev.frame)
+                    routeOptions.frameMessage.call(ctx, { mode: "websocket", ctx, wss, ws, wsf, req, peers }, ev.frame)
 
                     /*  process frame of expected type only  */
                     if (ev.frame.type === routeOptions.frameRequest) {
@@ -322,7 +322,7 @@ const register = async (server, pluginOptions) => {
             /*  hook into WebSocket disconnection  */
             ws.on("close", () => {
                 /*  allow application to hook into WebSocket disconnection  */
-                routeOptions.disconnect.call(ctx, { ctx, wss, ws, wsf, req, peers })
+                routeOptions.disconnect.call(ctx, { mode: "websocket", ctx, wss, ws, wsf, req, peers })
 
                 /*  stop tracking the peer  */
                 const idx = routePeers[routeId].indexOf(ws)
@@ -331,11 +331,11 @@ const register = async (server, pluginOptions) => {
 
             /*  allow application to hook into WebSocket error processing  */
             ws.on("error", (error) => {
-                routeOptions.error.call(ctx, { ctx, wss, ws, wsf, req, peers }, error)
+                routeOptions.error.call(ctx, { mode: "websocket", ctx, wss, ws, wsf, req, peers }, error)
             })
             if (routeOptions.frame === true) {
                 wsf.on("error", (error) => {
-                    routeOptions.error.call(ctx, { ctx, wss, ws, wsf, req, peers }, error)
+                    routeOptions.error.call(ctx, { mode: "websocket", ctx, wss, ws, wsf, req, peers }, error)
                 })
             }
         })