Skip to content

Commit

Permalink
[PM-16837] Fix agent only loading when featureflag is on during start…
Browse files Browse the repository at this point in the history
…up (#12742)

* Fix ssh generation and import not being available when agent feature-flag is disabled

* Fix agent only loading when featureflag is on during startup
  • Loading branch information
quexten authored Jan 8, 2025
1 parent 1b08335 commit 244539c
Show file tree
Hide file tree
Showing 4 changed files with 172 additions and 155 deletions.
9 changes: 2 additions & 7 deletions apps/desktop/src/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
// @ts-strict-ignore
import * as path from "path";

import { app, ipcMain } from "electron";
import { app } from "electron";
import { Subject, firstValueFrom } from "rxjs";

import { AccountServiceImplementation } from "@bitwarden/common/auth/services/account.service";
Expand Down Expand Up @@ -257,12 +257,7 @@ export class Main {
this.clipboardMain = new ClipboardMain();
this.clipboardMain.init();

ipcMain.handle("sshagent.init", async (event: any, message: any) => {
if (this.sshAgentService == null) {
this.sshAgentService = new MainSshAgentService(this.logService, this.messagingService);
this.sshAgentService.init();
}
});
this.sshAgentService = new MainSshAgentService(this.logService, this.messagingService);

new EphemeralValueStorageService();
new SSOLocalhostCallbackService(this.environmentService, this.messagingService);
Expand Down
8 changes: 8 additions & 0 deletions apps/desktop/src/platform/main/main-ssh-agent.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,14 @@ export class MainSshAgentService {
return sshagent.importKey(privateKey, password);
},
);

ipcMain.handle("sshagent.init", async (event: any, message: any) => {
this.init();
});

ipcMain.handle("sshagent.isloaded", async (event: any) => {
return this.agentState != null;
});
}

init() {
Expand Down
3 changes: 3 additions & 0 deletions apps/desktop/src/platform/preload.ts
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,9 @@ const sshAgent = {
});
return res;
},
isLoaded(): Promise<boolean> {
return ipcRenderer.invoke("sshagent.isloaded");
},
};

