From 58c1be57fa92cbf7ccd002cf04d5cc62b182c481 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Thu, 10 Dec 2020 15:59:24 -0600
Subject: [PATCH] Implement callback endpoints

VS Code uses these during the authentication flow.
---
 src/common/util.ts        |  11 ++++
 src/node/routes/static.ts |   5 +-
 src/node/routes/vscode.ts | 107 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 120 insertions(+), 3 deletions(-)

diff --git a/src/common/util.ts b/src/common/util.ts
index 7baa355a..67e182ce 100644
--- a/src/common/util.ts
+++ b/src/common/util.ts
@@ -101,3 +101,14 @@ export const arrayify = <T>(value?: T | T[]): T[] => {
   }
   return [value]
 }
+
+/**
+ * Get the first string. If there's no string return undefined.
+ */
+export const getFirstString = (value: string | string[] | object | undefined): string | undefined => {
+  if (Array.isArray(value)) {
+    return value[0]
+  }
+
+  return typeof value !== "object" ? value : undefined
+}
diff --git a/src/node/routes/static.ts b/src/node/routes/static.ts
index 22bdd8d2..30eed031 100644
--- a/src/node/routes/static.ts
+++ b/src/node/routes/static.ts
@@ -6,6 +6,7 @@ import { Readable } from "stream"
 import * as tarFs from "tar-fs"
 import * as zlib from "zlib"
 import { HttpCode, HttpError } from "../../common/http"
+import { getFirstString } from "../../common/util"
 import { rootPath } from "../constants"
 import { authenticated, ensureAuthenticated, replaceTemplates } from "../http"
 import { getMediaMime, pathToFsPath } from "../util"
@@ -15,8 +16,8 @@ export const router = Router()
 // The commit is for caching.
 router.get("/(:commit)(/*)?", async (req, res) => {
   // Used by VS Code to load extensions into the web worker.
-  const tar = Array.isArray(req.query.tar) ? req.query.tar[0] : req.query.tar
-  if (typeof tar === "string") {
+  const tar = getFirstString(req.query.tar)
+  if (tar) {
     ensureAuthenticated(req)
     let stream: Readable = tarFs.pack(pathToFsPath(tar))
     if (req.headers["accept-encoding"] && req.headers["accept-encoding"].includes("gzip")) {
diff --git a/src/node/routes/vscode.ts b/src/node/routes/vscode.ts
index 373dd4ce..85d902d3 100644
--- a/src/node/routes/vscode.ts
+++ b/src/node/routes/vscode.ts
@@ -1,7 +1,11 @@
 import * as crypto from "crypto"
-import { Router } from "express"
+import { Request, Router } from "express"
 import { promises as fs } from "fs"
 import * as path from "path"
+import qs from "qs"
+import { Emitter } from "../../common/emitter"
+import { HttpCode, HttpError } from "../../common/http"
+import { getFirstString } from "../../common/util"
 import { commit, rootPath, version } from "../constants"
 import { authenticated, ensureAuthenticated, redirect, replaceTemplates } from "../http"
 import { getMediaMime, pathToFsPath } from "../util"
@@ -86,6 +90,107 @@ router.get("/webview/*", ensureAuthenticated, async (req, res) => {
   )
 })
 
+interface Callback {
+  uri: {
+    scheme: string
+    authority?: string
+    path?: string
+    query?: string
+    fragment?: string
+  }
+  timeout: NodeJS.Timeout
+}
+
+const callbacks = new Map<string, Callback>()
+const callbackEmitter = new Emitter<{ id: string; callback: Callback }>()
+
+/**
+ * Get vscode-requestId from the query and throw if it's missing or invalid.
+ */
+const getRequestId = (req: Request): string => {
+  if (!req.query["vscode-requestId"]) {
+    throw new HttpError("vscode-requestId is missing", HttpCode.BadRequest)
+  }
+
+  if (typeof req.query["vscode-requestId"] !== "string") {
+    throw new HttpError("vscode-requestId is not a string", HttpCode.BadRequest)
+  }
+
+  return req.query["vscode-requestId"]
+}
+
+// Matches VS Code's fetch timeout.
+const fetchTimeout = 5 * 60 * 1000
+
+// The callback endpoints are used during authentication. A URI is stored on
+// /callback and then fetched later on /fetch-callback.
+// See ../../../lib/vscode/resources/web/code-web.js
+router.get("/callback", ensureAuthenticated, async (req, res) => {
+  const uriKeys = [
+    "vscode-requestId",
+    "vscode-scheme",
+    "vscode-authority",
+    "vscode-path",
+    "vscode-query",
+    "vscode-fragment",
+  ]
+
+  const id = getRequestId(req)
+
+  // Move any query variables that aren't URI keys into the URI's query
+  // (importantly, this will include the code for oauth).
+  const query: qs.ParsedQs = {}
+  for (const key in req.query) {
+    if (!uriKeys.includes(key)) {
+      query[key] = req.query[key]
+    }
+  }
+
+  const callback = {
+    uri: {
+      scheme: getFirstString(req.query["vscode-scheme"]) || "code-oss",
+      authority: getFirstString(req.query["vscode-authority"]),
+      path: getFirstString(req.query["vscode-path"]),
+      query: (getFirstString(req.query.query) ? getFirstString(req.query.query) + "&" : "") + qs.stringify(query),
+      fragment: getFirstString(req.query["vscode-fragment"]),
+    },
+    // Make sure the map doesn't leak if nothing fetches this URI.
+    timeout: setTimeout(() => callbacks.delete(id), fetchTimeout),
+  }
+
+  callbacks.set(id, callback)
+  callbackEmitter.emit({ id, callback })
+
+  res.sendFile(path.join(rootPath, "lib/vscode/resources/web/callback.html"))
+})
+
+router.get("/fetch-callback", ensureAuthenticated, async (req, res) => {
+  const id = getRequestId(req)
+
+  const send = (callback: Callback) => {
+    clearTimeout(callback.timeout)
+    callbacks.delete(id)
+    res.json(callback.uri)
+  }
+
+  const callback = callbacks.get(id)
+  if (callback) {
+    return send(callback)
+  }
+
+  // VS Code will try again if the route returns no content but it seems more
+  // efficient to just wait on this request for as long as possible?
+  const handler = callbackEmitter.event(({ id: emitId, callback }) => {
+    if (id === emitId) {
+      handler.dispose()
+      send(callback)
+    }
+  })
+
+  // If the client closes the connection.
+  req.on("close", () => handler.dispose())
+})
+
 export const wsRouter = WsRouter()
 
 wsRouter.ws("/", ensureAuthenticated, async (req) => {