Tweaks to transformer

src/vs/base/common/uriIpc.ts

@@ -8,9 +8,17 @@ import { URI, UriComponents } from 'vs/base/common/uri';
 export interface IURITransformer {
 	transformIncoming(uri: UriComponents): UriComponents;
 	transformOutgoing(uri: URI): URI;
+	transformOutgoing(uri: UriComponents): UriComponents;
 }
 
-export const DefaultURITransformer: IURITransformer = {
-	transformIncoming: (uri: UriComponents) => uri,
-	transformOutgoing: (uri: URI) => uri,
+export const DefaultURITransformer: IURITransformer = new class {
+	transformIncoming(uri: UriComponents) {
+		return uri;
+	}
+
+	transformOutgoing(uri: URI): URI;
+	transformOutgoing(uri: UriComponents): UriComponents;
+	transformOutgoing(uri: URI | UriComponents): URI | UriComponents {
+		return uri;
+	}
 };
\ No newline at end of file

src/vs/workbench/node/extensionHostMain.ts

@@ -316,11 +316,11 @@ export class ExtensionHostMain {
 	}
 
 	private transform(initData: IInitData, rpcProtocol: RPCProtocol): IInitData {
-		initData.extensions.forEach((ext) => (<any>ext).extensionLocation = URI.revive(ext.extensionLocation));
-		initData.environment.appRoot = URI.revive(initData.environment.appRoot);
-		initData.environment.appSettingsHome = URI.revive(initData.environment.appSettingsHome);
-		initData.environment.extensionDevelopmentLocationURI = URI.revive(initData.environment.extensionDevelopmentLocationURI);
-		initData.logsLocation = URI.revive(initData.logsLocation);
+		initData.extensions.forEach((ext) => (<any>ext).extensionLocation = URI.revive(rpcProtocol.transformIncomingURIs(ext.extensionLocation)));
+		initData.environment.appRoot = URI.revive(rpcProtocol.transformIncomingURIs(initData.environment.appRoot));
+		initData.environment.appSettingsHome = URI.revive(rpcProtocol.transformIncomingURIs(initData.environment.appSettingsHome));
+		initData.environment.extensionDevelopmentLocationURI = URI.revive(rpcProtocol.transformIncomingURIs(initData.environment.extensionDevelopmentLocationURI));
+		initData.logsLocation = URI.revive(rpcProtocol.transformIncomingURIs(initData.logsLocation));
 		initData.workspace = rpcProtocol.transformIncomingURIs(initData.workspace);
 		return initData;
 	}

src/vs/workbench/services/extensions/node/rpcProtocol.ts

@@ -17,6 +17,30 @@ import { IMessagePassingProtocol } from 'vs/base/parts/ipc/node/ipc';
 import { LazyPromise } from 'vs/workbench/services/extensions/node/lazyPromise';
 import { IRPCProtocol, ProxyIdentifier, getStringIdentifierForProxy } from 'vs/workbench/services/extensions/node/proxyIdentifier';
 
+export interface JSONStringifyReplacer {
+	(key: string, value: any): any;
+}
+
+function safeStringify(obj: any, replacer: JSONStringifyReplacer | null): string {
+	try {
+		return JSON.stringify(obj, <(key: string, value: any) => any>replacer);
+	} catch (err) {
+		return 'null';
+	}
+}
+
+function createURIReplacer(transformer: IURITransformer | null): JSONStringifyReplacer | null {
+	if (!transformer) {
+		return null;
+	}
+	return (key: string, value: any): any => {
+		if (value && value.$mid === 1) {
+			return transformer.transformOutgoing(value);
+		}
+		return value;
+	};
+}
+
 function _transformOutgoingURIs(obj: any, transformer: IURITransformer, depth: number): any {
 
 	if (!obj || depth > 200) {
@@ -113,6 +137,7 @@ export class RPCProtocol extends Disposable implements IRPCProtocol {
 	private readonly _protocol: IMessagePassingProtocol;
 	private readonly _logger: IRPCProtocolLogger | null;
 	private readonly _uriTransformer: IURITransformer | null;
+	private readonly _uriReplacer: JSONStringifyReplacer | null;
 	private _isDisposed: boolean;
 	private readonly _locals: any[];
 	private readonly _proxies: any[];
@@ -129,6 +154,7 @@ export class RPCProtocol extends Disposable implements IRPCProtocol {
 		this._protocol = protocol;
 		this._logger = logger;
 		this._uriTransformer = transformer;
+		this._uriReplacer = createURIReplacer(this._uriTransformer);
 		this._isDisposed = false;
 		this._locals = [];
 		this._proxies = [];
@@ -353,10 +379,7 @@ export class RPCProtocol extends Disposable implements IRPCProtocol {
 
 		promise.then((r) => {
 			delete this._cancelInvokedHandlers[callId];
-			if (this._uriTransformer) {
-				r = transformOutgoingURIs(r, this._uriTransformer);
-			}
-			const msg = MessageIO.serializeReplyOK(req, r);
+			const msg = MessageIO.serializeReplyOK(req, r, this._uriReplacer);
 			if (this._logger) {
 				this._logger.logOutgoing(msg.byteLength, req, RequestInitiator.OtherSide, `reply:`, r);
 			}
@@ -469,10 +492,7 @@ export class RPCProtocol extends Disposable implements IRPCProtocol {
 
 		this._pendingRPCReplies[callId] = result;
 		this._onWillSendRequest(req);
-		if (this._uriTransformer) {
-			args = transformOutgoingURIs(args, this._uriTransformer);
-		}
-		const msg = MessageIO.serializeRequest(req, rpcId, methodName, args, !!cancellationToken);
+		const msg = MessageIO.serializeRequest(req, rpcId, methodName, args, !!cancellationToken, this._uriReplacer);
 		if (this._logger) {
 			this._logger.logOutgoing(msg.byteLength, req, RequestInitiator.LocalSide, `request: ${getStringIdentifierForProxy(rpcId)}.${methodName}(`, args);
 		}
@@ -630,7 +650,7 @@ class MessageIO {
 		return false;
 	}
 
-	public static serializeRequest(req: number, rpcId: number, method: string, args: any[], usesCancellationToken: boolean): Buffer {
+	public static serializeRequest(req: number, rpcId: number, method: string, args: any[], usesCancellationToken: boolean, replacer: JSONStringifyReplacer | null): Buffer {
 		if (this._arrayContainsBuffer(args)) {
 			let massagedArgs: (string | Buffer)[] = new Array(args.length);
 			let argsLengths: number[] = new Array(args.length);
@@ -640,13 +660,13 @@ class MessageIO {
 					massagedArgs[i] = arg;
 					argsLengths[i] = arg.byteLength;
 				} else {
-					massagedArgs[i] = safeStringify(arg);
+					massagedArgs[i] = safeStringify(arg, replacer);
 					argsLengths[i] = Buffer.byteLength(massagedArgs[i], 'utf8');
 				}
 			}
 			return this._requestMixedArgs(req, rpcId, method, massagedArgs, argsLengths, usesCancellationToken);
 		}
-		return this._requestJSONArgs(req, rpcId, method, safeStringify(args), usesCancellationToken);
+		return this._requestJSONArgs(req, rpcId, method, safeStringify(args, replacer), usesCancellationToken);
 	}
 
 	private static _requestJSONArgs(req: number, rpcId: number, method: string, args: string, usesCancellationToken: boolean): Buffer {
@@ -719,14 +739,14 @@ class MessageIO {
 		return MessageBuffer.alloc(MessageType.Cancel, req, 0).buffer;
 	}
 
-	public static serializeReplyOK(req: number, res: any): Buffer {
+	public static serializeReplyOK(req: number, res: any, replacer: JSONStringifyReplacer | null): Buffer {
 		if (typeof res === 'undefined') {
 			return this._serializeReplyOKEmpty(req);
 		}
 		if (Buffer.isBuffer(res)) {
 			return this._serializeReplyOKBuffer(req, res);
 		}
-		return this._serializeReplyOKJSON(req, safeStringify(res));
+		return this._serializeReplyOKJSON(req, safeStringify(res, replacer));
 	}
 
 	private static _serializeReplyOKEmpty(req: number): Buffer {
@@ -772,7 +792,7 @@ class MessageIO {
 	}
 
 	private static _serializeReplyErrEror(req: number, _err: Error): Buffer {
-		const err = safeStringify(errors.transformErrorForSerialization(_err));
+		const err = safeStringify(errors.transformErrorForSerialization(_err), null);
 		const errByteLength = Buffer.byteLength(err, 'utf8');
 
 		let len = 0;
@@ -793,14 +813,6 @@ class MessageIO {
 	}
 }
 
-function safeStringify(obj: any): string {
-	try {
-		return JSON.stringify(obj);
-	} catch (err) {
-		return 'null';
-	}
-}
-
 const enum MessageType {
 	RequestJSONArgs = 1,
 	RequestJSONArgsWithCancellation = 2,