const powermonitor = {
Expand Down
307 changes: 159 additions & 148 deletions apps/desktop/src/platform/services/ssh-agent.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,153 +61,87 @@ export class SshAgentService implements OnDestroy {
) {}

async init() {
const isSshAgentFeatureEnabled = await this.configService.getFeatureFlag(FeatureFlag.SSHAgent);
if (isSshAgentFeatureEnabled) {
await ipc.platform.sshAgent.init();

this.messageListener
.messages$(new CommandDefinition("sshagent.signrequest"))
.pipe(
withLatestFrom(this.authService.activeAccountStatus$),
// This switchMap handles unlocking the vault if it is locked:
// - If the vault is locked, we will wait for it to be unlocked.
// - If the vault is not unlocked within the timeout, we will abort the flow.
// - If the vault is unlocked, we will continue with the flow.
// switchMap is used here to prevent multiple requests from being processed at the same time,
// and will cancel the previous request if a new one is received.
switchMap(([message, status]) => {
if (status !== AuthenticationStatus.Unlocked) {
ipc.platform.focusWindow();
this.toastService.showToast({
variant: "info",
title: null,
message: this.i18nService.t("sshAgentUnlockRequired"),
});
return this.authService.activeAccountStatus$.pipe(
filter((status) => status === AuthenticationStatus.Unlocked),
timeout({
first: this.SSH_VAULT_UNLOCK_REQUEST_TIMEOUT,
}),
catchError((error: unknown) => {
if (error instanceof TimeoutError) {
this.toastService.showToast({
variant: "error",
title: null,
message: this.i18nService.t("sshAgentUnlockTimeout"),
});
const requestId = message.requestId as number;
// Abort flow by sending a false response.
// Returning an empty observable this will prevent the rest of the flow from executing
return from(ipc.platform.sshAgent.signRequestResponse(requestId, false)).pipe(
map(() => EMPTY),
);
}

throw error;
}),
map(() => message),
);
}

return of(message);
}),
// This switchMap handles fetching the ciphers from the vault.
switchMap((message) =>
from(this.cipherService.getAllDecrypted()).pipe(
map((ciphers) => [message, ciphers] as const),
),
),
// This concatMap handles showing the dialog to approve the request.
concatMap(async ([message, ciphers]) => {
const cipherId = message.cipherId as string;
const isListRequest = message.isListRequest as boolean;
const requestId = message.requestId as number;
let application = message.processName as string;
if (application == "") {
application = this.i18nService.t("unknownApplication");
}

if (isListRequest) {
const sshCiphers = ciphers.filter(
(cipher) => cipher.type === CipherType.SshKey && !cipher.isDeleted,
);
const keys = sshCiphers.map((cipher) => {
return {
name: cipher.name,
privateKey: cipher.sshKey.privateKey,
cipherId: cipher.id,
};
});
await ipc.platform.sshAgent.setKeys(keys);
await ipc.platform.sshAgent.signRequestResponse(requestId, true);
return;
}

if (ciphers === undefined) {
ipc.platform.sshAgent
.signRequestResponse(requestId, false)
.catch((e) => this.logService.error("Failed to respond to SSH request", e));
}

const cipher = ciphers.find((cipher) => cipher.id == cipherId);
this.configService
.getFeatureFlag$(FeatureFlag.SSHAgent)
.pipe(
concatMap(async (enabled) => {
if (enabled && !(await ipc.platform.sshAgent.isLoaded())) {
return this.initSshAgent();
}
}),
takeUntil(this.destroy$),
)
.subscribe();
}

private async initSshAgent() {
await ipc.platform.sshAgent.init();

this.messageListener
.messages$(new CommandDefinition("sshagent.signrequest"))
.pipe(
withLatestFrom(this.authService.activeAccountStatus$),
// This switchMap handles unlocking the vault if it is locked:
// - If the vault is locked, we will wait for it to be unlocked.
// - If the vault is not unlocked within the timeout, we will abort the flow.
// - If the vault is unlocked, we will continue with the flow.
// switchMap is used here to prevent multiple requests from being processed at the same time,
// and will cancel the previous request if a new one is received.
switchMap(([message, status]) => {
if (status !== AuthenticationStatus.Unlocked) {
ipc.platform.focusWindow();
const dialogRef = ApproveSshRequestComponent.open(
this.dialogService,
cipher.name,
application,
this.toastService.showToast({
variant: "info",
title: null,
message: this.i18nService.t("sshAgentUnlockRequired"),
});
return this.authService.activeAccountStatus$.pipe(
filter((status) => status === AuthenticationStatus.Unlocked),
timeout({
first: this.SSH_VAULT_UNLOCK_REQUEST_TIMEOUT,
}),
catchError((error: unknown) => {
if (error instanceof TimeoutError) {
this.toastService.showToast({
variant: "error",
title: null,
message: this.i18nService.t("sshAgentUnlockTimeout"),
});
const requestId = message.requestId as number;
// Abort flow by sending a false response.
// Returning an empty observable this will prevent the rest of the flow from executing
return from(ipc.platform.sshAgent.signRequestResponse(requestId, false)).pipe(
map(() => EMPTY),
);
}

throw error;
}),
map(() => message),
);

const result = await firstValueFrom(dialogRef.closed);
return ipc.platform.sshAgent.signRequestResponse(requestId, result);
}),
takeUntil(this.destroy$),
)
.subscribe();

this.accountService.activeAccount$.pipe(skip(1), takeUntil(this.destroy$)).subscribe({
next: (account) => {
this.logService.info("Active account changed, clearing SSH keys");
ipc.platform.sshAgent
.clearKeys()
.catch((e) => this.logService.error("Failed to clear SSH keys", e));
},
error: (e: unknown) => {
this.logService.error("Error in active account observable", e);
ipc.platform.sshAgent
.clearKeys()
.catch((e) => this.logService.error("Failed to clear SSH keys", e));
},
complete: () => {
this.logService.info("Active account observable completed, clearing SSH keys");
ipc.platform.sshAgent
.clearKeys()
.catch((e) => this.logService.error("Failed to clear SSH keys", e));
},
});

combineLatest([
timer(0, this.SSH_REFRESH_INTERVAL),
this.desktopSettingsService.sshAgentEnabled$,
])
.pipe(
concatMap(async ([, enabled]) => {
if (!enabled) {
await ipc.platform.sshAgent.clearKeys();
return;
}

const ciphers = await this.cipherService.getAllDecrypted();
if (ciphers == null) {
await ipc.platform.sshAgent.lock();
return;
}

}

return of(message);
}),
// This switchMap handles fetching the ciphers from the vault.
switchMap((message) =>
from(this.cipherService.getAllDecrypted()).pipe(
map((ciphers) => [message, ciphers] as const),
),
),
// This concatMap handles showing the dialog to approve the request.
concatMap(async ([message, ciphers]) => {
const cipherId = message.cipherId as string;
const isListRequest = message.isListRequest as boolean;
const requestId = message.requestId as number;
let application = message.processName as string;
if (application == "") {
application = this.i18nService.t("unknownApplication");
}

if (isListRequest) {
const sshCiphers = ciphers.filter(
(cipher) =>
cipher.type === CipherType.SshKey &&
!cipher.isDeleted &&
cipher.organizationId === null,
(cipher) => cipher.type === CipherType.SshKey && !cipher.isDeleted,
);
const keys = sshCiphers.map((cipher) => {
return {
Expand All @@ -217,11 +151,88 @@ export class SshAgentService implements OnDestroy {
};
});
await ipc.platform.sshAgent.setKeys(keys);
}),
takeUntil(this.destroy$),
)
.subscribe();
}
await ipc.platform.sshAgent.signRequestResponse(requestId, true);
return;
}

if (ciphers === undefined) {
ipc.platform.sshAgent
.signRequestResponse(requestId, false)
.catch((e) => this.logService.error("Failed to respond to SSH request", e));
}

const cipher = ciphers.find((cipher) => cipher.id == cipherId);

ipc.platform.focusWindow();
const dialogRef = ApproveSshRequestComponent.open(
this.dialogService,
cipher.name,
application,
);

const result = await firstValueFrom(dialogRef.closed);
return ipc.platform.sshAgent.signRequestResponse(requestId, result);
}),
takeUntil(this.destroy$),
)
.subscribe();

this.accountService.activeAccount$.pipe(skip(1), takeUntil(this.destroy$)).subscribe({
next: (account) => {
this.logService.info("Active account changed, clearing SSH keys");
ipc.platform.sshAgent
.clearKeys()
.catch((e) => this.logService.error("Failed to clear SSH keys", e));
},
error: (e: unknown) => {
this.logService.error("Error in active account observable", e);
ipc.platform.sshAgent
.clearKeys()
.catch((e) => this.logService.error("Failed to clear SSH keys", e));
},
complete: () => {
this.logService.info("Active account observable completed, clearing SSH keys");
ipc.platform.sshAgent
.clearKeys()
.catch((e) => this.logService.error("Failed to clear SSH keys", e));
},
});

combineLatest([
timer(0, this.SSH_REFRESH_INTERVAL),
this.desktopSettingsService.sshAgentEnabled$,
])
.pipe(
concatMap(async ([, enabled]) => {
if (!enabled) {
await ipc.platform.sshAgent.clearKeys();
return;
}

const ciphers = await this.cipherService.getAllDecrypted();
if (ciphers == null) {
await ipc.platform.sshAgent.lock();
return;
}

const sshCiphers = ciphers.filter(
(cipher) =>
cipher.type === CipherType.SshKey &&
!cipher.isDeleted &&
cipher.organizationId === null,
);
const keys = sshCiphers.map((cipher) => {
return {
name: cipher.name,
privateKey: cipher.sshKey.privateKey,
cipherId: cipher.id,
};
});
await ipc.platform.sshAgent.setKeys(keys);
}),
takeUntil(this.destroy$),
)
.subscribe();
}

ngOnDestroy() {
Expand Down

0 comments on commit 244539c

Please sign in to comment.