@microsoft/windows-admin-center-sdk/core/data/powershell-stream.js

Microsoft - Windows Admin Center Shell

import { Observable, throwError } from 'rxjs';
import { catchError, filter, mergeMap, switchMap, take } from 'rxjs/operators';
import { LogLevel } from '../diagnostics/log-level';
import { Logging } from '../diagnostics/logging';
import { PowerShell } from './powershell';
import { WebsocketStreamConnectionState, WebsocketStreamDataRequestState, WebsocketStreamDataState, WebsocketStreamName, WebsocketStreamProcessor } from './websocket-stream';
/**
 * PowerShell Processor interface.
 */
class PowerShellProcessor extends WebsocketStreamProcessor {
}
/**
 * The PowerShell stream class.
 */
export class PowerShellStream {
    websocketStream;
    authorizationManager;
    static logSourceName = 'PowerShellStream';
    /**
     * The collection of set of monitors.
     */
    static monitorSets = [];
    static maxRunPerNode = 5;
    processors = new Map();
    queues = new Map();
    strings = MsftSme.getStrings().MsftSmeShell.Core.WebsocketStream.PowerShellStream;
    /**
     * Register the set of monitors.
     *
     * @param monitorSet The set of monitors.
     */
    static registerMonitorSet(monitorSet) {
        const found = PowerShellStream.monitorSets.find(monitors => monitors.name === monitorSet.name);
        if (found) {
            return;
        }
        PowerShellStream.monitorSets.push(monitorSet);
    }
    /**
     * Unregister the set of monitors.
     *
     * @param name The name of set of monitors.
     * @returns boolean true if unregistered the named set.
     */
    static unregisterMonitors(name) {
        const found = PowerShellStream.monitorSets.find(monitors => monitors.name === name);
        if (found) {
            PowerShellStream.monitorSets.remove(found);
            return true;
        }
        return false;
    }
    /**
     * Initializes a new instance of the PowerShellStream class.
     *
     * @param websocketStream the websocket stream object.
     * @param authorizationManager the authorization manager object.
     */
    constructor(websocketStream, authorizationManager) {
        this.websocketStream = websocketStream;
        this.authorizationManager = authorizationManager;
        websocketStream.registerProcessor(WebsocketStreamName.PowerShellStreamName, this);
    }
    /**
     * PowerShell script run.
     *
     * @param nodeName the node name.
     * @param script the script to run.
     * @param options the options for this request.
     * @return Observable<PowerShellResult> the query observable.
     */
    run(nodeName, commandOrScript, options) {
        const command = PowerShell.getPowerShellCommand(commandOrScript);
        return this.monitorCreateRequest(nodeName, command, PowerShell.newEndpointOptions(options));
    }
    /**
     * Cancel active powershell script.
     * Result response comes back to the original query to end.
     *
     * @param nodeName the node name.
     * @param id the id of original request specified as options.queryId.
     */
    cancel(nodeName, id, options) {
        const target = this.getTarget(nodeName, PowerShell.newEndpointOptions(options));
        const requestState = WebsocketStreamDataRequestState.Cancel;
        const request = { id, target, requestState, script: null };
        // remove from queue if not submitted yet.
        const queue = this.queues.get(target.nodeName);
        if (queue) {
            const pendingRequest = queue.pendingRequests.find(entry => entry.id === id);
            if (pendingRequest) {
                queue.pendingRequests.remove(pendingRequest);
                queue.outstandingCount--;
                const processor = this.processors.get(id);
                this.processors.delete(id);
                processor.complete();
                return;
            }
        }
        this.websocketStream.sendNext(WebsocketStreamName.PowerShellStreamName, request);
    }
    /**
     * Reset data for connection cleanup.
     */
    reset() {
        Logging.log({ level: LogLevel.Warning, message: this.strings.ResetError.message, source: PowerShellStream.logSourceName });
        const processors = [];
        this.processors.forEach((value) => processors.push(value));
        this.processors.clear();
        processors.forEach((processor) => {
            processor.error(new Error(this.strings.ResetError.message));
        });
    }
    /**
     * Process the socket message.
     *
     * @param message the socket message.
     */
    process(message) {
        if (!message) {
            throw new Error(this.strings.NoContentError.message);
        }
        const processor = this.processors.get(message.id);
        if (!processor) {
            Logging.log({ level: LogLevel.Warning, message: this.strings.UnexpectedReceivedError.message, source: PowerShellStream.logSourceName });
            return;
        }
        switch (message.state) {
            case WebsocketStreamDataState.Data:
                this.operationNext(processor, message.response);
                break;
            case WebsocketStreamDataState.Completed:
                this.operationComplete(processor, message.response);
                this.operationEnd(message.id);
                break;
            case WebsocketStreamDataState.Error:
                this.operationError(processor, { xhr: message });
                this.operationEnd(message.id);
                break;
            case WebsocketStreamDataState.Noop:
                break;
        }
    }
    /**
     * Gets the JEA powershell endpoint, if it exists
     *
     * @param nodeName The node name
     */
    getJeaEndpoint(nodeName) {
        return this.authorizationManager.getJeaEndpoint(nodeName);
    }
    /**
     * Gets websocket stream connection state.
     */
    getWebsocketStreamConnectionState() {
        return this.websocketStream.socketStateRaw;
    }
    /**
     * Gets WebSocket State.
     */
    get websocketState() {
        return this.websocketStream.socketState;
    }
    operationNext(processor, response) {
        const partial = processor.options && processor.options.partial;
        // buffering result.
        if (!partial) {
            if (!processor.response) {
                processor.response = response;
            }
            else {
                if (response.errors) {
                    if (!processor.response.errors) {
                        processor.response.errors = response.errors;
                    }
                    else {
                        response.errors.forEach(value => processor.response.errors.push(value));
                    }
                }
                if (response.progress) {
                    if (!processor.response.progress) {
                        processor.response.progress = response.progress;
                    }
                    else {
                        response.progress.forEach(value => processor.response.progress.push(value));
                    }
                }
                if (response.results) {
                    if (!processor.response.results) {
                        processor.response.results = response.results;
                    }
                    else {
                        response.results.forEach(value => processor.response.results.push(value));
                    }
                }
            }
        }
        else {
            processor.next(response);
        }
        return !partial;
    }
    operationComplete(processor, response) {
        if (this.operationNext(processor, response)) {
            processor.next(processor.response);
        }
        processor.complete();
    }
    operationError(processor, error) {
        processor.error(error);
    }
    operationEnd(id) {
        const processor = this.processors.get(id);
        this.processors.delete(id);
        const queue = this.queues.get(processor.target.nodeName);
        if (--queue.outstandingCount === 0) {
            this.queues.delete(processor.target.nodeName);
        }
        if (queue.pendingRequests.length > 0) {
            // if there is queued item, then send request.
            const request = queue.pendingRequests.shift();
            this.websocketStream.sendNext(WebsocketStreamName.PowerShellStreamName, request);
        }
    }
    createRequest(nodeName, command, options) {
        // publish object is created two ways.
        // 1) socket is connected so submit the request immediately with simple observable.
        //   (if-block and this is the most of cases.)
        // 2) socket is not connected so wait for the socket to ready and submit request with
        //    complex observable. Initial connect and re-connection takes this observable.
        //   (else-block and this is a few cases.)
        let publish;
        const endpoint = this.authorizationManager.getJeaEndpoint(nodeName);
        const newOptions = { ...(options || {}) };
        if (endpoint) {
            newOptions.powerShellEndpoint = endpoint;
        }
        if (this.websocketStream.socketStateRaw === WebsocketStreamConnectionState.Connected) {
            publish = this.createRequestSimple(nodeName, command, newOptions);
        }
        else {
            publish = this.websocketStream.socketState
                .pipe(filter(state => state === WebsocketStreamConnectionState.Connected
                || state === WebsocketStreamConnectionState.Failed
                || state === WebsocketStreamConnectionState.NotConfigured), take(1), mergeMap(state => {
                if (state === WebsocketStreamConnectionState.Connected) {
                    return this.createRequestSimple(nodeName, command, newOptions);
                }
                return throwError(() => new Error(this.strings.ConnectionError.message));
            }));
        }
        return publish
            .pipe(catchError((error) => {
            // retry if reset connection of socket was observed.
            if (error && error.message === this.strings.ResetError.message) {
                return this.monitorCreateRequest(nodeName, command, newOptions);
            }
            if ((!options || options.noAuth !== true) && !this.authorizationManager.signOnManager.isSignOnTokenEnabled) {
                if (this.authorizationManager.canHandleStreamFailure(error.xhr)) {
                    return this.authorizationManager.handleStreamFailure(nodeName, options, error.xhr)
                        .pipe(switchMap(updatedOptions => this.monitorCreateRequest(nodeName, command, updatedOptions)));
                }
            }
            if (this.authorizationManager.signOnManager.isSignOnTokenEnabled) {
                if (this.authorizationManager.signOnManager.canHandleStreamUnauthorizedLogin(error.xhr)) {
                    return this.authorizationManager.signOnManager.handleStreamUnauthorizedLogin(options, error.xhr)
                        .pipe(switchMap(updatedOptions => this.monitorCreateRequest(nodeName, command, updatedOptions)));
                }
            }
            return throwError(() => error);
        }));
    }
    createRequestSimple(nodeName, command, options) {
        return new Observable(observer => {
            const target = this.getTarget(nodeName, options);
            const requestState = WebsocketStreamDataRequestState.Normal;
            const id = this.sendRequest(observer, target, requestState, command, options);
            return () => {
                const processor = this.processors.get(id);
                if (processor) {
                    processor.end = true;
                    if (!processor.closed && !processor.closing) {
                        this.cancel(processor.target.nodeName, id);
                    }
                }
            };
        });
    }
    sendRequest(observer, target, requestState, command, options) {
        const id = (options && options.queryId) || MsftSme.getUniqueId();
        const request = {
            ...{
                id,
                target,
                requestState,
                options
            },
            ...command
        };
        const processor = new PowerShellProcessor(observer, target, options);
        const queue = this.queues.get(target.nodeName);
        this.processors.set(id, processor);
        // During a send request, if caller provides 'options.close' as true,
        // we shouldn't manage the request via a queue and on Gateway, we should create a
        // new, one time use Runspace, which is disposed after use, instead of using one from the pool.
        // As currently this is not handled on Gateway, just ignore the 'options.close' for now.
        /*
        if (options && options.close) {
            // disposing session.
            this.websocketStream.sendNext(WebsocketStreamName.PowerShellStreamName, request);
            return id;
        }
        */
        if (++queue.outstandingCount > PowerShellStream.maxRunPerNode) {
            queue.pendingRequests.push(request);
            return id;
        }
        this.websocketStream.sendNext(WebsocketStreamName.PowerShellStreamName, request);
        return id;
    }
    getTarget(nodeName, options) {
        if (!this.queues.has(nodeName)) {
            const queue = { outstandingCount: 0, pendingRequests: [] };
            this.queues.set(nodeName, queue);
        }
        return this.websocketStream.getTarget(this.authorizationManager, nodeName, options.powerShellEndpoint);
    }
    monitorCreateRequest(nodeName, command, options) {
        let monitored = (nodeName1, command1, options1) => this.createRequest(nodeName1, command1, options1);
        for (const monitorSet of PowerShellStream.monitorSets) {
            monitored = this.monitor(monitored, monitorSet);
        }
        return monitored(nodeName, command, options);
    }
    monitor(target, monitorSet) {
        return function (nodeName, command, options) {
            let context;
            return monitorSet.preMonitor(nodeName, command, options)
                .pipe(switchMap(packet => {
                context = packet;
                return target(packet.nodeName, packet.command, packet.options);
            }), catchError((error) => monitorSet.errorMonitor(error, context)), switchMap(response => monitorSet.successMonitor(response, context)));
        };
    }
}

//# sourceMappingURL=powershell-stream.js.map