From 4199e4e3cccaaa1d9d1e9dba9d60a14817614404 Mon Sep 17 00:00:00 2001
From: "Christopher N. KATOYI" <christopher.katoyi@icloud.com>
Date: Fri, 28 Jun 2019 23:41:09 +0200
Subject: [PATCH] [koa-websocket] Added generic types for ctx (#36086)

* added generic types for ctx

* fix: linting warning
---
 types/koa-websocket/index.d.ts             | 22 ++++++++++++----------
 types/koa-websocket/koa-websocket-tests.ts | 22 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/types/koa-websocket/index.d.ts b/types/koa-websocket/index.d.ts
index c23c9036b8..b44be656fc 100644
--- a/types/koa-websocket/index.d.ts
+++ b/types/koa-websocket/index.d.ts
@@ -4,6 +4,7 @@
 //                 Jaco Greeff <https://github.com/jacogr>
 //                 Martin Ždila <https://github.com/zdila>
 //                 Eunchong Yu <https://github.com/Kroisse>
+//                 Christopher N. Katoyi-Kaba <https://github.com/Christopher2K>
 // Definitions: https://github.com/DefinitelyTyped/DefinitelyTyped
 // TypeScript Version: 2.3
 
@@ -21,31 +22,32 @@ declare module "koa" {
 }
 
 declare namespace KoaWebsocket {
-    type Middleware = compose.Middleware<MiddlewareContext>;
+    type Middleware<StateT, CustomT> = compose.Middleware<MiddlewareContext<StateT> & CustomT>;
 
-    interface MiddlewareContext extends Koa.Context {
+    interface MiddlewareContext<StateT> extends Koa.Context {
         // Limitation: Declaration merging cannot overwrap existing properties.
         // That's why this property is here, not in the merged declaration above.
         app: App;
+        state: StateT;
     }
 
-    class Server {
-        app: App;
-        middleware: Middleware[];
+    class Server<StateT = any, CustomT = {}> {
+        app: App<StateT, CustomT>;
+        middleware: Array<Middleware<StateT, CustomT>>;
         server?: ws.Server;
 
-        constructor(app: Koa);
+        constructor(app: Koa<StateT, CustomT>);
 
         listen(options: ws.ServerOptions): ws.Server;
         onConnection(socket: ws, request: http.IncomingMessage): void;
-        use(middleware: Middleware): this;
+        use(middleware: Middleware<StateT, CustomT>): this;
     }
 
-    interface App extends Koa {
-        ws: Server;
+    interface App<StateT = any, CustomT = {}> extends Koa<StateT, CustomT> {
+        ws: Server<StateT, CustomT>;
     }
 }
 
-declare function KoaWebsocket(app: Koa, wsOptions?: ws.ServerOptions, httpsOptions?: https.ServerOptions): KoaWebsocket.App;
+declare function KoaWebsocket<StateT = any, CustomT = {}>(app: Koa<StateT, CustomT>, wsOptions?: ws.ServerOptions, httpsOptions?: https.ServerOptions): KoaWebsocket.App<StateT, CustomT>;
 
 export = KoaWebsocket;
diff --git a/types/koa-websocket/koa-websocket-tests.ts b/types/koa-websocket/koa-websocket-tests.ts
index 502b366652..68044ed26f 100644
--- a/types/koa-websocket/koa-websocket-tests.ts
+++ b/types/koa-websocket/koa-websocket-tests.ts
@@ -20,3 +20,25 @@ app.ws.use(async (ctx, next) => {
 });
 
 app.listen(3000);
+
+interface MyState {
+    persist: string;
+}
+
+const typedApp = websocket(new Koa<MyState>());
+
+typedApp.ws.use(async (ctx, next) => {
+    ctx.websocket.on('message', (message) => {
+        console.log(message + ctx.state.persist);
+        const server = ctx.app.ws.server;
+        if (server) {
+            server.clients.forEach(client => {
+                if (client !== ctx.websocket) {
+                    client.send(message);
+                }
+            });
+        }
+    });
+    ctx.websocket.send('Hello world');
+    await next();
+});