Add support for client websockets

src/api/loginclient.ts

@@ -108,15 +108,40 @@ export class LoginClient {
 		} else if (this.#step.type === "client_http") {
 			const headers = Object.fromEntries(Object.entries(this.#step.client_http.headers ?? {})
 				.map(([key, values]) => [key, values[0]]))
-			const body = this.#step.client_http.body ? Uint8Array.fromBase64(this.#step.client_http.body) : undefined
-			window.mautrixAPI.doClientHTTP(this.#step.client_http.url, {
-				method: this.#step.client_http.method,
-				headers,
-				body,
-			}).then(
-				resp => this.submitClientHTTP(resp),
-				err => this.submitClientHTTP({ error: err.toString() }),
-			)
+			if (
+				this.#step.client_http.method === "GET"
+				&& headers["Connection"] === "Upgrade"
+				&& headers["Upgrade"] === "websocket"
+				&& headers["Sec-Websocket-Key"]
+				&& headers["Sec-Websocket-Version"]
+				&& this.#step.client_http.stream_id
+			) {
+				const streamID = this.#step.client_http.stream_id
+				window.mautrixAPI.doClientWebsocket(
+					this.client.getWebsocketTunnelURL(streamID),
+					this.#step.client_http.url,
+					headers,
+				).then(
+					resp => {
+						if ("status_code" in resp) {
+							resp.stream_id = streamID
+						}
+						this.submitClientHTTP(resp)
+					},
+					err => this.submitClientHTTP({ error: err.toString() }),
+				)
+			} else {
+				const body = this.#step.client_http.body
+					? Uint8Array.fromBase64(this.#step.client_http.body) : undefined
+				window.mautrixAPI.doClientHTTP(this.#step.client_http.url, {
+					method: this.#step.client_http.method,
+					headers,
+					body,
+				}).then(
+					resp => this.submitClientHTTP(resp),
+					err => this.submitClientHTTP({ error: err.toString() }),
+				)
+			}
 		} else if (this.#step.type === "complete") {
 			setTimeout(this.onCompleteRefresh, 200)
 		}

src/api/provisionclient.ts

@@ -32,6 +32,16 @@ export class ProvisioningClient extends BaseAPIClient {
 		return new ProvisioningClient(this.baseURL, this.matrixClient, this.external, loginID)
 	}
 
+	getWebsocketTunnelURL(streamID: string): string {
+		const url = new URL(this.baseURL)
+		url.protocol = url.protocol === "http:" ? "ws" : "wss:"
+		url.pathname = `${url.pathname}${this.pathPrefix}/v3/login/client_websocket/${this.loginID}/${streamID}`
+		if (this.userID) {
+			url.searchParams.set("user_id", this.userID)
+		}
+		return url.toString()
+	}
+
 	whoami(signal?: AbortSignal): Promise<RespWhoami> {
 		return this.request("GET", "/v3/whoami", { signal })
 	}
