import * as vscode from 'vscode' ; import * as WebSocket from 'ws' ; import * as http from 'http' ; import * as fs from 'fs' ; import * as path from 'path' ; import { URL } from 'url' ; import { randomBytes } from 'crypto' ; import { Disposable } from '../utils/dispose' ; import { isFileInjectable } from '../utils/utils' ; import TelemetryReporter from 'vscode-extension-telemetry' ; import { EndpointManager } from '../infoManagers/endpointManager' ; import { WorkspaceManager } from '../infoManagers/workspaceManager' ; import { UriSchemes } from '../utils/constants' ; import { ConnectionManager } from '../infoManagers/connectionManager' ; description export class WSServerWithOriginCheck extends WebSocket . Server { public externalHostName : string | undefined ; param { http . IncomingMessage } req returns { boolean } public shouldHandle ( req ) { const origin = req . headers [ 'origin' ] ; return < boolean > ( ( super . shouldHandle ( req ) && origin && ( origin . startsWith ( UriSchemes . vscode_webview ) || ( this . externalHostName && origin === this . externalHostName ) ) ) ) ; } } description export class WSServer extends Disposable { private _wss : WSServerWithOriginCheck | undefined ; private _wsPath ! : string ; private _wsPort = 0 ; public set externalHostName ( hostName ) { if ( this . _wss ) { this . _wss . externalHostName = hostName ; } } constructor ( private readonly _reporter , private readonly _endpointManager , private readonly _workspaceManager , private readonly _connectionManager ) { super ( ) ; this . _register ( _connectionManager . onConnected ( ( e ) => { this . externalHostName = `${ e . httpURI . scheme }://${ e . httpURI . authority }` ; } ) ) ; } description private get _basePath ( ) : string | undefined { return this . _workspaceManager . workspacePath ; } public get wsPath ( ) { return this . _wsPath ; } public get wsPort ( ) { return this . _wsPort ; } public set wsPort ( portNum ) { this . _wsPort = portNum ; } private readonly _onConnected = this . _register ( new vscode . EventEmitter < void > ( ) ) ; public readonly onConnected = this . _onConnected . event ; description param { number } wsPort public start ( wsPort ) { this . _wsPort = wsPort ; this . _wsPath = `/${ randomBytes ( 20 ) . toString ( 'hex' ) }` ; this . startWSServer ( this . _basePath ?? '' ) ; } description public close ( ) { if ( this . _wss != null ) { this . _wss . close ( ) ; } } description public refreshBrowsers ( ) { if ( this . _wss ) { this . _wss . clients . forEach ( ( client ) => client . send ( JSON . stringify ( { command : 'reload' } ) ) ) ; } } param { string } basePath returns { boolean } private startWSServer ( basePath ) { this . _wss = new WSServerWithOriginCheck ( { port : this . _wsPort , host : this . _connectionManager . host , path : this . _wsPath } ) ; this . _wss . on ( 'connection' , ( ws ) => this . handleWSConnection ( basePath , ws ) ) ; this . _wss . on ( 'error' , ( err ) => this . handleWSError ( basePath , err ) ) ; this . _wss . on ( 'listening' , ( ) => this . handleWSListen ( ) ) ; return true ; } param { string } basePath param { any } err private handleWSError ( basePath , err ) { if ( err . code == 'EADDRINUSE' ) { this . _wsPort ++ ; this . startWSServer ( basePath ) ; } else if ( err . code == 'EADDRNOTAVAIL' ) { this . _connectionManager . resetHostToDefault ( ) ; this . startWSServer ( basePath ) ; } else { this . _reporter . sendTelemetryErrorEvent ( 'server.err' , { type : 'ws' , err : err , } ) ; console . log ( `Unknown error: ${ err }` ) ; } } description private handleWSListen ( ) { console . log ( `Websocket server is running on port ${ this . _wsPort }` ) ; this . _onConnected . fire ( ) ; } description param { string } basePath param { WebSocket } ws private handleWSConnection ( basePath , ws ) { ws . on ( 'message' , ( message ) => { const parsedMessage = JSON . parse ( message ) ; switch ( parsedMessage . command ) { case 'urlCheck' : { const results = this . performTargetInjectableCheck ( basePath , parsedMessage . url ) ; if ( ! results . injectable ) { this . _reporter . sendTelemetryEvent ( 'server.ws.foundNonInjectable' ) ; const sendData = { command : 'foundNonInjectable' , path : results . pathname , } ; ws . send ( JSON . stringify ( sendData ) ) ; } } } } ) ; } description param { string } basePath param { string } urlString returns { boolean private performTargetInjectableCheck ( basePath , urlString ) { const url = new URL ( urlString ) ; let absolutePath = path . join ( basePath , url . pathname ) ; if ( ! fs . existsSync ( absolutePath ) ) { const decodedLocation = this . _endpointManager . decodeLooseFileEndpoint ( absolutePath ) ; if ( ! decodedLocation || ! fs . existsSync ( decodedLocation ) ) { return { injectable : false , pathname : url . pathname } ; } else { absolutePath = decodedLocation ; } } if ( fs . statSync ( absolutePath ) . isDirectory ( ) || isFileInjectable ( absolutePath ) ) { return { injectable : true , pathname : url . pathname } ; } return { injectable : false , pathname : url . pathname } ; } }