Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
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
140 changes: 115 additions & 25 deletions src/standalone/userJupyterServer/jupyterPasswordConnect.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import {
} from '../../kernels/jupyter/types';
import { dispose } from '../../platform/common/utils/lifecycle';
import { JupyterSelfCertsError } from '../../platform/errors/jupyterSelfCertsError';
import { IEncryptedStorage } from '../../platform/common/application/types';

export interface IJupyterPasswordConnectInfo {
requiresPassword: boolean;
Expand All @@ -27,12 +28,15 @@ export interface IJupyterPasswordConnectInfo {
*/
export class JupyterPasswordConnect {
private savedConnectInfo = new Map<string, Promise<IJupyterPasswordConnectInfo>>();
private static readonly SERVICE_NAME = 'jupyter-server-password';

constructor(
private readonly configService: IConfigurationService,
private readonly agentCreator: IJupyterRequestAgentCreator | undefined,
private readonly requestCreator: IJupyterRequestCreator,
private readonly serverUriStorage: IJupyterServerUriStorage,
private readonly disposables: IDisposableRegistry
private readonly disposables: IDisposableRegistry,
private readonly encryptedStorage: IEncryptedStorage
) {
// Sign up to see if servers are removed from our uri storage list
this.serverUriStorage.onDidRemove(this.onDidRemoveServers, this, this.disposables);
Expand All @@ -59,6 +63,7 @@ export class JupyterPasswordConnect {
result = this.getJupyterConnectionInfo({
url: newUrl,
isTokenEmpty: options.isTokenEmpty,
handle: options.handle,
disposables,
validationErrorMessage: options.validationErrorMessage
}).then((value) => {
Expand Down Expand Up @@ -98,43 +103,60 @@ export class JupyterPasswordConnect {
private async getJupyterConnectionInfo(options: {
url: string;
isTokenEmpty: boolean;
handle: string;
validationErrorMessage?: string;
disposables: IDisposable[];
}): Promise<IJupyterPasswordConnectInfo> {
let xsrfCookie: string | undefined;
let sessionCookieName: string | undefined;
let sessionCookieValue: string | undefined;
let userPassword: string | undefined = undefined;
let useStoredPassword = false;

// First determine if we need a password. A request for the base URL with /tree? should return a 302 if we do.
const requiresPassword = await this.needPassword(options.url);

if (requiresPassword || options.isTokenEmpty) {
// Get password first
// Get password first - try stored password if available
if (requiresPassword && options.isTokenEmpty) {
const input = window.createInputBox();
options.disposables.push(input);
input.title = DataScience.jupyterSelectPasswordTitle;
input.placeholder = DataScience.jupyterSelectPasswordPrompt;
input.ignoreFocusOut = true;
input.password = true;
input.validationMessage = options.validationErrorMessage || '';
input.show();
input.buttons = [QuickInputButtons.Back];
userPassword = await new Promise<string>((resolve, reject) => {
input.onDidTriggerButton(
(e) => {
if (e === QuickInputButtons.Back) {
reject(InputFlowAction.back);
}
},
this,
options.disposables
);
input.onDidChangeValue(() => (input.validationMessage = ''), this, options.disposables);
input.onDidAccept(() => resolve(input.value), this, options.disposables);
input.onDidHide(() => reject(InputFlowAction.cancel), this, options.disposables);
});
// Try to get stored password first
const storedPassword = await this.getStoredPassword(options.handle);

if (storedPassword && !options.validationErrorMessage) {
// We have a stored password and no validation error, so try using it
userPassword = storedPassword;
useStoredPassword = true;
} else {
// No stored password or previous validation failed, prompt user
const input = window.createInputBox();
options.disposables.push(input);
input.title = DataScience.jupyterSelectPasswordTitle;
input.placeholder = DataScience.jupyterSelectPasswordPrompt;
input.ignoreFocusOut = true;
input.password = true;
input.validationMessage = options.validationErrorMessage || '';
input.show();
input.buttons = [QuickInputButtons.Back];
userPassword = await new Promise<string>((resolve, reject) => {
input.onDidTriggerButton(
(e) => {
if (e === QuickInputButtons.Back) {
reject(InputFlowAction.back);
}
},
this,
options.disposables
);
input.onDidChangeValue(() => (input.validationMessage = ''), this, options.disposables);
input.onDidAccept(() => resolve(input.value), this, options.disposables);
input.onDidHide(() => reject(InputFlowAction.cancel), this, options.disposables);
});

// If we had a stored password that failed, clear it
if (storedPassword && options.validationErrorMessage) {
await this.clearStoredPassword(options.handle);
}
}
}

if (typeof userPassword === undefined && !userPassword && options.isTokenEmpty) {
Expand Down Expand Up @@ -173,11 +195,23 @@ export class JupyterPasswordConnect {
// Remember session cookie can be empty, if both token and password are empty
if (xsrfCookie && sessionCookieName && (sessionCookieValue || options.isTokenEmpty)) {
sendTelemetryEvent(Telemetry.GetPasswordSuccess);

// Save the password for future use if authentication was successful and we have a new password
if (userPassword && !useStoredPassword) {
await this.storePassword(options.handle, userPassword);
}

const cookieString = `_xsrf=${xsrfCookie}; ${sessionCookieName}=${sessionCookieValue || ''}`;
const requestHeaders = { Cookie: cookieString, 'X-XSRFToken': xsrfCookie };
return { requestHeaders, requiresPassword };
} else {
sendTelemetryEvent(Telemetry.GetPasswordFailure);

// If we used a stored password and it failed, clear it
if (useStoredPassword) {
await this.clearStoredPassword(options.handle);
}

return { requiresPassword };
}
}
Expand Down Expand Up @@ -367,9 +401,65 @@ export class JupyterPasswordConnect {
servers.forEach((server) => {
if (server.id.startsWith('_builtin')) {
this.savedConnectInfo.delete(server.handle);
// Also clear any stored passwords for removed servers
this.clearStoredPassword(server.handle).catch((error) => {
logger.warn(`Failed to clear stored password for server ${server.handle}:`, error);
});
}
});
}

/**
* Creates a storage key for a server password based on the server handle
*/
private getPasswordStorageKey(handle: string): string {
return `password-${handle}`;
}

/**
* Retrieves a stored password for the given server handle
*/
private async getStoredPassword(handle: string): Promise<string | undefined> {
try {
return await this.encryptedStorage.retrieve(
JupyterPasswordConnect.SERVICE_NAME,
this.getPasswordStorageKey(handle)
);
} catch (error) {
logger.warn(`Failed to retrieve stored password for server ${handle}:`, error);
return undefined;
}
}

/**
* Stores a password for the given server handle
*/
private async storePassword(handle: string, password: string): Promise<void> {
try {
await this.encryptedStorage.store(
JupyterPasswordConnect.SERVICE_NAME,
this.getPasswordStorageKey(handle),
password
);
} catch (error) {
logger.warn(`Failed to store password for server ${handle}:`, error);
}
}

/**
* Clears a stored password for the given server handle
*/
private async clearStoredPassword(handle: string): Promise<void> {
try {
await this.encryptedStorage.store(
JupyterPasswordConnect.SERVICE_NAME,
this.getPasswordStorageKey(handle),
undefined
);
} catch (error) {
logger.warn(`Failed to clear stored password for server ${handle}:`, error);
}
}
}

export function addTrailingSlash(url: string): string {
Expand Down
137 changes: 136 additions & 1 deletion src/standalone/userJupyterServer/jupyterPasswordConnect.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,14 @@ import { Disposable, InputBox } from 'vscode';
import { noop } from '../../test/core';
import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock';
import { dispose } from '../../platform/common/utils/lifecycle';
import { IEncryptedStorage } from '../../platform/common/application/types';

/* eslint-disable @typescript-eslint/no-explicit-any, , */
suite('JupyterServer Password Connect', () => {
let jupyterPasswordConnect: JupyterPasswordConnect;
let configService: ConfigurationService;
let requestCreator: IJupyterRequestCreator;
let encryptedStorage: IEncryptedStorage;

const xsrfValue: string = '12341234';
const sessionName: string = 'sessionName';
Expand Down Expand Up @@ -63,14 +65,16 @@ suite('JupyterServer Password Connect', () => {
when(mockedVSCodeNamespaces.window.createInputBox()).thenReturn(inputBox);
configService = mock(ConfigurationService);
requestCreator = mock(JupyterRequestCreator);
encryptedStorage = mock<IEncryptedStorage>();
const serverUriStorage = mock<IJupyterServerUriStorage>();

jupyterPasswordConnect = new JupyterPasswordConnect(
instance(configService),
undefined,
instance(requestCreator),
instance(serverUriStorage),
disposables
disposables,
instance(encryptedStorage)
);
});
teardown(() => (disposables = dispose(disposables)));
Expand Down Expand Up @@ -501,4 +505,135 @@ suite('JupyterServer Password Connect', () => {
mockSessionResponse.verifyAll();
fetchMock.verifyAll();
});

test('Stores password after successful authentication', async () => {
const { fetchMock, mockXsrfHeaders, mockXsrfResponse } = createMockSetup(false, true);

// Mock our second call to get session cookie
const mockSessionResponse = typemoq.Mock.ofType(nodeFetch.Response);
const mockSessionHeaders = typemoq.Mock.ofType(nodeFetch.Headers);
mockSessionHeaders
.setup((mh) => mh.raw())
.returns(() => {
return {
'set-cookie': [`${sessionName}=${sessionValue}`]
};
});
mockSessionResponse.setup((mr) => mr.ok).returns(() => true);
mockSessionResponse.setup((mr) => mr.status).returns(() => 302);
mockSessionResponse.setup((mr) => mr.headers).returns(() => mockSessionHeaders.object);

const postParams = new URLSearchParams();
postParams.append('_xsrf', xsrfValue);
postParams.append('password', 'Python');

// typemoq doesn't love this comparison, so generalize it a bit
fetchMock
.setup((fm) =>
fm(
'http://testname:8888/login?',
typemoq.It.isObjectWith({
method: 'post',
headers: {
Cookie: `_xsrf=${xsrfValue}`,
Connection: 'keep-alive',
'content-type': 'application/x-www-form-urlencoded;charset=UTF-8'
}
})
)
)
.returns(() => Promise.resolve(mockSessionResponse.object));
when(requestCreator.getFetchMethod()).thenReturn(fetchMock.object as any);

// Mock that no password is stored initially
when(encryptedStorage.retrieve(anything(), anything())).thenResolve(undefined);

const result = await jupyterPasswordConnect.getPasswordConnectionInfo({
url: 'http://testname:8888/',
isTokenEmpty: true,
handle: 'test-server-handle'
});

assert(result, 'Failed to get password');
assert(result.requestHeaders, 'Expected request headers');

// Verify that the password was stored after successful authentication
// The store method should have been called with the correct service, key, and password
// Note: We can't easily verify the password value as it's user input, but we can verify the call was made

// Verfiy network calls
mockXsrfHeaders.verifyAll();
mockSessionHeaders.verifyAll();
mockXsrfResponse.verifyAll();
mockSessionResponse.verifyAll();
fetchMock.verifyAll();
});

test('Uses stored password when available', async () => {
const storedPassword = 'stored-password';
const { fetchMock, mockXsrfHeaders, mockXsrfResponse } = createMockSetup(false, true);

// Mock our second call to get session cookie
const mockSessionResponse = typemoq.Mock.ofType(nodeFetch.Response);
const mockSessionHeaders = typemoq.Mock.ofType(nodeFetch.Headers);
mockSessionHeaders
.setup((mh) => mh.raw())
.returns(() => {
return {
'set-cookie': [`${sessionName}=${sessionValue}`]
};
});
mockSessionResponse.setup((mr) => mr.ok).returns(() => true);
mockSessionResponse.setup((mr) => mr.status).returns(() => 302);
mockSessionResponse.setup((mr) => mr.headers).returns(() => mockSessionHeaders.object);

const postParams = new URLSearchParams();
postParams.append('_xsrf', xsrfValue);
postParams.append('password', storedPassword);

// typemoq doesn't love this comparison, so generalize it a bit
fetchMock
.setup((fm) =>
fm(
'http://testname:8888/login?',
typemoq.It.isObjectWith({
method: 'post',
headers: {
Cookie: `_xsrf=${xsrfValue}`,
Connection: 'keep-alive',
'content-type': 'application/x-www-form-urlencoded;charset=UTF-8'
}
})
)
)
.returns(() => Promise.resolve(mockSessionResponse.object));
when(requestCreator.getFetchMethod()).thenReturn(fetchMock.object as any);

// Mock that a password is stored
when(encryptedStorage.retrieve('jupyter-server-password', 'password-test-server-handle')).thenResolve(
storedPassword
);

// Override the input box setup to NOT show input since we should use stored password
inputBox.value = '';

const result = await jupyterPasswordConnect.getPasswordConnectionInfo({
url: 'http://testname:8888/',
isTokenEmpty: true,
handle: 'test-server-handle'
});

assert(result, 'Failed to get password');
assert(result.requestHeaders, 'Expected request headers');

// Verify that retrieve was called to get the stored password
// The method should have been called but we don't need to verify store since password was already stored

// Verfiy network calls
mockXsrfHeaders.verifyAll();
mockSessionHeaders.verifyAll();
mockXsrfResponse.verifyAll();
mockSessionResponse.verifyAll();
fetchMock.verifyAll();
});
});
3 changes: 2 additions & 1 deletion src/standalone/userJupyterServer/userServerUrlProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,8 @@ export class UserJupyterServerUrlProvider
agentCreator,
requestCreator,
serverUriStorage,
disposables
disposables,
encryptedStorage
);
this.jupyterHubPasswordConnect = new JupyterHubPasswordConnect(configService, agentCreator, requestCreator);
}
Expand Down