feat: Change parseCookie command to read from HTTP cookies including …

…AccessToken & StudioSessionToken
aws · Dec 22, 2024 · 9cc5276 · 9cc5276
1 parent 1acde5e
commit 9cc5276
Show file tree

Hide file tree

Showing 3 changed files with 154 additions and 1 deletion.
diff --git a/patched-vscode/src/vs/server/node/remoteExtensionHostAgentServer.ts b/patched-vscode/src/vs/server/node/remoteExtensionHostAgentServer.ts
@@ -169,6 +169,46 @@ class RemoteExtensionHostAgentServer extends Disposable implements IServerAPI {
 			return serveFile(filePath, CacheControl.ETAG, this._logService, req, res, responseHeaders);
 		}
 
+		if (pathname === '/load-cookies') {
+			let entries = parsedUrl.query['entry'] ?? [];
+			let workspace = (Array.isArray(parsedUrl.query['workspace']) ? parsedUrl.query['workspace'][0] : parsedUrl.query['workspace']) ?? '/';
+
+			if (!Array.isArray(entries)) {
+				entries = [entries];
+			}
+
+			const requestCookie = req.headers.cookie;
+			if (!requestCookie) {
+				res.writeHead(200, { 'Content-Type' : 'text/plain' });
+				return void res.end('');
+			}
+
+			const parsedCookies = requestCookie.split(';').reduce((result, cookie) => {
+				const [name, ...rest] = cookie.split('=');
+				const value = rest.join('=').trim();
+				const key = name.trim();
+				if (entries.includes(key)) {
+					result[key] = decodeURIComponent(value);
+				}
+				return result;
+			}, {} as Record<string, string>);
+
+			const dirPath = join(workspace, '.aws/sso');
+			const filePath = join(dirPath, '/cookies.json');
+			try {
+				if (!fs.existsSync(dirPath)) {
+					fs.mkdirSync(dirPath, { recursive: true });
+				}
+				fs.writeFileSync(filePath,  JSON.stringify(parsedCookies));	
+			} catch {
+				res.writeHead(200, { 'Content-Type' : 'text/plain' });
+				return void res.end('');			
+			}
+
+			res.writeHead(200, { 'Content-Type' : 'text/plain' });
+			return void res.end(filePath);
+		}
+
 		// workbench web UI
 		if (this._webClientServer) {
 			this._webClientServer.handle(req, res, parsedUrl);

diff --git a/patched-vscode/src/vs/workbench/browser/client.ts b/patched-vscode/src/vs/workbench/browser/client.ts
@@ -16,6 +16,7 @@ export class SagemakerServerClient extends Disposable {
 
 	static LOGOUT_COMMAND_ID = 'sagemaker.logout';
 	static COOKIE_COMMAND_ID = 'sagemaker.parseCookies';
+	static COOKIE_LOAD_COMMAND_ID = 'sagemaker.loadCookies';
 
 	private registerSagemakerCommands() {
 		const authMode: string | undefined = this.getCookieValue('authMode');
@@ -24,6 +25,16 @@ export class SagemakerServerClient extends Disposable {
 		const ssoExpiryTimestamp: string | undefined = this.getCookieValue('ssoExpiryTimestamp')
 		const redirectURL: string | undefined = this.getCookieValue('redirectURL')
 
+		const cookieEntries = [
+			'authMode',
+			'expiryTime',
+			'studioUserProfileName',
+			'ssoExpiryTimestamp',
+			'redirectURL',
+			'AccessToken',
+			'StudioSessionToken',
+		];
+
 		this.logService.debug('Registering sagemaker commands...');
 
 		CommandsRegistry.registerCommand(SagemakerServerClient.COOKIE_COMMAND_ID, () => {
@@ -36,6 +47,26 @@ export class SagemakerServerClient extends Disposable {
 			};
 		});
 
+		CommandsRegistry.registerCommand(
+			SagemakerServerClient.COOKIE_LOAD_COMMAND_ID,
+			async () => {
+				try {
+					const entriesQueryParam = cookieEntries
+						.map((entry) => `entry=${entry}`)
+						.join("&");
+					const urlParams = new URLSearchParams(window.location.search);
+					const workspaceFolder = urlParams.get('folder');
+					if (workspaceFolder) {
+						const response = await fetch(`/load-cookies?workspace=${encodeURIComponent(workspaceFolder)}&${entriesQueryParam}`);
+						return await response.text();
+					}
+					return '';
+				} catch (error) {
+					return '';
+				}
+			}
+		);
+
 		CommandsRegistry.registerCommand(SagemakerServerClient.LOGOUT_COMMAND_ID, () => {
 			const currentUrl = new URL(window.location.href);
 			const hostname = currentUrl.hostname;

diff --git a/patches/sagemaker-integration.diff b/patches/sagemaker-integration.diff
@@ -2,7 +2,7 @@ Index: sagemaker-code-editor/vscode/src/vs/workbench/browser/client.ts
 ===================================================================
 --- /dev/null
 +++ sagemaker-code-editor/vscode/src/vs/workbench/browser/client.ts
-@@ -0,0 +1,61 @@
+@@ -0,0 +1,92 @@
 +import { Disposable } from 'vs/base/common/lifecycle';
 +import { CommandsRegistry } from 'vs/platform/commands/common/commands';
 +import { MenuId, MenuRegistry } from "vs/platform/actions/common/actions";
@@ -21,6 +21,7 @@ Index: sagemaker-code-editor/vscode/src/vs/workbench/browser/client.ts
 +
 +	static LOGOUT_COMMAND_ID = 'sagemaker.logout';
 +	static COOKIE_COMMAND_ID = 'sagemaker.parseCookies';
++	static COOKIE_LOAD_COMMAND_ID = 'sagemaker.loadCookies';
 +
 +	private registerSagemakerCommands() {
 +		const authMode: string | undefined = this.getCookieValue('authMode');
@@ -29,6 +30,16 @@ Index: sagemaker-code-editor/vscode/src/vs/workbench/browser/client.ts
 +		const ssoExpiryTimestamp: string | undefined = this.getCookieValue('ssoExpiryTimestamp')
 +		const redirectURL: string | undefined = this.getCookieValue('redirectURL')
 +
++		const cookieEntries = [
++			'authMode',
++			'expiryTime',
++			'studioUserProfileName',
++			'ssoExpiryTimestamp',
++			'redirectURL',
++			'AccessToken',
++			'StudioSessionToken',
++		];
++
 +		this.logService.debug('Registering sagemaker commands...');
 +
 +		CommandsRegistry.registerCommand(SagemakerServerClient.COOKIE_COMMAND_ID, () => {
@@ -41,6 +52,26 @@ Index: sagemaker-code-editor/vscode/src/vs/workbench/browser/client.ts
 +			};
 +		});
 +
++		CommandsRegistry.registerCommand(
++			SagemakerServerClient.COOKIE_LOAD_COMMAND_ID,
++			async () => {
++				try {
++					const entriesQueryParam = cookieEntries
++						.map((entry) => `entry=${entry}`)
++						.join("&");
++					const urlParams = new URLSearchParams(window.location.search);
++					const workspaceFolder = urlParams.get('folder');
++					if (workspaceFolder) {
++						const response = await fetch(`/load-cookies?workspace=${encodeURIComponent(workspaceFolder)}&${entriesQueryParam}`);
++						return await response.text();
++					}
++					return '';
++				} catch (error) {
++					return '';
++				}
++			}
++		);
++
 +		CommandsRegistry.registerCommand(SagemakerServerClient.LOGOUT_COMMAND_ID, () => {
 +			const currentUrl = new URL(window.location.href);
 +			const hostname = currentUrl.hostname;
@@ -65,6 +96,57 @@ Index: sagemaker-code-editor/vscode/src/vs/workbench/browser/client.ts
 +	}
 +}
 \ No newline at end of file
+Index: sagemaker-code-editor/vscode/src/vs/server/node/remoteExtensionHostAgentServer.ts
+===================================================================
+--- sagemaker-code-editor.orig/vscode/src/vs/server/node/remoteExtensionHostAgentServer.ts
++++ sagemaker-code-editor/vscode/src/vs/server/node/remoteExtensionHostAgentServer.ts
+@@ -168,6 +168,46 @@
+			}
+			return serveFile(filePath, CacheControl.ETAG, this._logService, req, res, responseHeaders);
+		}
++
++		if (pathname === '/load-cookies') {
++			let entries = parsedUrl.query['entry'] ?? [];
++			let workspace = (Array.isArray(parsedUrl.query['workspace']) ? parsedUrl.query['workspace'][0] : parsedUrl.query['workspace']) ?? '/';
++
++			if (!Array.isArray(entries)) {
++				entries = [entries];
++			}
++
++			const requestCookie = req.headers.cookie;
++			if (!requestCookie) {
++				res.writeHead(200, { 'Content-Type' : 'text/plain' });
++				return void res.end('');
++			}
++
++			const parsedCookies = requestCookie.split(';').reduce((result, cookie) => {
++				const [name, ...rest] = cookie.split('=');
++				const value = rest.join('=').trim();
++				const key = name.trim();
++				if (entries.includes(key)) {
++					result[key] = decodeURIComponent(value);
++				}
++				return result;
++			}, {} as Record<string, string>);
++
++			const dirPath = join(workspace, '.aws/sso');
++			const filePath = join(dirPath, '/cookies.json');
++			try {
++				if (!fs.existsSync(dirPath)) {
++					fs.mkdirSync(dirPath, { recursive: true });
++				}
++				fs.writeFileSync(filePath,  JSON.stringify(parsedCookies));	
++			} catch {
++				res.writeHead(200, { 'Content-Type' : 'text/plain' });
++				return void res.end('');			
++			}
++
++			res.writeHead(200, { 'Content-Type' : 'text/plain' });
++			return void res.end(filePath);
++		}
+
+		// workbench web UI
+		if (this._webClientServer) {
 Index: sagemaker-code-editor/vscode/src/vs/workbench/browser/web.main.ts
 ===================================================================
 --- sagemaker-code-editor.orig/vscode/src/vs/workbench/browser/web.main.ts