Merge pull request #725 from frsechet/websocket-fixes

Add DELETE route to close a websocket connection
dherault · Jul 2, 2019 · 28266f2 · 28266f2
2 parents 1ba79f2 + e70db4b
commit 28266f2
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 0 deletions.
diff --git a/manual_test_websocket/main/test/e2e/ws.e2e.js b/manual_test_websocket/main/test/e2e/ws.e2e.js
@@ -370,5 +370,29 @@ describe('serverless', () => {
 
       expect(res).to.have.status(410);
     }).timeout(timeout);
+
+    it('should be able to close connections via REST API', async () => {
+      await createClient();
+      const c2 = await createClient();
+      const url = new URL(endpoint);
+      const signature = { service: 'execute-api', host: url.host, path: `${url.pathname}/@connections/${c2.id}`, method: 'DELETE', body: 'Hello World!', headers: { 'Content-Type': 'text/plain'/* 'application/text' */ } };
+      aws4.sign(signature, { accessKeyId: cred.accessKeyId, secretAccessKey: cred.secretAccessKey });
+      const res = await req.del(signature.path.replace(url.pathname, '')).set('X-Amz-Date', signature.headers['X-Amz-Date']).set('Authorization', signature.headers.Authorization).set('Content-Type', signature.headers['Content-Type']);
+
+      expect(res).to.have.status(200);
+    }).timeout(timeout);
+
+    it('should receive error code when deleting a previously closed client via REST API', async () => {
+      const c = await createClient();
+      const cId = c.id;
+      c.ws.close();
+      const url = new URL(endpoint);
+      const signature = { service: 'execute-api', host: url.host, path: `${url.pathname}/@connections/${cId}`, method: 'DELETE', body: 'Hello World!', headers: { 'Content-Type': 'text/plain'/* 'application/text' */ } };
+      aws4.sign(signature, { accessKeyId: cred.accessKeyId, secretAccessKey: cred.secretAccessKey });
+      const res = await req.del(signature.path.replace(url.pathname, '')).set('X-Amz-Date', signature.headers['X-Amz-Date']).set('Authorization', signature.headers.Authorization).set('Content-Type', signature.headers['Content-Type']);
+
+      expect(res).to.have.status(410);
+    }).timeout(timeout);
+
   });
 });
diff --git a/src/ApiGatewayWebSocket.js b/src/ApiGatewayWebSocket.js
@@ -24,6 +24,10 @@ module.exports = class ApiGatewayWebSocket {
     this.websocketsApiRouteSelectionExpression = serverless.service.provider.websocketsApiRouteSelectionExpression || '$request.body.action';
   }
 
+  printBlankLine() {
+    console.log();
+  }
+
   _createWebSocket() {
     // start COPY PASTE FROM HTTP SERVER CODE
     const serverOptions = {
@@ -248,6 +252,33 @@ module.exports = class ApiGatewayWebSocket {
         return '';
       },
     });
+
+    this.wsServer.route({
+      method: 'DELETE',
+      path: '/@connections/{connectionId}',
+      config: { payload: { parse: false } },
+      handler: (request, h) => {
+        debugLog(`got DELETE to ${request.url}`);
+
+        const getByConnectionId = (map, searchValue) => {
+          for (const [key, connection] of map.entries()) {
+            if (connection.connectionId === searchValue) return key;
+          }
+
+          return undefined;
+        };
+
+        const ws = getByConnectionId(this.clients, request.params.connectionId);
+
+        if (!ws) return h.response().code(410);
+
+        ws.close();
+
+        debugLog(`closed connection:${request.params.connectionId}`);
+
+        return '';
+      },
+    });
   }
 
   _createWsAction(fun, funName, servicePath, funOptions, event) {