Resolve vscode ipc sockets based on folder match.

Author: sleexyz

patches/store-socket.diff

@@ -24,14 +24,37 @@ Index: code-server/lib/vscode/src/vs/workbench/api/node/extHostExtensionService.
  import * as performance from 'vs/base/common/performance';
  import { createApiFactoryAndRegisterActors } from 'vs/workbench/api/common/extHost.api.impl';
  import { RequireInterceptor } from 'vs/workbench/api/common/extHostRequireInterceptor';
-@@ -72,6 +74,10 @@ export class ExtHostExtensionService ext
- 		if (this._initData.remote.isRemote && this._initData.remote.authority) {
- 			const cliServer = this._instaService.createInstance(CLIServer);
- 			process.env['VSCODE_IPC_HOOK_CLI'] = cliServer.ipcHandlePath;
-+
-+			fs.writeFile(path.join(os.tmpdir(), 'vscode-ipc'), cliServer.ipcHandlePath).catch((error) => {
-+				this._logService.error(error);
+@@ -17,6 +19,7 @@ import { ExtensionRuntime } from 'vs/wor
+ import { CLIServer } from 'vs/workbench/api/node/extHostCLIServer';
+ import { realpathSync } from 'vs/base/node/extpath';
+ import { ExtHostConsoleForwarder } from 'vs/workbench/api/node/extHostConsoleForwarder';
++import { IExtHostWorkspace } from '../common/extHostWorkspace';
+ 
+ class NodeModuleRequireInterceptor extends RequireInterceptor {
+ 
+@@ -83,6 +86,25 @@ export class ExtHostExtensionService ext
+ 		const configProvider = await this._extHostConfiguration.getConfigProvider();
+ 		await connectProxyResolver(this._extHostWorkspace, configProvider, this, this._logService, this._mainThreadTelemetryProxy, this._initData);
+ 		performance.mark('code/extHost/didInitProxyResolver');
++	
++		// TODO: confirm that this is the best place to put this.
++		(async () => {
++			let handlePath = process.env['VSCODE_IPC_HOOK_CLI'];
++			if (!handlePath) {
++				return
++			}
++			const workspace = this._instaService.invokeFunction((accessor) => {
++				const workspaceService = accessor.get(IExtHostWorkspace);
++				return workspaceService.workspace;
 +			});
- 		}
++			const entry = {
++				workspace,
++				ipcHandlePath: handlePath
++			};
++			fs.appendFile(path.join(os.tmpdir(), 'vscode-ipc'), '\n' + JSON.stringify(entry), 'utf-8');
++		})().catch(error => {
++			this._logService.error(error);
++		});
+ 	}
 
- 		// Module loading tricks
+ 	protected _getEntryPoint(extensionDescription: IExtensionDescription): string | undefined {

src/node/cli.ts

@@ -3,17 +3,9 @@ import { promises as fs } from "fs"
 import { load } from "js-yaml"
 import * as os from "os"
 import * as path from "path"
-import {
-  canConnect,
-  generateCertificate,
-  generatePassword,
-  humanPath,
-  paths,
-  isNodeJSErrnoException,
-  splitOnFirstEquals,
-} from "./util"
-
-const DEFAULT_SOCKET_PATH = path.join(os.tmpdir(), "vscode-ipc")
+import { generateCertificate, generatePassword, humanPath, paths, splitOnFirstEquals } from "./util"
+import { DEFAULT_SOCKET_PATH, VscodeSocketResolver } from "./vscodeSocket"
+import { getResolvedPathsFromArgs } from "./main"
 
 export enum Feature {
   // No current experimental features!
@@ -732,27 +724,6 @@ function bindAddrFromAllSources(...argsConfig: UserProvidedArgs[]): Addr {
   return addr
 }
 
-/**
- * Reads the socketPath based on path passed in.
- *
- * The one usually passed in is the DEFAULT_SOCKET_PATH.
- *
- * If it can't read the path, it throws an error and returns undefined.
- */
-export async function readSocketPath(path: string): Promise<string | undefined> {
-  try {
-    return await fs.readFile(path, "utf8")
-  } catch (error) {
-    // If it doesn't exist, we don't care.
-    // But if it fails for some reason, we should throw.
-    // We want to surface that to the user.
-    if (!isNodeJSErrnoException(error) || error.code !== "ENOENT") {
-      throw error
-    }
-  }
-  return undefined
-}
-
 /**
  * Determine if it looks like the user is trying to open a file or folder in an
  * existing instance. The arguments here should be the arguments the user
@@ -765,14 +736,18 @@ export const shouldOpenInExistingInstance = async (args: UserProvidedArgs): Prom
     return process.env.VSCODE_IPC_HOOK_CLI
   }
 
+  const paths = getResolvedPathsFromArgs(args)
+  const resolver = new VscodeSocketResolver(DEFAULT_SOCKET_PATH)
+
   // If these flags are set then assume the user is trying to open in an
   // existing instance since these flags have no effect otherwise.
   const openInFlagCount = ["reuse-window", "new-window"].reduce((prev, cur) => {
     return args[cur as keyof UserProvidedArgs] ? prev + 1 : prev
   }, 0)
   if (openInFlagCount > 0) {
     logger.debug("Found --reuse-window or --new-window")
-    return readSocketPath(DEFAULT_SOCKET_PATH)
+    await resolver.load()
+    return await resolver.getConnectedSocketPath(paths)
   }
 
   // It's possible the user is trying to spawn another instance of code-server.
@@ -781,8 +756,9 @@ export const shouldOpenInExistingInstance = async (args: UserProvidedArgs): Prom
   // 2. That a file or directory was passed.
   // 3. That the socket is active.
   if (Object.keys(args).length === 1 && typeof args._ !== "undefined" && args._.length > 0) {
-    const socketPath = await readSocketPath(DEFAULT_SOCKET_PATH)
-    if (socketPath && (await canConnect(socketPath))) {
+    await resolver.load()
+    const socketPath = await resolver.getConnectedSocketPath(paths)
+    if (socketPath) {
       logger.debug("Found existing code-server socket")
       return socketPath
     }

src/node/main.ts

@@ -61,6 +61,11 @@ export const runCodeCli = async (args: DefaultedArgs): Promise<void> => {
   process.exit(0)
 }
 
+export const getResolvedPathsFromArgs = (args: UserProvidedArgs): string[] => {
+  const paths = args._ || []
+  return paths.map((p) => path.resolve(p))
+}
+
 export const openInExistingInstance = async (args: DefaultedArgs, socketPath: string): Promise<void> => {
   const pipeArgs: OpenCommandPipeArgs & { fileURIs: string[] } = {
     type: "open",
@@ -70,9 +75,9 @@ export const openInExistingInstance = async (args: DefaultedArgs, socketPath: st
     forceNewWindow: args["new-window"],
     gotoLineMode: true,
   }
-  const paths = args._ || []
+  const paths = getResolvedPathsFromArgs(args)
   for (let i = 0; i < paths.length; i++) {
-    const fp = path.resolve(paths[i])
+    const fp = paths[i]
     if (await isDirectory(fp)) {
       pipeArgs.folderURIs.push(fp)
     } else {

src/node/vscodeSocket.ts

@@ -0,0 +1,142 @@
+import { logger } from "@coder/logger"
+import * as os from "os"
+import * as path from "path"
+import { promises as fs } from "fs"
+import { isNodeJSErrnoException } from "./util"
+import { canConnect } from "./util"
+
+export const DEFAULT_SOCKET_PATH = path.join(os.tmpdir(), "vscode-ipc")
+
+export interface VscodeSocket {
+  workspace: {
+    id: string
+    folders: {
+      uri: {
+        path: string
+      }
+    }[]
+  }
+
+  socketPath: string
+}
+
+export class VscodeSocketResolver {
+  // Map from folder path to socket path.
+  private pathMap = new Map<string, string[]>()
+
+  // Map from socket path to VscodeSocket.
+  private entries = new Map<string, { index: number; entry: VscodeSocket }>()
+
+  constructor(private readonly socketRegistryFilePath: string) {}
+
+  async load(): Promise<void> {
+    try {
+      const data = await fs.readFile(this.socketRegistryFilePath, "utf8")
+      const lines = data.split("\n")
+      let entries = []
+      for (const line of lines) {
+        if (!line) {
+          continue
+        }
+        entries.push(JSON.parse(line) as VscodeSocket)
+      }
+      this.loadFromEntries(entries)
+    } catch (error) {
+      // If it doesn't exist, we don't care.
+      // But if it fails for some reason, we should throw.
+      // We want to surface that to the user.
+      if (!isNodeJSErrnoException(error) || error.code !== "ENOENT") {
+        throw error
+      }
+    }
+  }
+
+  private loadFromEntries(entries: VscodeSocket[]): void {
+    for (let i = 0; i < entries.length; i++) {
+      const entry = entries[i]
+      this.entries.set(entry.socketPath, { index: i, entry })
+      for (const folder of entry.workspace.folders) {
+        if (this.pathMap.has(folder.uri.path)) {
+          this.pathMap.get(folder.uri.path)!.push(entry.socketPath)
+        } else {
+          this.pathMap.set(folder.uri.path, [entry.socketPath])
+        }
+      }
+    }
+  }
+
+  /**
+   * Returns matching and then non-matching socket paths, sorted by most recently registered.
+   */
+  getSocketPaths(paths: string[] = []): string[] {
+    const matches = new Set<string>()
+    for (const socketPath of paths.flatMap((p) => this.getMatchingSocketPaths(p))) {
+      matches.add(socketPath)
+    }
+
+    // Sort by most recently registered, preferring matches.
+    return Array.from(this.entries.values())
+      .sort((a, b) => {
+        if (matches.has(a.entry.socketPath) && matches.has(b.entry.socketPath)) {
+          return b.index - a.index
+        }
+        if (matches.has(a.entry.socketPath)) {
+          return -1
+        }
+        if (matches.has(b.entry.socketPath)) {
+          return 1
+        }
+        return b.index - a.index
+      })
+      .map(({ entry }) => entry.socketPath)
+  }
+
+  /**
+   * Returns the best socket path that we can connect to.
+   * See getSocketPaths for the order of preference.
+   * We also delete any sockets that we can't connect to.
+   */
+  async getConnectedSocketPath(paths: string[] = []): Promise<string | undefined> {
+    const socketPaths = this.getSocketPaths(paths)
+    let ret = undefined
+
+    const toDelete = new Set<string>()
+    // Connect to the most recently registered socket first.
+    for (let i = 0; i < socketPaths.length; i++) {
+      const socketPath = socketPaths[i]
+      if (await canConnect(socketPath)) {
+        ret = socketPath
+        break
+      }
+      toDelete.add(socketPath)
+    }
+
+    if (toDelete.size === 0) {
+      return ret
+    }
+
+    // Remove any sockets that we couldn't connect to.
+    for (const socketPath of toDelete) {
+      logger.debug(`Deleting stale socket from socket registry: ${socketPath}`)
+      this.entries.delete(socketPath)
+    }
+    const newEntries = Array.from(this.entries.values()).map(({ entry }) => entry)
+    this.pathMap = new Map()
+    this.entries = new Map()
+    this.loadFromEntries(newEntries)
+    await this.save()
+
+    return ret
+  }
+
+  private getMatchingSocketPaths(p: string): string[] {
+    return Array.from(this.pathMap.entries())
+      .filter(([folder]) => p.startsWith(folder))
+      .flatMap(([, socketPaths]) => socketPaths)
+  }
+
+  async save(): Promise<void> {
+    const lines = Array.from(this.entries.values()).map(({ entry }) => JSON.stringify(entry))
+    return fs.writeFile(this.socketRegistryFilePath, lines.join("\n"))
+  }
+}