Skip to content

Commit

Permalink
refactor: Mutex class & use in AuthHandler; TSDoc
Browse files Browse the repository at this point in the history
- Adds a new `Mutex` class that wraps around
a `DeferredPromise`. See the `@brief` for details.
- Refactors the `AuthHandler` class to use a map
of `Mutex` for each profile used in critical
sections.
- Added doc: `AuthHandler, DeferredPromise, Mutex`

Signed-off-by: Trae Yelovich <[email protected]>
  • Loading branch information
traeok committed Dec 20, 2024
1 parent 0dce568 commit 167977a
Show file tree
Hide file tree
Showing 3 changed files with 199 additions and 55 deletions.
146 changes: 94 additions & 52 deletions packages/zowe-explorer-api/src/profiles/AuthHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,28 +10,29 @@
*/

import { Gui } from "../globals";
import { CorrelatedError, DeferredPromise, reloadWorkspacesForProfile } from "../utils";
import { CorrelatedError, refreshActiveEditorForProfile, refreshWorkspacesForProfile } from "../utils";
import * as imperative from "@zowe/imperative";
import { IZoweDatasetTreeNode, IZoweJobTreeNode, IZoweTree, IZoweTreeNode, IZoweUSSTreeNode } from "../tree";
import { IZoweTree, IZoweTreeNode } from "../tree";
import { commands } from "vscode";
import { Mutex } from "../utils/Mutex";

export interface IAuthMethods {
/* Method for establishing SSO login with a given profile name */
ssoLogin: (node?: IZoweTreeNode, profileName?: string) => PromiseLike<boolean>;
/* Method that prompts the user for credentials, sets them on the profile and returns them to the caller if set */
promptCredentials: (profile: string | imperative.IProfileLoaded, rePrompt?: boolean) => PromiseLike<string[]>;
}

