Add Studio Token Auth Flow (#4931)

Author: julieg18

extension/src/__mocks__/vscode.ts

@@ -24,6 +24,7 @@ export enum TreeItemCollapsibleState {
 export const Uri = {
   file: URI.file,
   from: URI.from,
+  parse: URI.parse,
   joinPath: Utils.joinPath
 }
 export const window = {

extension/src/setup/index.ts

@@ -424,7 +424,8 @@ export class Setup
       () => this.initializeGit(),
       (offline: boolean) => this.studio.updateStudioOffline(offline),
       () => this.isPythonExtensionUsed(),
-      () => this.updatePythonEnvironment()
+      () => this.updatePythonEnvironment(),
+      () => this.studio.requestStudioTokenAuthentication()
     )
     this.dispose.track(
       this.onDidReceivedWebviewMessage(message =>

extension/src/setup/studio.ts

@@ -1,9 +1,13 @@
 import { Event, EventEmitter } from 'vscode'
+import fetch from 'node-fetch'
+import { STUDIO_URL } from './webview/contract'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
 import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
 import { Args, ConfigKey, Flag } from '../cli/dvc/constants'
 import { ContextKey, setContextValue } from '../vscode/context'
 import { Disposable } from '../class/dispose'
+import { getCallBackUrl, openUrl, waitForUriResponse } from '../vscode/external'
+import { Modal } from '../vscode/modal'
 
 export const isStudioAccessToken = (text?: string): boolean => {
   if (!text) {
@@ -102,6 +106,80 @@ export class Studio extends Disposable {
     )
   }
 
+  public async requestStudioTokenAuthentication() {
+    const response = await this.fetchFromStudio(
+      `${STUDIO_URL}/api/device-login`,
+      {
+        client_name: 'VS Code'
+      }
+    )
+
+    const {
+      token_uri: tokenUri,
+      verification_uri: verificationUri,
+      user_code: userCode,
+      device_code: deviceCode
+    } = (await response.json()) as {
+      token_uri: string
+      verification_uri: string
+      user_code: string
+      device_code: string
+    }
+
+    const callbackUrl = await getCallBackUrl('/studio-complete-auth')
+    const verificationUrlWithParams = new URL(verificationUri)
+
+    verificationUrlWithParams.searchParams.append('redirect_uri', callbackUrl)
+    verificationUrlWithParams.searchParams.append('code', userCode)
+
+    await openUrl(verificationUrlWithParams.toString())
+    void waitForUriResponse('/studio-complete-auth', () => {
+      void this.requestStudioToken(deviceCode, tokenUri)
+    })
+  }
+
+  private fetchFromStudio(reqUri: string, body: Record<string, unknown>) {
+    return fetch(reqUri, {
+      body: JSON.stringify(body),
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      method: 'POST'
+    })
+  }
+
+  private async fetchStudioToken(deviceCode: string, tokenUri: string) {
+    const response = await this.fetchFromStudio(tokenUri, {
+      code: deviceCode
+    })
+
+    if (response.status !== 200) {
+      const { detail } = (await response.json()) as {
+        detail: string
+      }
+      return Modal.errorWithOptions(
+        `Unable to get token. Failed with "${detail}"`
+      )
+    }
+
+    const { access_token: accessToken } = (await response.json()) as {
+      access_token: string
+    }
+
+    return accessToken
+  }
+
+  private async requestStudioToken(deviceCode: string, tokenUri: string) {
+    const token = await this.fetchStudioToken(deviceCode, tokenUri)
+    const cwd = this.getCwd()
+
+    if (!token || !cwd) {
+      return
+    }
+
+    return this.saveStudioAccessTokenInConfig(cwd, token)
+  }
+
   private async setStudioValues() {
     const cwd = this.getCwd()
 

extension/src/setup/webview/messages.ts

@@ -1,5 +1,5 @@
 import { commands } from 'vscode'
-import { STUDIO_URL, SetupData, SetupData as TSetupData } from './contract'
+import { SetupData, SetupData as TSetupData } from './contract'
 import { Logger } from '../../common/logger'
 import {
   MessageFromWebview,
@@ -12,7 +12,6 @@ import {
   RegisteredCliCommands,
   RegisteredCommands
 } from '../../commands/external'
-import { openUrl } from '../../vscode/external'
 import { autoInstallDvc, autoUpgradeDvc } from '../autoInstall'
 
 export class WebviewMessages {
@@ -21,19 +20,22 @@ export class WebviewMessages {
   private readonly updateStudioOffline: (offline: boolean) => Promise<void>
   private readonly isPythonExtensionUsed: () => Promise<boolean>
   private readonly updatePythonEnv: () => Promise<void>
+  private readonly requestToken: () => Promise<void>
 
   constructor(
     getWebview: () => BaseWebview<TSetupData> | undefined,
     initializeGit: () => void,
     updateStudioOffline: (shareLive: boolean) => Promise<void>,
     isPythonExtensionUsed: () => Promise<boolean>,
-    updatePythonEnv: () => Promise<void>
+    updatePythonEnv: () => Promise<void>,
+    requestStudioToken: () => Promise<void>
   ) {
     this.getWebview = getWebview
     this.initializeGit = initializeGit
     this.updateStudioOffline = updateStudioOffline
     this.isPythonExtensionUsed = isPythonExtensionUsed
     this.updatePythonEnv = updatePythonEnv
+    this.requestToken = requestStudioToken
   }
 
   public sendWebviewMessage(data: SetupData) {
@@ -64,10 +66,6 @@ export class WebviewMessages {
         return commands.executeCommand(
           RegisteredCommands.EXTENSION_SETUP_WORKSPACE
         )
-      case MessageFromWebviewType.OPEN_STUDIO:
-        return this.openStudio()
-      case MessageFromWebviewType.OPEN_STUDIO_PROFILE:
-        return this.openStudioProfile()
       case MessageFromWebviewType.SAVE_STUDIO_TOKEN:
         return commands.executeCommand(
           RegisteredCommands.ADD_STUDIO_ACCESS_TOKEN
@@ -78,6 +76,8 @@ export class WebviewMessages {
         )
       case MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE:
         return this.updateStudioOffline(message.payload)
+      case MessageFromWebviewType.REQUEST_STUDIO_TOKEN:
+        return this.requestStudioToken()
       case MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW:
         return commands.executeCommand(RegisteredCommands.EXPERIMENT_SHOW)
       case MessageFromWebviewType.REMOTE_ADD:
@@ -131,11 +131,12 @@ export class WebviewMessages {
     return autoInstallDvc(isPythonExtensionUsed)
   }
 
-  private openStudio() {
-    return openUrl(STUDIO_URL)
-  }
-
-  private openStudioProfile() {
-    return openUrl(`${STUDIO_URL}/user/_/profile?section=accessToken`)
+  private requestStudioToken() {
+    sendTelemetryEvent(
+      EventName.VIEWS_SETUP_REQUEST_STUDIO_TOKEN,
+      undefined,
+      undefined
+    )
+    return this.requestToken()
   }
 }

extension/src/telemetry/constants.ts

@@ -109,6 +109,7 @@ export const EventName = Object.assign(
     VIEWS_SETUP_FOCUS_CHANGED: 'views.setup.focusChanged',
     VIEWS_SETUP_INIT_GIT: 'views.setup.initializeGit',
     VIEWS_SETUP_INSTALL_DVC: 'views.setup.installDvc',
+    VIEWS_SETUP_REQUEST_STUDIO_TOKEN: 'view.setup.requestStudioToken',
     VIEWS_SETUP_SHOW_SCM_FOR_COMMIT: 'views.setup.showScmForCommit',
     VIEWS_SETUP_UPDATE_PYTHON_ENVIRONMENT:
       'views.setup.updatePythonEnvironment',
@@ -323,6 +324,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_SETUP_CLOSE]: undefined
   [EventName.VIEWS_SETUP_CREATED]: undefined
   [EventName.VIEWS_SETUP_FOCUS_CHANGED]: undefined
+  [EventName.VIEWS_SETUP_REQUEST_STUDIO_TOKEN]: undefined
   [EventName.VIEWS_SETUP_UPDATE_PYTHON_ENVIRONMENT]: undefined
   [EventName.VIEWS_SETUP_SHOW_SCM_FOR_COMMIT]: undefined
   [EventName.VIEWS_SETUP_INIT_GIT]: undefined

extension/src/test/suite/setup/index.test.ts

@@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, it, suite } from 'mocha'
 import { ensureFileSync, remove } from 'fs-extra'
 import { expect } from 'chai'
 import { SinonStub, restore, spy, stub } from 'sinon'
+import * as Fetch from 'node-fetch'
 import {
   MessageItem,
   QuickPickItem,
@@ -43,12 +44,15 @@ import { MIN_CLI_VERSION } from '../../../cli/dvc/contract'
 import { run } from '../../../setup/runner'
 import * as Python from '../../../extensions/python'
 import { ContextKey } from '../../../vscode/context'
+import * as ExternalUtil from '../../../vscode/external'
 import { Setup } from '../../../setup'
-import { SetupSection } from '../../../setup/webview/contract'
+import { STUDIO_URL, SetupSection } from '../../../setup/webview/contract'
 import { getFirstWorkspaceFolder } from '../../../vscode/workspaceFolders'
 import { Response } from '../../../vscode/response'
 import { DvcConfig } from '../../../cli/dvc/config'
 import * as QuickPickUtil from '../../../setup/quickPick'
+import { EventName } from '../../../telemetry/constants'
+import { Modal } from '../../../vscode/modal'
 
 suite('Setup Test Suite', () => {
   const disposable = Disposable.fn()
@@ -850,49 +854,137 @@ suite('Setup Test Suite', () => {
       ).to.be.calledWithExactly('setContext', 'dvc.cli.incompatible', false)
     })
 
-    it('should handle a message from the webview to open Studio', async () => {
-      const { mockOpenExternal, setup, urlOpenedEvent } = buildSetup({
+    it('should handle a message from the webview to request a token from studio', async () => {
+      const { setup, mockFetch, studio } = buildSetup({
         disposer: disposable
       })
 
+      const mockConfig = stub(DvcConfig.prototype, 'config')
+      mockConfig.resolves('')
       const webview = await setup.showWebview()
       await webview.isReady()
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      stub(Setup.prototype as any, 'getCliCompatible').returns(true)
 
       const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockGetCallbackUrl = stub(ExternalUtil, 'getCallBackUrl')
+      const mockOpenUrl = stub(ExternalUtil, 'openUrl')
+      const mockWaitForUriRes = stub(ExternalUtil, 'waitForUriResponse')
+      const mockStudioRes = {
+        device_code: 'Yi-NPd9ggvNUDBcam5bP8iivbtLhnqVgM_lSSbilqNw',
+        token_uri: 'https://studio.iterative.ai/api/device-login/token',
+        user_code: '40DWMKNA',
+        verification_uri: 'https://studio.iterative.ai/auth/device-login'
+      }
+      const mockCallbackUrl = 'url-to-vscode'
+
+      mockFetch.onFirstCall().resolves({
+        json: () => Promise.resolve(mockStudioRes)
+      } as Fetch.Response)
+      mockGetCallbackUrl.onFirstCall().resolves(mockCallbackUrl)
+
+      const callbackUriHandlerEvent: Promise<() => unknown> = new Promise(
+        resolve =>
+          mockWaitForUriRes.onFirstCall().callsFake((_, onResponse) => {
+            resolve(onResponse)
+          })
+      )
 
       mockMessageReceived.fire({
-        type: MessageFromWebviewType.OPEN_STUDIO
+        type: MessageFromWebviewType.REQUEST_STUDIO_TOKEN
       })
 
-      await urlOpenedEvent
-      expect(mockOpenExternal).to.be.calledWith(
-        Uri.parse('https://studio.iterative.ai')
+      const mockOnStudioResponse = await callbackUriHandlerEvent
+
+      expect(mockSendTelemetryEvent).to.be.calledOnce
+      expect(mockSendTelemetryEvent).to.be.calledWith(
+        EventName.VIEWS_SETUP_REQUEST_STUDIO_TOKEN,
+        undefined,
+        undefined
+      )
+      expect(mockFetch).to.be.calledOnce
+      expect(mockFetch).to.be.calledOnceWithExactly(
+        `${STUDIO_URL}/api/device-login`,
+        {
+          body: JSON.stringify({
+            client_name: 'VS Code'
+          }),
+          headers: {
+            'Content-Type': 'application/json'
+          },
+          method: 'POST'
+        }
+      )
+      expect(mockGetCallbackUrl).to.be.calledOnce
+      expect(mockGetCallbackUrl).to.be.calledWith('/studio-complete-auth')
+      expect(mockOpenUrl).to.be.calledOnce
+      expect(mockOpenUrl).to.be.calledWith(
+        `${mockStudioRes.verification_uri}?redirect_uri=${mockCallbackUrl}&code=${mockStudioRes.user_code}`
       )
-    }).timeout(WEBVIEW_TEST_TIMEOUT)
 
-    it("should handle a message from the webview to open the user's Studio profile", async () => {
-      const { setup, mockOpenExternal, urlOpenedEvent } = buildSetup({
-        disposer: disposable
-      })
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const mockGetCwd = stub(studio as any, 'getCwd')
+      const mockModalShowError = stub(Modal, 'errorWithOptions')
+      const mockSaveStudioToken = stub(studio, 'saveStudioAccessTokenInConfig')
+      mockFetch.onSecondCall().resolves({
+        json: () =>
+          Promise.resolve({ detail: 'Request failed for some reason.' }),
+        status: 500
+      } as Fetch.Response)
+
+      const failedTokenEvent = new Promise(resolve =>
+        mockGetCwd.onFirstCall().callsFake(() => {
+          resolve(undefined)
+          return dvcDemoPath
+        })
+      )
 
-      const webview = await setup.showWebview()
-      await webview.isReady()
+      mockOnStudioResponse()
 
-      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      await failedTokenEvent
 
-      mockMessageReceived.fire({
-        type: MessageFromWebviewType.OPEN_STUDIO_PROFILE
+      expect(mockFetch).to.be.calledTwice
+      expect(mockFetch).to.be.calledWithExactly(mockStudioRes.token_uri, {
+        body: JSON.stringify({
+          code: mockStudioRes.device_code
+        }),
+        headers: {
+          'Content-Type': 'application/json'
+        },
+        method: 'POST'
       })
+      expect(mockModalShowError).to.be.calledOnce
+      expect(mockModalShowError).to.be.calledWithExactly(
+        'Unable to get token. Failed with "Request failed for some reason."'
+      )
+      expect(mockSaveStudioToken).not.to.be.called
 
-      await urlOpenedEvent
-      expect(mockOpenExternal).to.be.calledWith(
-        Uri.parse(
-          'https://studio.iterative.ai/user/_/profile?section=accessToken'
-        )
+      const mockToken = 'isat_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA'
+      mockFetch.onThirdCall().resolves({
+        json: () => Promise.resolve({ access_token: mockToken }),
+        status: 200
+      } as Fetch.Response)
+      const tokenEvent = new Promise(resolve =>
+        mockGetCwd.onSecondCall().callsFake(() => {
+          resolve(undefined)
+          return dvcDemoPath
+        })
+      )
+
+      mockOnStudioResponse()
+
+      await tokenEvent
+
+      expect(mockFetch).to.be.calledThrice
+      expect(mockSaveStudioToken).to.be.calledOnce
+      expect(mockSaveStudioToken).to.be.calledWithExactly(
+        dvcDemoPath,
+        mockToken
       )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
-    it("should handle a message from the webview to save the user's Studio access token", async () => {
+    it("should handle a message from the webview to manually save the user's Studio access token", async () => {
       const mockToken = 'isat_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA'
 
       const { setup, mockExecuteCommand, messageSpy } = buildSetup({