Skip to content

Commit 7cb3972

Browse files
kevinwang5658kevinwang5658
committed
chore: refactored apply flow to use express
1 parent 83bc13f commit 7cb3972

File tree

6 files changed

+130
-144
lines changed

6 files changed

+130
-144
lines changed

‎src/config.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
export const config = {
22
loginServerPort: 51_039,
3+
connectServerPort: 51_040,
4+
35
corsAllowedOrigins: [
46
'https://codify-dashboard.com',
57
'https://https://codify-dashboard.kevinwang5658.workers.dev',

‎src/connect/apply.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { WsServerManager } from './server.js';
1010
export function connectApplyInitHandler(msg: any, initWs: WebSocket, manager: WsServerManager) {
1111
const sessionId = uuid();
1212

13-
manager.startAdhocWsServer(sessionId, async (ws) => {
13+
manager.addAdhocWsServer(sessionId, async (ws) => {
1414
console.log('connected apply ws');
1515

1616
const { config } = msg;

‎src/connect/http-route-handler.ts‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import { spawn } from '@homebridge/node-pty-prebuilt-multiarch';
2+
import { ResourceConfig } from 'codify-schemas';
3+
import { Router } from 'express';
4+
import * as fs from 'node:fs/promises';
5+
import * as os from 'node:os';
6+
import * as path from 'node:path';
7+
import { v4 as uuid } from 'uuid';
8+
import { WebSocket } from 'ws';
9+
10+
import { WsServerManager } from './server.js';
11+
12+
const router = Router({
13+
mergeParams: true,
14+
});
15+
16+
router.post('/apply', (req, res) => {
17+
const sessionId = uuid();
18+
const manager = WsServerManager.get();
19+
const config = req.body;
20+
21+
manager.addAdhocWsServer(sessionId, async (ws: WebSocket) => wsHandler(ws, config));
22+
23+
return res.status(200).json({ sessionId });
24+
});
25+
26+
async function wsHandler(ws: WebSocket, config: ResourceConfig): Promise<void> {
27+
const tmpDir = await fs.mkdtemp(os.tmpdir());
28+
const filePath = path.join(tmpDir, 'codify.json');
29+
30+
await fs.writeFile(filePath, JSON.stringify(config));
31+
32+
const pty = spawn('zsh', ['-c', 'codify apply'], {
33+
name: 'xterm-color',
34+
cols: 80,
35+
rows: 30,
36+
cwd: process.env.HOME,
37+
env: process.env
38+
});
39+
40+
pty.onData((data) => {
41+
ws.send(Buffer.from(data, 'utf8'));
42+
});
43+
44+
ws.on('message', (message) => {
45+
pty.write(message.toString('utf8'));
46+
})
47+
48+
pty.onExit(({ exitCode, signal }) => {
49+
console.log('pty exit', exitCode, signal);
50+
ws.terminate();
51+
})
52+
}
53+
54+
export default router;

‎src/connect/server.ts‎

Lines changed: 57 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
import { IncomingMessage, Server, createServer } from 'node:http';
2+
import { Duplex } from 'node:stream';
23
import { v4 as uuid } from 'uuid';
34
import { WebSocket, WebSocketServer } from 'ws';
5+
import { config } from '../config.js';
46

5-
const DEFAULT_PORT = 51_040;
7+
let instance: WsServerManager | undefined;
68

79
export class WsServerManager {
810

@@ -14,19 +16,27 @@ export class WsServerManager {
1416

1517
private connectionSecret;
1618

17-
constructor(connectionSecret?: string) {
18-
this.server = createServer();
19+
static init(server: Server, connectionSecret?: string): WsServerManager {
20+
instance = new WsServerManager(server, connectionSecret);
21+
return instance;
22+
}
23+
24+
static get(): WsServerManager {
25+
if (!instance) {
26+
throw new Error('You must call WsServerManager.init before using it');
27+
}
28+
29+
return instance;
30+
}
31+
32+
private constructor(server: Server, connectionSecret?: string) {
33+
this.server = server
1934
this.connectionSecret = connectionSecret;
2035
this.wsServerMap.set('default', this.createWssServer());
2136

22-
this.initServer();
37+
this.server.on('upgrade', this.onUpgrade)
2338
}
2439

25-
listen(cb?: () => void, port?: number, ) {
26-
this.port = port ?? DEFAULT_PORT
27-
this.server.listen(this.port, 'localhost', cb);
28-
}
29-
3040
setDefaultHandler(handler: (ws: WebSocket, manager: WsServerManager) => void): WsServerManager {
3141
const wss = this.createWssServer();
3242
this.wsServerMap.set('default', wss);
@@ -35,73 +45,50 @@ export class WsServerManager {
3545
return this;
3646
}
3747

38-
addAdditionalHandlers(path: string, handler: (ws: WebSocket) => void): WsServerManager {
39-
this.handlerMap.set(path, () => {
40-
const wss = this.addWebsocketServer();
41-
42-
});
43-
44-
return this;
45-
}
46-
47-
startAdhocWsServer(sessionId: string, handler: (ws: WebSocket, manager: WsServerManager) => void) {
48+
addAdhocWsServer(sessionId: string, handler: (ws: WebSocket, manager: WsServerManager) => void) {
4849
this.wsServerMap.set(sessionId, this.createWssServer());
4950
this.handlerMap.set(sessionId, handler);
5051
}
5152

52-
private addWebsocketServer(): string {
53-
const key = uuid();
54-
55-
const wss = new WebSocketServer({
56-
noServer: true
57-
})
58-
this.wsServerMap.set(key, wss);
59-
60-
wss.on('close', () => {
61-
this.wsServerMap.delete(key);
62-
})
63-
64-
return key;
53+
private onUpgrade = (request: IncomingMessage, socket: Duplex, head: Buffer): void => {
54+
const { pathname } = new URL(request.url!, 'ws://localhost:51040')
55+
56+
if (!this.validateOrigin(request.headers.origin!)
57+
|| this.validateConnectionSecret(request)) {
58+
console.error('Unauthorized request from', request.headers.origin);
59+
socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n')
60+
socket.destroy();
61+
return;
62+
}
63+
64+
if (pathname === '/ws' && this.handlerMap.has('default')) {
65+
const wss = this.wsServerMap.get('default');
66+
wss?.handleUpgrade(request, socket, head, (ws, request) => this.handlerMap.get('default')!(ws, this, request));
67+
return;
68+
}
69+
70+
const pathSections = pathname.split('/').filter(Boolean);
71+
if (
72+
pathSections[0] === 'ws'
73+
&& pathSections[1] === 'session'
74+
&& pathSections[2]
75+
&& this.handlerMap.has(pathSections[2])
76+
) {
77+
const sessionId = pathSections[2];
78+
console.log('session found, upgrading', sessionId);
79+
80+
const wss = this.wsServerMap.get(sessionId)!;
81+
82+
wss.handleUpgrade(request, socket, head, (ws, request) => this.handlerMap.get(sessionId)!(ws, this, request));
83+
}
6584
}
6685

67-
private initServer() {
68-
this.server.on('upgrade', (request, socket, head) => {
69-
console.log('upgrade')
70-
71-
const { pathname } = new URL(request.url!, 'ws://localhost:51040')
72-
console.log('Pathname:', pathname)
73-
74-
const code = request.headers['sec-websocket-protocol']
75-
if (this.connectionSecret && code !== this.connectionSecret) {
76-
console.log('Auth failed');
77-
socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n')
78-
socket.destroy()
79-
return;
80-
}
81-
82-
if (pathname === '/' && this.handlerMap.has('default')) {
83-
const wss = this.wsServerMap.get('default');
84-
wss?.handleUpgrade(request, socket, head, (ws, request) => this.handlerMap.get('default')!(ws, this, request));
85-
return;
86-
}
87-
88-
const pathSections = pathname.split('/').filter(Boolean);
89-
console.log(pathSections);
90-
console.log('available sessions', this.handlerMap)
91-
92-
if (pathSections[0] === 'session'
93-
&& pathSections[1]
94-
&& this.handlerMap.has(pathSections[1])
95-
) {
96-
const sessionId = pathSections[1];
97-
console.log('session found, upgrading', sessionId);
98-
99-
const wss = this.wsServerMap.get(sessionId)!;
100-
101-
wss.handleUpgrade(request, socket, head, (ws, request) => this.handlerMap.get(sessionId)!(ws, this, request));
102-
return;
103-
}
104-
})
86+
private validateOrigin = (origin: string): boolean =>
87+
config.corsAllowedOrigins.includes(origin)
88+
89+
private validateConnectionSecret = (request: IncomingMessage): boolean => {
90+
const connectionSecret = request.headers['connection-secret'] as string;
91+
return connectionSecret === this.connectionSecret;
10592
}
10693

10794
private createWssServer(): WebSocketServer {

src/connect/route-handler.ts renamed to src/connect/ws-route-handler.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import { connectApplyInitHandler } from './apply.js';
44
import { WsServerManager } from './server.js';
55

66
export async function defaultWsHandler(ws: WebSocket, manager: WsServerManager) {
7+
console.log('[WS] Connection opened');
8+
79
ws.on('message', (message) => {
810
let msg;
911
try {

‎src/orchestrators/connect.ts‎

Lines changed: 14 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -1,90 +1,31 @@
1-
import { spawn } from '@homebridge/node-pty-prebuilt-multiarch';
1+
import cors from 'cors';
2+
import express, { json } from 'express';
23
import { randomBytes } from 'node:crypto';
34
import open from 'open';
45
import { WebSocket } from 'ws';
56

6-
import { defaultWsHandler } from '../connect/route-handler.js';
7+
import { config } from '../config.js';
8+
import HttpRouteHandler from '../connect/http-route-handler.js';
79
import { WsServerManager } from '../connect/server.js';
10+
import { defaultWsHandler } from '../connect/ws-route-handler.js';
811

912
export class ConnectOrchestrator {
1013
static async run() {
1114
const connectionSecret = ConnectOrchestrator.tokenGenerate()
12-
console.log(connectionSecret)
13-
14-
const server = new WsServerManager(connectionSecret)
15-
.setDefaultHandler(defaultWsHandler)
16-
.addAdditionalHandlers('/apply-logs', () => {})
17-
.addAdditionalHandlers('/import-logs', () => {})
18-
.addAdditionalHandlers('/terminal', () => {})
19-
20-
server.listen(() => {
15+
const app = express();
16+
17+
app.use(cors({ origin: config.corsAllowedOrigins }))
18+
app.use(json())
19+
app.use(HttpRouteHandler);
20+
21+
const server = app.listen(config.connectServerPort, () => {
2122
open(`http://localhost:3000/connection/success?code=${connectionSecret}`)
2223
});
23-
}
24-
25-
private static onConnection(ws: WebSocket) {
26-
console.log('[WS] Connection opened');
27-
28-
ws.on('apply', (message) => {
29-
let data;
30-
try {
31-
data = JSON.parse(message.toString('utf8'));
32-
console.log(data);
33-
} catch (error) {
34-
console.error(error);
35-
}
3624

37-
});
38-
39-
ws.on('close', () => {
40-
41-
});
25+
const wsManager = WsServerManager.init(server, connectionSecret)
26+
.setDefaultHandler(defaultWsHandler)
4227
}
4328

44-
/*
45-
private static async onApply(ws: WebSocket, data: any) {
46-
const { config } = data;
47-
const tmpDir = await fs.mkdtemp(os.tmpdir());
48-
const filePath = path.join(tmpDir, 'codify.json');
49-
50-
await fs.writeFile(filePath, JSON.stringify(config));
51-
52-
const server = createServer()
53-
const wss = new WebSocketServer({
54-
55-
})
56-
57-
server.on('upgrade', (request, socket, head) => {
58-
wss.handleUpgrade(request, socket, head, (ws2) => {
59-
const pty = spawn('zsh', ['-c', '"codify apply"'], {
60-
name: 'xterm-color',
61-
cols: 80,
62-
rows: 30,
63-
cwd: process.env.HOME,
64-
env: process.env
65-
});
66-
67-
pty.onData((data) => {
68-
ws2.send(Buffer.from(data, 'utf8'));
69-
});
70-
71-
ws2.on('message', (message) => {
72-
pty.write(message.toString('utf8'));
73-
})
74-
75-
pty.onExit((code) => {
76-
ws2.close(code, code);
77-
})
78-
})
79-
});
80-
81-
server.listen(2123, () => {
82-
ws.emit('apply_Response', {
83-
wsPass: 'pass',
84-
})
85-
})
86-
} */
87-
8829
private static tokenGenerate(length = 20): string {
8930
return Buffer.from(randomBytes(length)).toString('hex')
9031
}

0 commit comments

Comments
 (0)