Add optional secret to websocket communication

bin/chii.js

@@ -17,6 +17,7 @@ program
   .option('--https', 'serve chii over https')
   .option('--ssl-cert <cert>', 'provide an ssl certificate')
   .option('--ssl-key <key>', 'provide an ssl key')
+  .option('--secrets <secret1,secret2...>', 'provide secrets', (value) => !value ? [] : value.split(','))
   .action(options => {
     server.start(options);
   });

server/index.js

@@ -18,14 +18,14 @@ async function start({
   sslCert,
   sslKey,
   basePath = '/',
+  secrets = [],
 } = {}) {
   domain = domain || 'localhost:' + port;
   if (!endWith(basePath, '/')) {
     basePath += '/';
   }
-
   const app = new Koa();
-  const wss = new WebSocketServer();
+  const wss = new WebSocketServer( secrets );
 
   app.use(compress()).use(router(wss.channelManager, domain, cdn, basePath));
 

server/lib/WebSocketServer.js

@@ -4,15 +4,24 @@ const ChannelManager = require('./ChannelManager');
 const query = require('licia/query');
 
 module.exports = class WebSocketServer {
-  constructor() {
+  constructor( secrets ) {
     this.channelManager = new ChannelManager();
 
     const wss = (this._wss = new WebSocket.Server({ noServer: true }));
 
     wss.on('connection', (ws, req) => {
       const type = ws.type;
       if (type === 'target') {
-        const { id, chiiUrl, title, favicon } = ws;
+        const { id, chiiUrl, title, favicon, secret: targetSecret } = ws;
+
+        const isSecretRequired = Array.isArray(secrets) && secrets.length > 0;
+        const isSecretValid = targetSecret && secrets.includes(targetSecret);
+
+        if (isSecretRequired && !isSecretValid) {
+          console.log('wrong or no secret provided');
+          ws.close();
+        }
+
         let ip = req.socket.remoteAddress;
         const userAgent = req.headers['user-agent'];
         if (req.headers['x-forwarded-for']) {
@@ -47,6 +56,7 @@ module.exports = class WebSocketServer {
             ws.favicon = q.favicon;
             ws.userAgent = q.userAgent;
             ws.rtc = q.rtc === 'true';
+            ws.secret = q.secret;
           } else {
             ws.target = q.target;
           }

src/target/config.ts

@@ -31,6 +31,7 @@ if (!startWith(serverUrl, 'http')) {
 let embedded = false;
 let rtc = false;
 let cdn = '';
+let secret = '';
 
 const element = getTargetScriptEl();
 if (element) {
@@ -41,6 +42,7 @@ if (element) {
     rtc = true;
   }
   cdn = element.getAttribute('cdn') || '';
+  secret = element.getAttribute('secret') || '';
 }
 
 if (cdn && endWith(cdn, '/')) {
@@ -51,7 +53,7 @@ const sessionStore = safeStorage('session');
 
 let id = sessionStore.getItem('chii-id');
 if (!id) {
-  id = randomId(6);
+  id = randomId(32);
   sessionStore.setItem('chii-id', id);
 }
 
@@ -62,4 +64,5 @@ export {
   rtc,
   cdn,
   id,
+  secret
 };

src/target/connectServer.ts

@@ -1,7 +1,7 @@
 import Socket from 'licia/Socket';
 import query from 'licia/query';
 import chobitsu from 'chobitsu';
-import { serverUrl, id } from './config';
+import { serverUrl, id, secret } from './config';
 import { getFavicon } from './util';
 
 let isInit = false;
@@ -24,6 +24,7 @@ export default function () {
       title: (window as any).ChiiTitle || document.title,
       favicon: getFavicon(),
       '__chobitsu-hide__': true,
+      ...(secret ? { secret } : {}),
     })}`
   );