@@ -43,10 +53,10 @@ export class ProvisioningClient extends BaseAPIClient {
 	async startLogin(flowID: string, refresh: () => void): Promise<LoginClient> {
 		const resp: RespSubmitLogin = await this.request(
 			"POST", `/v3/login/start/${encodeURIComponent(flowID)}`, {}, {
-				query: { client_http: "1" },
+				query: { client_http: "1", client_websocket: "2" },
 			},
 		)
-		return new LoginClient(this, resp, refresh)
+		return new LoginClient(this.withLogin(resp.login_id), resp, refresh)
 	}
 
 	logout(loginID: string | "all"): Promise<RespLogout> {

src/electron.ts

@@ -1,5 +1,7 @@
+import https from "https"
+import { Socket } from "net"
 import path from "path"
-import { BrowserWindow, app, ipcMain, shell } from "electron"
+import { BrowserWindow, app, ipcMain, net, shell } from "electron"
 import started from "electron-squirrel-startup"
 import type { AccessTokenChangedParams } from "./preload"
 import type { LoginClientHTTPResponse } from "./types/loginstep.ts"
@@ -27,6 +29,81 @@ ipcMain.handle("mautrix:open-in-browser", (_event, url: string) => {
 	return shell.openExternal(url)
 })
 
+ipcMain.handle("mautrix:client-websocket", (
+	_event, tunnelURL: string, url: string, headers: Record<string, string>,
+): Promise<LoginClientHTTPResponse> => {
+	const makeTunnel = (socket: Socket, head: Buffer | null) => {
+		socket.pause()
+		console.log("Opening client websocket tunnel to", tunnelURL)
+		const tunnel = new net.WebSocket(tunnelURL, {
+			headers: {
+				"Authorization": `Bearer ${accessToken}`,
+			},
+		})
+		tunnel.binaryType = "nodebuffer"
+		socket.on("data", data => {
+			console.log("Client websocket upstream -> tunnel message", data.length)
+			tunnel.send(data)
+		})
+		tunnel.onmessage = (msg: MessageEvent) => {
+			console.log("Client websocket tunnel -> upstream message", msg.data.length)
+			socket.write(msg.data)
+		}
+		tunnel.onopen = () => {
+			console.log("Client websocket tunnel opened")
+			if (head) {
+				console.log("Sending", head.length, "initial bytes")
+				tunnel.send(head)
+				head = null
+			}
+			socket.resume()
+		}
+		tunnel.onerror = (err: Error) => {
+			console.log("Client websocket tunnel errored", err)
+			socket.destroy()
+		}
+		tunnel.onclose = () => {
+			console.log("Client websocket tunnel closed")
+			if (!socket.closed) {
+				socket.end()
+				socket.setTimeout(2000, () => socket.destroy())
+			}
+		}
+		const closeTunnel = (err?: Error) => {
+			console.log("Client websocket upstream closed", err)
+			if (tunnel.readyState !== WebSocket.CLOSED && tunnel.readyState !== WebSocket.CLOSING) {
+				tunnel.close()
+			}
+		}
+		socket.on("close", closeTunnel)
+		socket.on("error", closeTunnel)
+	}
+	return new Promise<LoginClientHTTPResponse>((resolve, reject) => {
+		console.log("Opening client websocket to", url, "with headers", headers)
+		const wsReq = https.request(url, { headers })
+		wsReq.on("upgrade", (res, socket, head) => {
+			console.log("Client websocket upstream opened successfully", res.statusCode, headers)
+			resolve({
+				status_code: res.statusCode ?? 101,
+				headers: Object.fromEntries(Object.entries(res.headers).map(([name, value]) => {
+					if (value == undefined) {
+						value = [""]
+					} else if (typeof value === "string") {
+						value = [value]
+					}
+					return [name, value]
+				})),
+			})
+			makeTunnel(socket, head)
+		})
+		wsReq.on("error", err => {
+			console.log("Client websocket upstream errored", err)
+			reject(err)
+		})
+		wsReq.end()
+	})
+})
+
 ipcMain.handle("mautrix:client-http", async (
 	_event, url: string, init: Parameters<typeof fetch>[1],
 ): Promise<LoginClientHTTPResponse> => {

src/preload.ts

@@ -16,6 +16,9 @@ interface MautrixAPI {
 	closeWebview: () => Promise<void>
 	doWebAuthn: (params: LoginWebAuthnParams) => Promise<AuthenticationResponseJSON>
 	doClientHTTP: (...params: Parameters<typeof fetch>) => Promise<LoginClientHTTPResponse>,
+	doClientWebsocket: (
+		tunnelURL: string, url: string, headers: Record<string, string>,
+	) => Promise<LoginClientHTTPResponse>,
 	accessTokenChanged: (newDetails: AccessTokenChangedParams) => Promise<void>
 	openInBrowser: (url: string) => Promise<void>
 	isDevBuild: boolean,
@@ -26,6 +29,8 @@ const api: MautrixAPI = {
 	closeWebview: () => ipcRenderer.invoke("mautrix:close-webview"),
 	doWebAuthn: params => ipcRenderer.invoke("mautrix:webauthn", params),
 	doClientHTTP: (url, init) => ipcRenderer.invoke("mautrix:client-http", url, init),
+	doClientWebsocket: (tunnelURL, url, headers) =>
+		ipcRenderer.invoke("mautrix:client-websocket", tunnelURL, url, headers),
 	accessTokenChanged: newDetails => ipcRenderer.invoke("mautrix:access-token-changed", newDetails),
 	openInBrowser: url => ipcRenderer.invoke("mautrix:open-in-browser", url),
 	isDevBuild: process.env.NODE_ENV === "development",

src/types/loginstep.ts

@@ -85,13 +85,15 @@ export interface LoginClientHTTPParams {
 	url: string
 	headers?: Record<string, string[]>
 	body?: string
+	stream_id?: string
 }
 
 export type LoginClientHTTPResponse = {
 	status_code: number
 	final_url?: string
 	headers?: Record<string, string[]>
 	body?: string
+	stream_id?: string
 } | {
 	error: string
 }