1+ import * as http from 'http' ;
12import _ from 'lodash' ;
23import * as events from 'events' ;
34import express from 'express' ;
45import cors from 'cors' ;
56import corsGate from 'cors-gate' ;
7+ import * as WebSocket from 'ws' ;
68
79import { HtkConfig } from '../config' ;
810import { buildInterceptors } from '../interceptors' ;
@@ -13,6 +15,7 @@ import { ApiModel } from './api-model';
1315import { exposeGraphQLAPI } from './graphql-api' ;
1416import { exposeRestAPI } from './rest-api' ;
1517import { HttpClient } from '../client/http-client' ;
18+ import { UiOperationBridge , getSocketPath } from './ui-operation-bridge' ;
1619
1720/**
1821 * This file contains the core server API, used by the UI to query
@@ -39,6 +42,8 @@ import { HttpClient } from '../client/http-client';
3942export class HttpToolkitServerApi extends events . EventEmitter {
4043
4144 private server : express . Application ;
45+ private bridge : UiOperationBridge ;
46+ private authToken : string | undefined ;
4247
4348 constructor (
4449 config : HtkConfig ,
@@ -47,6 +52,9 @@ export class HttpToolkitServerApi extends events.EventEmitter {
4752 ) {
4853 super ( ) ;
4954
55+ this . bridge = new UiOperationBridge ( ) ;
56+ this . authToken = config . authToken ;
57+
5058 const interceptors = buildInterceptors ( config ) ;
5159
5260 this . server = express ( ) ;
@@ -134,8 +142,64 @@ export class HttpToolkitServerApi extends events.EventEmitter {
134142
135143 start ( ) {
136144 return new Promise < void > ( ( resolve , reject ) => {
137- this . server . listen ( 45457 , '127.0.0.1' , resolve ) ; // Localhost only
138- this . server . once ( 'error' , reject ) ;
145+ const httpServer : http . Server = this . server . listen ( 45457 , '127.0.0.1' , resolve ) ;
146+ httpServer . once ( 'error' , reject ) ;
147+
148+ this . attachWebSocketBridge ( httpServer ) ;
149+ this . startBridgeApiServer ( ) ;
150+ } ) ;
151+ }
152+
153+ private attachWebSocketBridge ( httpServer : http . Server ) {
154+ const wss = new WebSocket . Server ( { noServer : true } ) ;
155+
156+ httpServer . on ( 'upgrade' , ( req : http . IncomingMessage , socket , head ) => {
157+ const url = new URL ( req . url ! , `http://localhost` ) ;
158+
159+ if ( url . pathname !== '/ui-operations' ) {
160+ socket . destroy ( ) ;
161+ return ;
162+ }
163+
164+ // Enforce the same origin restrictions as the REST/GraphQL API
165+ const origin = req . headers [ 'origin' ] || '' ;
166+ if ( ! ALLOWED_ORIGINS . some ( pattern => pattern . test ( origin ) ) ) {
167+ socket . write ( 'HTTP/1.1 403 Forbidden\r\n\r\n' ) ;
168+ socket . destroy ( ) ;
169+ return ;
170+ }
171+
172+ // Check auth token if configured. Browser WebSocket can't send
173+ // headers, so accept token via query param or Authorization header.
174+ if ( this . authToken ) {
175+ const authHeader = req . headers [ 'authorization' ] || '' ;
176+ const tokenMatch = authHeader . match ( / B e a r e r ( \S + ) / ) || [ ] ;
177+ const headerToken = tokenMatch [ 1 ] ;
178+ const queryToken = url . searchParams . get ( 'token' ) ;
179+
180+ if ( headerToken !== this . authToken && queryToken !== this . authToken ) {
181+ socket . write ( 'HTTP/1.1 403 Forbidden\r\n\r\n' ) ;
182+ socket . destroy ( ) ;
183+ return ;
184+ }
185+ }
186+
187+ wss . handleUpgrade ( req , socket as any , head , ( ws ) => {
188+ this . bridge . setWebSocket ( ws as any ) ;
189+ } ) ;
139190 } ) ;
140191 }
192+
193+ private startBridgeApiServer ( ) {
194+ try {
195+ const socketPath = getSocketPath ( ) ;
196+ this . bridge . startApiServer ( socketPath ) ;
197+ } catch ( err : any ) {
198+ console . warn (
199+ `Failed to start UI Bridge socket server: ${ err . message } . ` +
200+ `MCP & remote control will not be available.`
201+ ) ;
202+ }
203+ }
204+
141205} ;
0 commit comments