|
| 1 | +import { logger } from "@coder/logger" |
| 2 | +import * as os from "os" |
| 3 | +import * as path from "path" |
| 4 | +import { promises as fs } from "fs" |
| 5 | +import { isNodeJSErrnoException } from "./util" |
| 6 | +import { canConnect } from "./util" |
| 7 | + |
| 8 | +export const DEFAULT_SOCKET_PATH = path.join(os.tmpdir(), "vscode-ipc") |
| 9 | + |
| 10 | +export interface VscodeSocket { |
| 11 | + workspace: { |
| 12 | + id: string |
| 13 | + folders: { |
| 14 | + uri: { |
| 15 | + path: string |
| 16 | + } |
| 17 | + }[] |
| 18 | + } |
| 19 | + |
| 20 | + socketPath: string |
| 21 | +} |
| 22 | + |
| 23 | +export class VscodeSocketResolver { |
| 24 | + // Map from folder path to socket path. |
| 25 | + private pathMap = new Map<string, string[]>() |
| 26 | + |
| 27 | + // Map from socket path to VscodeSocket. |
| 28 | + private entries = new Map<string, { index: number; entry: VscodeSocket }>() |
| 29 | + |
| 30 | + constructor(private readonly socketRegistryFilePath: string) {} |
| 31 | + |
| 32 | + async load(): Promise<void> { |
| 33 | + try { |
| 34 | + const data = await fs.readFile(this.socketRegistryFilePath, "utf8") |
| 35 | + const lines = data.split("\n") |
| 36 | + let entries = [] |
| 37 | + for (const line of lines) { |
| 38 | + if (!line) { |
| 39 | + continue |
| 40 | + } |
| 41 | + entries.push(JSON.parse(line) as VscodeSocket) |
| 42 | + } |
| 43 | + this.loadFromEntries(entries) |
| 44 | + } catch (error) { |
| 45 | + // If it doesn't exist, we don't care. |
| 46 | + // But if it fails for some reason, we should throw. |
| 47 | + // We want to surface that to the user. |
| 48 | + if (!isNodeJSErrnoException(error) || error.code !== "ENOENT") { |
| 49 | + throw error |
| 50 | + } |
| 51 | + } |
| 52 | + } |
| 53 | + |
| 54 | + private loadFromEntries(entries: VscodeSocket[]): void { |
| 55 | + for (let i = 0; i < entries.length; i++) { |
| 56 | + const entry = entries[i] |
| 57 | + this.entries.set(entry.socketPath, { index: i, entry }) |
| 58 | + for (const folder of entry.workspace.folders) { |
| 59 | + if (this.pathMap.has(folder.uri.path)) { |
| 60 | + this.pathMap.get(folder.uri.path)!.push(entry.socketPath) |
| 61 | + } else { |
| 62 | + this.pathMap.set(folder.uri.path, [entry.socketPath]) |
| 63 | + } |
| 64 | + } |
| 65 | + } |
| 66 | + } |
| 67 | + |
| 68 | + /** |
| 69 | + * Returns matching and then non-matching socket paths, sorted by most recently registered. |
| 70 | + */ |
| 71 | + getSocketPaths(paths: string[] = []): string[] { |
| 72 | + const matches = new Set<string>() |
| 73 | + for (const socketPath of paths.flatMap((p) => this.getMatchingSocketPaths(p))) { |
| 74 | + matches.add(socketPath) |
| 75 | + } |
| 76 | + |
| 77 | + // Sort by most recently registered, preferring matches. |
| 78 | + return Array.from(this.entries.values()) |
| 79 | + .sort((a, b) => { |
| 80 | + if (matches.has(a.entry.socketPath) && matches.has(b.entry.socketPath)) { |
| 81 | + return b.index - a.index |
| 82 | + } |
| 83 | + if (matches.has(a.entry.socketPath)) { |
| 84 | + return -1 |
| 85 | + } |
| 86 | + if (matches.has(b.entry.socketPath)) { |
| 87 | + return 1 |
| 88 | + } |
| 89 | + return b.index - a.index |
| 90 | + }) |
| 91 | + .map(({ entry }) => entry.socketPath) |
| 92 | + } |
| 93 | + |
| 94 | + /** |
| 95 | + * Returns the best socket path that we can connect to. |
| 96 | + * See getSocketPaths for the order of preference. |
| 97 | + * We also delete any sockets that we can't connect to. |
| 98 | + */ |
| 99 | + async getConnectedSocketPath(paths: string[] = []): Promise<string | undefined> { |
| 100 | + const socketPaths = this.getSocketPaths(paths) |
| 101 | + let ret = undefined |
| 102 | + |
| 103 | + const toDelete = new Set<string>() |
| 104 | + // Connect to the most recently registered socket first. |
| 105 | + for (let i = 0; i < socketPaths.length; i++) { |
| 106 | + const socketPath = socketPaths[i] |
| 107 | + if (await canConnect(socketPath)) { |
| 108 | + ret = socketPath |
| 109 | + break |
| 110 | + } |
| 111 | + toDelete.add(socketPath) |
| 112 | + } |
| 113 | + |
| 114 | + if (toDelete.size === 0) { |
| 115 | + return ret |
| 116 | + } |
| 117 | + |
| 118 | + // Remove any sockets that we couldn't connect to. |
| 119 | + for (const socketPath of toDelete) { |
| 120 | + logger.debug(`Deleting stale socket from socket registry: ${socketPath}`) |
| 121 | + this.entries.delete(socketPath) |
| 122 | + } |
| 123 | + const newEntries = Array.from(this.entries.values()).map(({ entry }) => entry) |
| 124 | + this.pathMap = new Map() |
| 125 | + this.entries = new Map() |
| 126 | + this.loadFromEntries(newEntries) |
| 127 | + await this.save() |
| 128 | + |
| 129 | + return ret |
| 130 | + } |
| 131 | + |
| 132 | + private getMatchingSocketPaths(p: string): string[] { |
| 133 | + return Array.from(this.pathMap.entries()) |
| 134 | + .filter(([folder]) => p.startsWith(folder)) |
| 135 | + .flatMap(([, socketPaths]) => socketPaths) |
| 136 | + } |
| 137 | + |
| 138 | + async save(): Promise<void> { |
| 139 | + const lines = Array.from(this.entries.values()).map(({ entry }) => JSON.stringify(entry)) |
| 140 | + return fs.writeFile(this.socketRegistryFilePath, lines.join("\n")) |
| 141 | + } |
| 142 | +} |
0 commit comments