Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Studio Token Auth Flow #4931

Merged
merged 22 commits into from
Nov 13, 2023
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Start updating code for studio change
julieg18 committed Nov 7, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
commit 0eb294bfa0571a0cefea036e723f5d5a91e9e598
3 changes: 2 additions & 1 deletion extension/src/setup/index.ts
Original file line number Diff line number Diff line change
@@ -434,7 +434,8 @@ export class Setup
() => this.isPythonExtensionUsed(),
() => this.updatePythonEnvironment(),
() => this.studio.requestStudioAuthentication(),
() => this.studio.openStudioVerifyUserUrl()
() => this.studio.openStudioVerifyUserUrl(),
() => this.sendDataToWebview()
)
this.dispose.track(
this.onDidReceivedWebviewMessage(message =>
63 changes: 42 additions & 21 deletions extension/src/setup/studio.ts
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
import { EventEmitter } from 'vscode'
import { isStudioAccessToken, pollForStudioToken } from './token'
import { isStudioAccessToken } from './token'
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 { openUrl } from '../vscode/external'
import { getCallBackUrl, openUrl, waitForUriResponse } from '../vscode/external'

export class Studio {
protected studioConnectionChanged: EventEmitter<void>
@@ -103,10 +103,18 @@ export class Studio {
)
}

public openStudioVerifyUserUrl() {
const url = this.getStudioVerifyUserUrl()
if (!url) {
return
}
void openUrl(url)
}

public async requestStudioAuthentication() {
const response = await fetch(`${STUDIO_URL}/api/device-login`, {
body: JSON.stringify({
client_name: 'vscode'
client_name: 'VS Code'
}),
headers: {
'Content-Type': 'application/json'
@@ -125,16 +133,19 @@ export class Studio {
user_code: string
device_code: string
}
this.updateStudioUserVerifyDetails(userCode, verificationUri)
void this.requestStudioToken(deviceCode, tokenUri)
}

public openStudioVerifyUserUrl() {
const url = this.getStudioVerifyUserUrl()
if (!url) {
return
}
void openUrl(url)
const callbackUrl = await getCallBackUrl('/studio-complete-auth')
const verificationUrlWithCallback = new URL(verificationUri)

verificationUrlWithCallback.searchParams.append('redirect_uri', callbackUrl)
verificationUrlWithCallback.searchParams.append('code', userCode)
this.updateStudioUserVerifyDetails(
userCode,
verificationUrlWithCallback.toString()
)
void waitForUriResponse('/studio-complete-auth', () =>
this.requestStudioToken(deviceCode, tokenUri)
)
}

private updateStudioUserVerifyDetails(
@@ -173,14 +184,24 @@ export class Studio {
}
}

private async requestStudioToken(
studioDeviceCode: string,
studioTokenRequestUri: string
) {
const token = await pollForStudioToken(
studioTokenRequestUri,
studioDeviceCode
)
private async requestStudioToken(deviceCode: string, tokenUri: string) {
const response = await fetch(tokenUri, {
body: JSON.stringify({
code: deviceCode
}),
headers: {
'Content-Type': 'application/json'
},
method: 'POST'
})

if (response.status !== 200) {
return
}

const { access_token: accessToken } = (await response.json()) as {
access_token: string
}

this.updateStudioUserVerifyDetails(null, undefined)

@@ -190,7 +211,7 @@ export class Studio {
return
}

return this.saveStudioAccessTokenInConfig(cwd, token)
return this.saveStudioAccessTokenInConfig(cwd, accessToken)
}

private accessConfig(cwd: string, ...args: Args) {
julieg18 marked this conversation as resolved.
Show resolved Hide resolved
32 changes: 0 additions & 32 deletions extension/src/setup/token.ts
Original file line number Diff line number Diff line change
@@ -4,35 +4,3 @@ export const isStudioAccessToken = (text?: string): boolean => {
}
return text.startsWith('isat_') && text.length >= 53
}

// chose the simplest way to do this, in reality we need a way to
// offer a straightforward way to stop the polling either because the
// user cancels or possibly because were getting errors
export const pollForStudioToken = async (
tokenUri: string,
deviceCode: string
): Promise<string> => {
const response = await fetch(tokenUri, {
body: JSON.stringify({
code: deviceCode
}),
headers: {
'Content-Type': 'application/json'
},
method: 'POST'
})

if (response.status === 400) {
await new Promise(resolve => setTimeout(resolve, 2000))
return pollForStudioToken(tokenUri, deviceCode)
}
if (response.status !== 200) {
await new Promise(resolve => setTimeout(resolve, 2000))
return pollForStudioToken(tokenUri, deviceCode)
}

const { access_token: accessToken } = (await response.json()) as {
access_token: string
}
return accessToken
}
17 changes: 15 additions & 2 deletions extension/src/setup/webview/messages.ts
Original file line number Diff line number Diff line change
@@ -22,6 +22,7 @@ export class WebviewMessages {
private readonly updatePythonEnv: () => Promise<void>
private readonly requestStudioAuth: () => Promise<void>
private readonly openStudioAuthLink: () => void
private readonly update: () => Promise<void>

constructor(
getWebview: () => BaseWebview<TSetupData> | undefined,
@@ -30,7 +31,8 @@ export class WebviewMessages {
isPythonExtensionUsed: () => Promise<boolean>,
updatePythonEnv: () => Promise<void>,
requestStudioAuth: () => Promise<void>,
openStudioAuthLink: () => void
openStudioAuthLink: () => void,
update: () => Promise<void>
) {
this.getWebview = getWebview
this.initializeGit = initializeGit
@@ -39,6 +41,7 @@ export class WebviewMessages {
this.updatePythonEnv = updatePythonEnv
this.requestStudioAuth = requestStudioAuth
this.openStudioAuthLink = openStudioAuthLink
this.update = update
}

public sendWebviewMessage(data: SetupData) {
@@ -82,7 +85,7 @@ export class WebviewMessages {
case MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE:
return this.updateStudioOffline(message.payload)
case MessageFromWebviewType.REQUEST_STUDIO_TOKEN:
return this.requestStudioAuth()
return this.requestStudioAuthentication()
case MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW:
return commands.executeCommand(RegisteredCommands.EXPERIMENT_SHOW)
case MessageFromWebviewType.REMOTE_ADD:
@@ -135,4 +138,14 @@ export class WebviewMessages {

return autoInstallDvc(isPythonExtensionUsed)
}

private async requestStudioAuthentication() {
sendTelemetryEvent(
EventName.VIEWS_SETUP_REQUEST_STUDIO_AUTH,
undefined,
undefined
)
await this.requestStudioAuth()
return this.update()
}
}
2 changes: 2 additions & 0 deletions extension/src/telemetry/constants.ts
Original file line number Diff line number Diff line change
@@ -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_AUTH: 'view.setup.requestStudioAuth',
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_AUTH]: undefined
[EventName.VIEWS_SETUP_UPDATE_PYTHON_ENVIRONMENT]: undefined
[EventName.VIEWS_SETUP_SHOW_SCM_FOR_COMMIT]: undefined
[EventName.VIEWS_SETUP_INIT_GIT]: undefined
92 changes: 91 additions & 1 deletion extension/src/test/suite/setup/index.test.ts
Original file line number Diff line number Diff line change
@@ -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,
@@ -854,7 +855,96 @@ suite('Setup Test Suite', () => {
).to.be.calledWithExactly('setContext', 'dvc.cli.incompatible', false)
})

it("should handle a message from the webview to save the user's Studio access token", async () => {
it('should handle a message to request a token from studio', async () => {
const { setup, mockFetch } = buildSetup({
disposer: disposable
})
const mockConfig = stub(DvcConfig.prototype, 'config')
mockConfig.resolves('')
const webview = await setup.showWebview()
await webview.isReady()

const mockMessageReceived = getMessageReceivedEmitter(webview)
const mockSendMessage = stub(BaseWebview.prototype, 'show')

// eslint-disable-next-line @typescript-eslint/no-explicit-any
stub(Setup.prototype as any, 'getCliCompatible').returns(true)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this private in the first place?

private getCliCompatible() {
    return this.cliCompatible
  }

That does not seem to have any use since it's accessing a local variable. Is this code for tests only? Is there no other way to change the value? Why not set it public then? cc. @mattseddon (I think you're the one that added it)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there no other way to change the value?

Yes, that is the reason that I remember... https://github.com/iterative/vscode-dvc/pull/3768/files#r1178629676


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 mockToken = 'isat_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA'
mockFetch
.onFirstCall()
.resolves({
json: () => Promise.resolve(mockStudioRes)
} as Fetch.Response)
.onSecondCall()
.resolves({
json: () =>
Promise.resolve({
access_token: mockToken
}),
status: 200
} as Fetch.Response)

const verifyUserStatusSent = new Promise(resolve =>
mockSendMessage.callsFake(data => {
resolve(undefined)
return Promise.resolve(!!data)
})
)

mockMessageReceived.fire({
type: MessageFromWebviewType.REQUEST_STUDIO_TOKEN
})

await verifyUserStatusSent

expect(mockSendMessage).to.be.calledOnce
expect(mockSendMessage).to.be.calledWithMatch({
isStudioConnected: false,
shareLiveToStudio: true,
studioVerifyUserCode: mockStudioRes.user_code
})

// next steps on this test
// check updateStudioUserVerifyDetails to make sure userCode and verifyUrl are updated
// stub waitForUriResponse (can probably wrap a stub in a event)
// call requestStudioToken in stub
// spy requestToken to make sure it's called with the right stuff
// check updateStudioUserVerifyDetails to make sure userCode and verifyUrl are updated
// config check time
})

it('should handle a message from the webview to open the studio verification url', async () => {
const { setup, studio, mockOpenExternal, urlOpenedEvent } = buildSetup({
disposer: disposable
})

const webview = await setup.showWebview()
await webview.isReady()

const mockStudioVerifyUrl =
'https://studio.iterative.ai/auth/device-login'
stub(studio, 'getStudioVerifyUserUrl')
.onFirstCall()
.returns(mockStudioVerifyUrl)

const mockMessageReceived = getMessageReceivedEmitter(webview)

mockMessageReceived.fire({
type: MessageFromWebviewType.OPEN_STUDIO_VERIFY_USER_LINK
})

await urlOpenedEvent
expect(mockOpenExternal).to.be.calledWith(Uri.parse(mockStudioVerifyUrl))
}).timeout(WEBVIEW_TEST_TIMEOUT)

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({
7 changes: 7 additions & 0 deletions extension/src/test/suite/setup/util.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { join } from 'path'
import { EventEmitter, commands, env } from 'vscode'
import * as Fetch from 'node-fetch'
import { Disposer } from '@hediet/std/disposable'
import { fake, spy, stub } from 'sinon'
import { ensureDirSync } from 'fs-extra'
@@ -19,6 +20,7 @@ import { Resource } from '../../../resourceLocator'
import { MIN_CLI_VERSION } from '../../../cli/dvc/contract'
import { Status } from '../../../status'
import { BaseWebview } from '../../../webview'
import { Studio } from '../../../setup/studio'

export const TEMP_DIR = join(dvcDemoPath, 'temp-empty-watcher-dir')

@@ -105,6 +107,8 @@ export const buildSetup = ({

const mockConfig = stub(dvcConfig, 'config').resolves('')

const mockFetch = stub(Fetch, 'default')

const setup = disposer.track(
new Setup(
config,
@@ -133,6 +137,7 @@ export const buildSetup = ({
mockAutoUpgradeDvc,
mockConfig,
mockExecuteCommand,
mockFetch,
mockGetGitRepositoryRoot,
mockGitVersion,
mockGlobalVersion,
@@ -144,6 +149,8 @@ export const buildSetup = ({
mockVersion,
resourceLocator,
setup,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
studio: (setup as any).studio as Studio,
urlOpenedEvent
}
}
23 changes: 22 additions & 1 deletion extension/src/vscode/external.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,24 @@
import { env, Uri } from 'vscode'
import { env, ProviderResult, Uri, window } from 'vscode'

export const openUrl = (url: string) => env.openExternal(Uri.parse(url))

export const getCallBackUrl = async (path: string) => {
const uri = await env.asExternalUri(
Uri.parse(`${env.uriScheme}://iterative.dvc${path}`)
)

return uri.toString()
}

export const waitForUriResponse = (
path: string,
onResponse: (uri: Uri) => unknown
) => {
window.registerUriHandler({
handleUri(uri: Uri): ProviderResult<void> {
if (uri.path.includes(path)) {
onResponse(uri)
}
}
})
}