diff --git a/bin/chii.js b/bin/chii.js index ac79a7a5..a409dbfe 100755 --- a/bin/chii.js +++ b/bin/chii.js @@ -17,6 +17,7 @@ program .option('--https', 'serve chii over https') .option('--ssl-cert ', 'provide an ssl certificate') .option('--ssl-key ', 'provide an ssl key') + .option('--secrets ', 'provide secrets', (value) => !value ? [] : value.split(',')) .action(options => { server.start(options); }); diff --git a/server/index.js b/server/index.js index 38a70b4c..30c96999 100644 --- a/server/index.js +++ b/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)); diff --git a/server/lib/WebSocketServer.js b/server/lib/WebSocketServer.js index 840979d9..93402a84 100644 --- a/server/lib/WebSocketServer.js +++ b/server/lib/WebSocketServer.js @@ -4,7 +4,7 @@ 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 })); @@ -12,7 +12,16 @@ module.exports = class WebSocketServer { 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; } diff --git a/src/target/config.ts b/src/target/config.ts index 0c377525..81aae6c6 100644 --- a/src/target/config.ts +++ b/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 }; diff --git a/src/target/connectServer.ts b/src/target/connectServer.ts index b3ad0321..6665df1b 100644 --- a/src/target/connectServer.ts +++ b/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 } : {}), })}` );