3 files changed, 108 insertions, 19 deletions

diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index bd99d696a..2896d90ac 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -95,6 +95,10 @@ async function tryActivate(context: vscode.ExtensionContext) {
         await activate(context).catch(log.error);
     });
 
+    ctx.registerCommand('updateGithubToken', ctx => async () => {
+        await queryForGithubToken(new PersistentState(ctx.globalState));
+    });
+
     ctx.registerCommand('analyzerStatus', commands.analyzerStatus);
     ctx.registerCommand('memoryUsage', commands.memoryUsage);
     ctx.registerCommand('reloadWorkspace', commands.reloadWorkspace);
@@ -173,7 +177,9 @@ async function bootstrapExtension(config: Config, state: PersistentState): Promi
         if (!shouldCheckForNewNightly) return;
     }
 
-    const release = await fetchRelease("nightly").catch((e) => {
+    const release = await downloadWithRetryDialog(state, async () => {
+        return await fetchRelease("nightly", state.githubToken);
+    }).catch((e) => {
         log.error(e);
         if (state.releaseId === undefined) { // Show error only for the initial download
             vscode.window.showErrorMessage(`Failed to download rust-analyzer nightly ${e}`);
@@ -192,10 +198,14 @@ async function bootstrapExtension(config: Config, state: PersistentState): Promi
     assert(!!artifact, `Bad release: ${JSON.stringify(release)}`);
 
     const dest = path.join(config.globalStoragePath, "rust-analyzer.vsix");
-    await download({
-        url: artifact.browser_download_url,
-        dest,
-        progressTitle: "Downloading rust-analyzer extension",
+
+    await downloadWithRetryDialog(state, async () => {
+        await download({
+            url: artifact.browser_download_url,
+            dest,
+            progressTitle: "Downloading rust-analyzer extension",
+            overwrite: true,
+        });
     });
 
     await vscode.commands.executeCommand("workbench.extensions.installExtension", vscode.Uri.file(dest));
@@ -308,21 +318,22 @@ async function getServer(config: Config, state: PersistentState): Promise<string
         if (userResponse !== "Download now") return dest;
     }
 
-    const release = await fetchRelease(config.package.releaseTag);
+    const releaseTag = config.package.releaseTag;
+    const release = await downloadWithRetryDialog(state, async () => {
+        return await fetchRelease(releaseTag, state.githubToken);
+    });
     const artifact = release.assets.find(artifact => artifact.name === `rust-analyzer-${platform}.gz`);
     assert(!!artifact, `Bad release: ${JSON.stringify(release)}`);
 
-    // Unlinking the exe file before moving new one on its place should prevent ETXTBSY error.
-    await fs.unlink(dest).catch(err => {
-        if (err.code !== "ENOENT") throw err;
-    });
-
-    await download({
-        url: artifact.browser_download_url,
-        dest,
-        progressTitle: "Downloading rust-analyzer server",
-        gunzip: true,
-        mode: 0o755
+    await downloadWithRetryDialog(state, async () => {
+        await download({
+            url: artifact.browser_download_url,
+            dest,
+            progressTitle: "Downloading rust-analyzer server",
+            gunzip: true,
+            mode: 0o755,
+            overwrite: true,
+        });
     });
 
     // Patching executable if that's NixOS.
@@ -333,3 +344,56 @@ async function getServer(config: Config, state: PersistentState): Promise<string
     await state.updateServerVersion(config.package.version);
     return dest;
 }
+
+async function downloadWithRetryDialog<T>(state: PersistentState, downloadFunc: () => Promise<T>): Promise<T> {
+    while (true) {
+        try {
+            return await downloadFunc();
+        } catch (e) {
+            const selected = await vscode.window.showErrorMessage("Failed to download: " + e.message, {}, {
+                title: "Update Github Auth Token",
+                updateToken: true,
+            }, {
+                title: "Retry download",
+                retry: true,
+            }, {
+                title: "Dismiss",
+            });
+
+            if (selected?.updateToken) {
+                await queryForGithubToken(state);
+                continue;
+            } else if (selected?.retry) {
+                continue;
+            }
+            throw e;
+        };
+    }
+}
+
+async function queryForGithubToken(state: PersistentState): Promise<void> {
+    const githubTokenOptions: vscode.InputBoxOptions = {
+        value: state.githubToken,
+        password: true,
+        prompt: `
+            This dialog allows to store a Github authorization token.
+            The usage of an authorization token will increase the rate
+            limit on the use of Github APIs and can thereby prevent getting
+            throttled.
+            Auth tokens can be created at https://github.com/settings/tokens`,
+    };
+
+    const newToken = await vscode.window.showInputBox(githubTokenOptions);
+    if (newToken === undefined) {
+        // The user aborted the dialog => Do not update the stored token
+        return;
+    }
+
+    if (newToken === "") {
+        log.info("Clearing github token");
+        await state.updateGithubToken(undefined);
+    } else {
+        log.info("Storing new github token");
+        await state.updateGithubToken(newToken);
+    }
+}
diff --git a/editors/code/src/net.ts b/editors/code/src/net.ts
index 5eba2728d..9ba17b7b5 100644
--- a/editors/code/src/net.ts
+++ b/editors/code/src/net.ts
@@ -18,7 +18,8 @@ const OWNER = "rust-analyzer";
 const REPO = "rust-analyzer";
 
 export async function fetchRelease(
-    releaseTag: string
+    releaseTag: string,
+    githubToken: string | null | undefined,
 ): Promise<GithubRelease> {
 
     const apiEndpointPath = `/repos/${OWNER}/${REPO}/releases/tags/${releaseTag}`;
@@ -27,7 +28,12 @@ export async function fetchRelease(
 
     log.debug("Issuing request for released artifacts metadata to", requestUrl);
 
-    const response = await fetch(requestUrl, { headers: { Accept: "application/vnd.github.v3+json" } });
+    const headers: Record<string, string> = { Accept: "application/vnd.github.v3+json" };
+    if (githubToken != null) {
+        headers.Authorization = "token " + githubToken;
+    }
+
+    const response = await fetch(requestUrl, { headers: headers });
 
     if (!response.ok) {
         log.error("Error fetching artifact release info", {
@@ -70,6 +76,7 @@ interface DownloadOpts {
     dest: string;
     mode?: number;
     gunzip?: boolean;
+    overwrite?: boolean;
 }
 
 export async function download(opts: DownloadOpts) {
@@ -79,6 +86,13 @@ export async function download(opts: DownloadOpts) {
     const randomHex = crypto.randomBytes(5).toString("hex");
     const tempFile = path.join(dest.dir, `${dest.name}${randomHex}`);
 
+    if (opts.overwrite) {
+        // Unlinking the exe file before moving new one on its place should prevent ETXTBSY error.
+        await fs.promises.unlink(opts.dest).catch(err => {
+            if (err.code !== "ENOENT") throw err;
+        });
+    }
+
     await vscode.window.withProgress(
         {
             location: vscode.ProgressLocation.Notification,
diff --git a/editors/code/src/persistent_state.ts b/editors/code/src/persistent_state.ts
index 5705eed81..afb652589 100644
--- a/editors/code/src/persistent_state.ts
+++ b/editors/code/src/persistent_state.ts
@@ -38,4 +38,15 @@ export class PersistentState {
     async updateServerVersion(value: string | undefined) {
         await this.globalState.update("serverVersion", value);
     }
+
+    /**
+     * Github authorization token.
+     * This is used for API requests against the Github API.
+     */
+    get githubToken(): string | undefined {
+        return this.globalState.get("githubToken");
+    }
+    async updateGithubToken(value: string | undefined) {
+        await this.globalState.update("githubToken", value);
+    }
 }