export interface AuthPromptOpts extends IAuthMethods {
export interface AuthPromptParams extends IAuthMethods {
// Whether the profile is using token-based authentication
isUsingTokenAuth?: boolean;
// API-friendly error correlation for the "Invalid Credentials" scenario
errorCorrelation?: CorrelatedError;
}

export interface LockProfileOpts extends AuthPromptOpts {
waitAfterLock?: boolean;
}

type ProfileLike = string | imperative.IProfileLoaded;
export class AuthHandler {
private static lockedProfiles: Map<string, DeferredPromise<void>> = new Map();
private static lockedProfiles: Map<string, Mutex> = new Map();

/**
* Function that checks whether a profile is using token based authentication
Expand All @@ -47,41 +48,59 @@ export class AuthHandler {
return baseSecureProfileProps?.includes("tokenValue") && !profileUsesBasicAuth;
}

public static unlockProfile(profile: ProfileLike): void {
/**
* Unlocks the given profile so it can be used again.
* @param profile {ProfileLike} The profile (name or {@link imperative.IProfileLoaded} object) to unlock
* @param refreshResources {boolean} Whether to refresh high-priority resources (active editor & virtual workspace) after unlocking
*/
public static unlockProfile(profile: ProfileLike, refreshResources?: boolean): void {
const profileName = typeof profile === "string" ? profile : profile.name;
const deferred = this.lockedProfiles.get(profileName);
if (deferred) {
deferred.resolve();
this.lockedProfiles.delete(profileName);
// reload virtual workspaces for the profile now that its usable
reloadWorkspacesForProfile(profileName);
deferred.unlock();
if (refreshResources) {
// refresh an active, unsaved editor if it contains the profile
refreshActiveEditorForProfile(profileName);
// refresh virtual workspaces for the profile
refreshWorkspacesForProfile(profileName);
}
}
}

/**
* Prompts the user to authenticate over SSO or a credential prompt in the event of an error.
* @param imperativeError The authentication error that was encountered
* @param profile The profile to authenticate
* @param opts {AuthPromptParams} Prompt parameters (login methods, using token auth, error correlation)
* @returns {boolean} Whether authentication was successful
*/
public static async promptForAuthentication(
imperativeError: imperative.ImperativeError,
profile: ProfileLike,
opts: AuthPromptOpts
opts: AuthPromptParams
): Promise<boolean> {
const profileName = typeof profile === "string" ? profile : profile.name;
if (imperativeError.mDetails.additionalDetails) {
const tokenError: string = imperativeError.mDetails.additionalDetails;
if (tokenError.includes("Token is not valid or expired.") || opts.isUsingTokenAuth) {
// Handle token-based authentication error through the given `ssoLogin` method.
const message = "Log in to Authentication Service";
const userResp = await Gui.showMessage(opts.errorCorrelation?.message ?? imperativeError.message, {
items: [message],
vsCodeOpts: { modal: true },
});
if (userResp === message) {
if (await opts.ssoLogin(null, profileName)) {
// SSO login was successful, unlock profile
AuthHandler.unlockProfile(profileName);
// SSO login was successful, unlock profile so it can be used again
AuthHandler.unlockProfile(profileName, true);
return true;
}
}
return false;
}
}

// Prompt the user to update their credentials using the given `promptCredentials` method.
const checkCredsButton = "Update Credentials";
const creds = await Gui.errorMessage(opts.errorCorrelation?.message ?? imperativeError.message, {
items: [checkCredsButton],
Expand All @@ -94,60 +113,82 @@ export class AuthHandler {
});

if (creds != null) {
// New creds were provided
// Propagate new creds to other profiles
// New creds were set, directly propagate new profile to other tree providers.

// TODO: If we can access extender tree providers (e.g. CICS), it would help to propagate profile updates here.
// For now we will propagate profile changes to core providers (Data Sets, USS, Jobs)
const treeProviders = (await commands.executeCommand("zowe.getTreeProviders")) as any;
const dsNode = (await (treeProviders.ds as IZoweTree<IZoweDatasetTreeNode>).getChildren()).find((n) => n.label === profileName);
if (dsNode && typeof profile !== "string") {
dsNode.setProfileToChoice(profile);
}
const ussNode = (await (treeProviders.uss as IZoweTree<IZoweUSSTreeNode>).getChildren()).find((n) => n.label === profileName);
if (ussNode && typeof profile !== "string") {
ussNode.setProfileToChoice(profile);
}
const jobsNode = (await (treeProviders.job as IZoweTree<IZoweJobTreeNode>).getChildren()).find((n) => n.label === profileName);
if (jobsNode && typeof profile !== "string") {
jobsNode.setProfileToChoice(profile);
for (const provider of [treeProviders.ds, treeProviders.uss, treeProviders.job]) {
const node = (await (provider as IZoweTree<IZoweTreeNode>).getChildren()).find((n) => n.label === profileName);
if (node && typeof profile !== "string") {
node.setProfileToChoice(profile);
}
}
AuthHandler.unlockProfile(profileName);
// Unlock profile so it can be used again
AuthHandler.unlockProfile(profileName, true);
return true;
}
return false;
}

public static async lockProfile(profile: ProfileLike, imperativeError?: imperative.ImperativeError, opts?: LockProfileOpts): Promise<void> {
/**
* Locks the given profile to prevent further use in asynchronous operations (at least where the lock is respected).
* Supports prompting for authentication if an Imperative error and prompt options are given.
* @param profile The profile to lock
* @param imperativeError The Imperative error that was encountered when using the profile
* @param opts Prompt parameters to use during authentication
* @returns Whether the profile was successfully locked
*/
public static async lockProfile(profile: ProfileLike, imperativeError?: imperative.ImperativeError, opts?: AuthPromptParams): Promise<boolean> {
const profileName = typeof profile === "string" ? profile : profile.name;
if (this.lockedProfiles.has(profileName)) {
return this.lockedProfiles.get(profileName)!.promise;

// If the mutex does not exist, make one for the profile and acquire the lock
if (!this.lockedProfiles.has(profileName)) {
this.lockedProfiles.set(profileName, new Mutex());
}
const deferred = new DeferredPromise<void>();
this.lockedProfiles.set(profileName, deferred);

// Prompt the user to re-authenticate
// Attempt to acquire the lock
const mutex = this.lockedProfiles.get(profileName);
await mutex!.lock();

// Prompt the user to re-authenticate if an error and options were provided
if (imperativeError && opts) {
const credsEntered = await AuthHandler.promptForAuthentication(imperativeError, profile, opts);

// If the user failed to re-authenticate, reject the promise
// TODO: more manual testing
if (!credsEntered) {
deferred.reject();
mutex!.unlock();
}
// Return `true` as the mutex was still locked successfully.
return true;
}

if (opts?.waitAfterLock) {
return deferred.promise;
}
return true;
}

/**
* Checks whether the given profile has its lock acquired.
* @param profile The profile to check
* @returns {boolean} `true` if the given profile is locked, `false` otherwise
*/
public static isLocked(profile: ProfileLike): boolean {
return this.lockedProfiles.has(typeof profile === "string" ? profile : profile.name);
const mutex = this.lockedProfiles.get(typeof profile === "string" ? profile : profile.name);
if (mutex == null) {
return false;
}

return mutex.locked;
}

public static async waitIfLocked(profile: ProfileLike): Promise<void> {
const deferred = this.lockedProfiles.get(typeof profile === "string" ? profile : profile.name);
if (deferred) {
await deferred.promise;
/**
* Non-blocking operation that attempts to acquire the lock for a profile if it can be acquired.
* @param profile The profile to acquire the lock for
* @returns {boolean} `true` if the lock was acquired, `false` otherwise
*/
public static tryToLockProfile(profile: ProfileLike): boolean {
const mutex = this.lockedProfiles.get(typeof profile === "string" ? profile : profile.name);
if (mutex) {
return mutex.tryLock();
}
return true;
}
}

Expand All @@ -157,10 +198,11 @@ export function withCredentialManagement<T extends (...args: any[]) => any | Pro
apiMethod: T
): T {
return async function (...args: any[]) {
await AuthHandler.waitIfLocked(profile);

await AuthHandler.lockProfile(profile);
try {
return await apiMethod(...args);
const res = await apiMethod(...args);
AuthHandler.unlockProfile(profile);
return res;
} catch (error) {
if (error instanceof imperative.ImperativeError) {
const imperativeError: imperative.ImperativeError = error as imperative.ImperativeError;
Expand All @@ -169,7 +211,7 @@ export function withCredentialManagement<T extends (...args: any[]) => any | Pro
httpErrorCode === imperative.RestConstants.HTTP_STATUS_401 ||
imperativeError.message.includes("All configured authentication methods failed")
) {
await AuthHandler.lockProfile(profile, imperativeError, { ...authMethods });
await AuthHandler.promptForAuthentication(imperativeError, profile, { ...authMethods });
return await apiMethod(...args);
} else {
throw error;
Expand Down
32 changes: 29 additions & 3 deletions packages/zowe-explorer-api/src/utils/DeferredPromise.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,41 @@
*
*/

/* Status of the deferred promise */
export type DeferredPromiseStatus = "pending" | "fulfilled" | "rejected";

/**
* @brief Externally control the resolution and rejection of a promise.
*
* @details
* Creates a promise with accessible `resolve` and `reject` methods, enabling external entities to
* settle the promise based on custom logic or asynchronous events. This is particularly useful when
* the promise's outcome depends on factors outside the immediate context.
*/
export class DeferredPromise<T> {
private mStatus: DeferredPromiseStatus = "pending";

public promise: Promise<T>;
public resolve: (value: T | PromiseLike<T>) => void;
public reject: (reason?: any) => void;

public constructor() {
this.promise = new Promise<T>((resolve, reject) => {
this.resolve = resolve;
this.reject = reject;
this.resolve = (value) => {
this.mStatus = "fulfilled";
resolve(value);
};
this.reject = (err) => {
this.mStatus = "rejected";
reject(err);
};
});
}
}

/**
* @returns {PromiseStatus} The status of the deferred promise
*/
public get status(): DeferredPromiseStatus {
return this.mStatus;
}
}
76 changes: 76 additions & 0 deletions packages/zowe-explorer-api/src/utils/Mutex.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
/**
* This program and the accompanying materials are made available under the terms of the
* Eclipse Public License v2.0 which accompanies this distribution, and is available at
* https://www.eclipse.org/legal/epl-v20.html
*
* SPDX-License-Identifier: EPL-2.0
*
* Copyright Contributors to the Zowe Project.
*
*/

import { DeferredPromise } from "./DeferredPromise";

/**
* @brief
* A class for managing mutual exclusion in asynchronous operations,
* ensuring that only one asynchronous function can access a critical section at a time.
*
* @details
* The lock state of the mutex is determined by the presence of a pending promise in the `mDeferred` property (deferred promise).
* If the lock is available, `mDeferred` is `null`.
*/
export class Mutex {
private mDeferred: DeferredPromise<void> | null = null;

/**
* Attempt to lock the mutex if possible.
* @returns Whether the mutex was successfully locked
*/
public tryLock(): boolean {
if (this.mDeferred === null) {
this.mDeferred = new DeferredPromise<void>();
return true;
}

return false;
}

/**
* Waits for the mutex to be available and immediately locks it thereafter.
*/
public async lock(): Promise<void> {
while (this.mDeferred !== null) {
await this.mDeferred.promise;
}
this.mDeferred = new DeferredPromise<void>();
}

/**
* Unlocks the mutex, resolving or rejecting the promise depending on whether an error is provided.
* @param error {string} An optional error message. When provided, the deferred promise is rejected - otherwise it is resolved.
*/
public unlock(error?: string): void {
if (this.mDeferred !== null) {
if (error) {
this.mDeferred.reject(error);
} else {
this.mDeferred.resolve();
}
this.mDeferred = null;
} else {
throw new Error("Cannot unlock a lock that hasn't been acquired.");
}
}

/**
* @returns {boolean} whether the mutex is locked
*/
public get locked(): boolean {
if (this.mDeferred === null) {
return false;
}

return true;
}
}

0 comments on commit 167977a

Please sign in to comment.