Skip to content
Closed
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
50 changes: 49 additions & 1 deletion workspaces/ballerina/ballerina-extension/src/stateMachine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,13 @@ import { StateMachinePopup } from './stateMachinePopup';
import { checkIsBallerinaPackage, checkIsBI, fetchScope, getOrgPackageName, UndoRedoManager, getProjectTomlValues } from './utils';
import { buildProjectArtifactsStructure } from './utils/project-artifacts';
import { activateDevantFeatures } from './features/devant/activator';
import {
isCloudEditorEnvironment,
getCurrentAccessToken,
shouldUpdateToken,
updateBallerinaSettingsWithStsToken,
getDevantStsToken
} from './utils/auth-utils';

interface MachineContext extends VisualizerLocation {
langClient: ExtendedLangClient | null;
Expand All @@ -27,6 +34,34 @@ export let history: History;
export let undoRedoManager: IUndoRedoManager;
let pendingProjectRootUpdateResolvers: Array<() => void> = [];

/**
* Helper function to update STS token when packages are being pulled
* This ensures authenticated access to private packages in cloud editor environments
*/
async function updateStsTokenForPackagePull(): Promise<void> {
// Only update in cloud editor environment
if (!isCloudEditorEnvironment()) {
return;
}

try {
// Get the STS token from platform extension
const stsToken = await getDevantStsToken();

if (stsToken && stsToken.trim() !== "") {
const currentToken = await getCurrentAccessToken();

// Only update if token needs updating
if (shouldUpdateToken(currentToken, stsToken)) {
await updateBallerinaSettingsWithStsToken(stsToken);
}
}
} catch (error) {
console.error('Failed to update STS token for package pull:', error);
// Don't throw - allow package pull to proceed even if token update fails
}
}

const stateMachine = createMachine<MachineContext>(
{
/** @xstate-layout N4IgpgJg5mDOIC5QDUCWsCuBDANqgXmAE4DEASgKIDKFAKgPq0Dy9FAGrRQHJUCSTXepQCCAEQCaAbQAMAXUSgADgHtYqAC6plAOwUgAnogC0ARgDsAVgB0AZgvSHANhPS7Ji44BMAGhAAPRHsAFgBfEN80TFwCYitUbQ1UaMISCB0wOO0AN2UAawzI7DxCIkzE5LAEeJyAYyxNHRlZJr0VNQbdJH9ETyCgqwAOO2legE4gxwtLTxtfQwQTT1HpW1G18aCLAaCZsIj0IpjS+PLisBJiImVSxRx6gDNrgFsrQorjhM0KquzlOo6mi0um1Ejo9AEEL1+kN7GMJlMLDM5sYTHYrI5LA4bNILOYpgM9iA3mdSlgaposvUwAAZKipdKZHL5V4Hd5WMkUqm0n61epabSAuStVSgzqgCF2UZWIJmTyIkwDcYDcxBZELGzjaVjUYWUbbbEmUaE4lHdnk1CU9Q0umXa5WW4PZ4sqIks2cq3c6p-PmNORApQijrgxCS6Wy+WKoLKsyqgyIExBPXSxNrSwuBXSAnhIms104WAUIhXUiUWhkKRC4GB-nBhYDTxWQ2LMwKpUqnxx9WanYp3X66SG42501gPxW7RqHRkMBYCD6EhMAAK3HoyF4FAA6v6QCCg10Ib0BtLhjYXPZpMsBo41QmLNZPAMY2tXI4caMzEOXaasqgwAB3YRzSyDIf3-XhPnpbQMi9ZkTRKKxQIAoCQN-P9wI0HlvQBP1KwDdoay6eZz2ldYBgsbFT0TCxfAlcwrA8R8PBcNwPE-Q54MQwCKQyP8wAAIzQf9qWUWd4igSDoN+WDhw41CuItHj+MEv9hNE7QoEw-5+UFeQq3wsF9x6bZ6McRjnFcXEPDVTYGxTPVyNcBMdTYtlOOQqxeIE1DVIgSAJMZPIChk2I3O4jylO8kTfIgTSfQFHDdLw0VayWBshkxUZ3ERMwbD6azpH6OyyIopyLBc11QoUhDUOnWd5yXFc103bddwI8UelPVYbEsbFcUsMwBjVO8j1y7q7BKqjyu-OT3MQ2q5xIJqN3oABVRdRGETgWurAz2shTqNW6hy+vxNVHEmWxHyjPUBmkAavCm2T-3k4Dqv-eb5wAMV4akKFYUReFobb9LFbp9pMLqepxPEBrVMxZWPRY+k2bZdmzOCQpmsLEIoCBEnUkgKABhhRAEChgeSwz9scKwzFGI7NnhxwbAxDt5kfKUZTsl83w-QltGUXz4C6DGiGFEHayMUZr07IxPGZqwRnPZYJnfPoyvR4KPlOGJxcpvbTClSZegGs8FTlDEhqjEi1mKxzJs1r94I5C0uSoPW9wNw10URGVlWh+sPDMOHsUbFMr3rZYH1CR32NifNC2LD22rBw2fZN-3cUDy3O0yiG6dtmFEzumxHtiUdx0nbQPuT3bU+942-bN7Pg87FwJisTwn2WRwgm6uU+f2J3Mee5Da9BiEjCjNVTxp+sUxsB8KMHnNh9KSrXsQ9D1HH2tpCGhwbfsibnNj1ysaqzzlJ8sTd6p-fO1MhsnL1J+o+2Mv14v16r8i2dIDvntYyUZHymRsGRFszhZidgGLAzuSMBo9nsCzT+b0kLYxqjOOcgCwYP3mO3awfRbY3TuleTwqCN4oX-LjfGUAcEQjwfGcwUpRjak8CYRw0sCqODCGEIAA */
Expand Down Expand Up @@ -63,6 +98,8 @@ const stateMachine = createMachine<MachineContext>(
projectUri: (context, event) => event.projectPath
}),
async (context, event) => {
// Update STS token before building project structure (for package pulls)
await updateStsTokenForPackagePull();
await buildProjectArtifactsStructure(event.projectPath, StateMachine.langClient(), true);
notifyCurrentWebview();
// Resolve the next pending promise waiting for project root update completion
Expand Down Expand Up @@ -100,6 +137,8 @@ const stateMachine = createMachine<MachineContext>(
package: (context, event) => event.data.packageName,
}),
async (context, event) => {
// Update STS token before building project structure (for package pulls)
await updateStsTokenForPackagePull();
await buildProjectArtifactsStructure(event.data.projectPath, StateMachine.langClient(), true);
notifyCurrentWebview();
}
Expand Down Expand Up @@ -344,6 +383,9 @@ const stateMachine = createMachine<MachineContext>(
// If the project uri is not set, we don't need to build the project structure
if (context.projectUri) {

// Update STS token before building project structure (for package pulls)
await updateStsTokenForPackagePull();

// Add a 2 second delay before registering artifacts
await new Promise(resolve => setTimeout(resolve, 1000));
// Initial Project Structure
Expand Down Expand Up @@ -406,6 +448,9 @@ const stateMachine = createMachine<MachineContext>(
return;
}

// Update STS token before running bal build (for package pulls)
await updateStsTokenForPackagePull();

const taskDefinition: TaskDefinition = {
type: 'shell',
task: 'run'
Expand Down Expand Up @@ -740,7 +785,10 @@ export function updateView(refreshTreeView?: boolean, projectUri?: string) {

stateService.send({ type: "VIEW_UPDATE", viewLocation: lastView ? newLocation : { view: "Overview" } });
if (refreshTreeView) {
buildProjectArtifactsStructure(projectUri || StateMachine.context().projectUri, StateMachine.langClient(), true);
// Update STS token before building project structure (for package pulls)
updateStsTokenForPackagePull().then(() => {
buildProjectArtifactsStructure(projectUri || StateMachine.context().projectUri, StateMachine.langClient(), true);
});
}
notifyCurrentWebview();
}
Expand Down
186 changes: 186 additions & 0 deletions workspaces/ballerina/ballerina-extension/src/utils/auth-utils.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,186 @@
/**
* Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com) All Rights Reserved.
*
* WSO2 LLC. licenses this file to you under the Apache License,
* Version 2.0 (the "License"); you may not use this file except
* in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

import * as fs from 'fs';
import * as path from 'path';
import { extensions } from 'vscode';
import { IWso2PlatformExtensionAPI } from "@wso2/wso2-platform-core";
import { extension } from '../BalExtensionContext';

/**
* Checks if we're running in a cloud editor environment
* @returns boolean - true if CLOUD_STS_TOKEN environment variable exists
*/
export function isCloudEditorEnvironment(): boolean {
return !!process.env.CLOUD_STS_TOKEN;
}

/**
* Reads the current access token from Settings.toml
* @returns Promise<string | undefined> - The current access token or undefined if not found
*/
export async function getCurrentAccessToken(): Promise<string | undefined> {
try {
const settingsPath = path.join(extension.ballerinaExtInstance.getBallerinaUserHome(), 'Settings.toml');

if (!fs.existsSync(settingsPath)) {
return undefined;
}

const content = fs.readFileSync(settingsPath, 'utf8');
const tokenMatch = content.match(/accesstoken\s*=\s*"([^"]*)"/);

return tokenMatch ? tokenMatch[1] : undefined;
} catch (error) {
console.error('Failed to read current access token:', error);
return undefined;
}
}

/**
* Checks if a JWT token is expired or about to expire (within 5 minutes)
* @param token - The JWT token to check
* @returns boolean - true if token is expired or about to expire
*/
export function isTokenExpiredOrExpiring(token: string): boolean {
try {
// JWT tokens have 3 parts separated by dots
const parts = token.split('.');
if (parts.length !== 3) {
return true; // Not a valid JWT format
}

// Decode the payload (second part)
const payload = JSON.parse(atob(parts[1]));

if (!payload.exp) {
return true; // No expiration time means we should refresh
}

// Check if token expires within 5 minutes (300 seconds)
const currentTime = Math.floor(Date.now() / 1000);
const expirationTime = payload.exp;
const bufferTime = 300; // 5 minutes buffer

return (expirationTime - currentTime) <= bufferTime;
} catch (error) {
console.warn('Failed to decode token, treating as expired:', error);
return true; // If we can't decode it, treat as expired
}
}

/**
* Checks if the token needs to be updated based on current vs new token
* @param currentToken - The token currently in Settings.toml
* @param newToken - The new STS token from platform extension
* @returns boolean - true if token needs updating
*/
export function shouldUpdateToken(currentToken: string | undefined, newToken: string): boolean {
// No current token exists
if (!currentToken) {
return true;
}

// Tokens are different
if (currentToken !== newToken) {
return true;
}

// Current token is expired or expiring soon
if (isTokenExpiredOrExpiring(currentToken)) {
return true;
}

return false;
}

/**
* Updates the Ballerina Settings.toml file with the provided STS token
* @param stsToken - The STS token to set as accesstoken
*/
export async function updateBallerinaSettingsWithStsToken(stsToken: string): Promise<void> {
try {
const settingsPath = path.join(extension.ballerinaExtInstance.getBallerinaUserHome(), 'Settings.toml');

if (!fs.existsSync(settingsPath)) {
// Create the .ballerina directory if it doesn't exist
const ballerinaDir = path.dirname(settingsPath);
if (!fs.existsSync(ballerinaDir)) {
fs.mkdirSync(ballerinaDir, { recursive: true });
}

// Create new Settings.toml with STS token
const newContent = `[central]\naccesstoken="${stsToken}"\n`;
fs.writeFileSync(settingsPath, newContent, 'utf8');
return;
}

// Read existing Settings.toml
let content = fs.readFileSync(settingsPath, 'utf8');

// Check if [central] section exists
if (content.includes('[central]')) {
// Replace existing accesstoken
content = content.replace(
/accesstoken\s*=\s*"[^"]*"/g,
`accesstoken="${stsToken}"`
);
} else {
// Add [central] section with accesstoken
content += `\n[central]\naccesstoken="${stsToken}"\n`;
}

// Write back to file
fs.writeFileSync(settingsPath, content, 'utf8');
console.log('Updated Settings.toml with STS token');

} catch (error) {
console.error('Failed to update Settings.toml with STS token:', error);
throw error;
}
}

/**
* Retrieves the STS token from the platform extension for authenticated requests
* @returns Promise<string | undefined> - The STS token if available, undefined otherwise
*/
export async function getDevantStsToken(): Promise<string | undefined> {
try {
// Try to get STS token from platform extension
const platformExt = extensions.getExtension("wso2.wso2-platform");
if (!platformExt) {
return undefined;
}

// Check if extension is already active before activating
if (!platformExt.isActive) {
await platformExt.activate();
}

const platformExtAPI: IWso2PlatformExtensionAPI = platformExt.exports;

const platformStsToken = await platformExtAPI.getStsToken();
if (platformStsToken && platformStsToken.trim() !== "") {
return platformStsToken;
}
} catch (error) {
console.error("Failed to get STS token from platform extension:", error);
}

return undefined;
}
Loading