From 9c3aab2f705fba387a9a5be5e89c1062d7ef9918 Mon Sep 17 00:00:00 2001
From: jake <77554505+brxken128@users.noreply.github.com>
Date: Thu, 9 Feb 2023 10:07:53 +0000
Subject: [PATCH] [ENG-357] Automatically mount associated key during
 decryption (#564)

* auto-mount associated key during decryption based on content salt

* use iterator directly

* cleanup
---
 core/src/object/fs/decrypt.rs                 | 18 ++++++++++++---
 packages/client/src/core.ts                   |  2 +-
 .../components/dialog/DecryptFileDialog.tsx   | 23 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/core/src/object/fs/decrypt.rs b/core/src/object/fs/decrypt.rs
index 9b93dd37d..fe28db435 100644
--- a/core/src/object/fs/decrypt.rs
+++ b/core/src/object/fs/decrypt.rs
@@ -19,6 +19,7 @@ pub struct FileDecryptorJobState {}
 pub struct FileDecryptorJobInit {
 	pub location_id: i32,
 	pub path_id: i32,
+	pub mount_associated_key: bool,
 	pub output_path: Option<PathBuf>,
 	pub password: Option<String>, // if this is set, we can assume the user chose password decryption
 	pub save_to_library: Option<bool>,
@@ -66,6 +67,7 @@ impl StatefulJob for FileDecryptorJob {
 	) -> Result<(), JobError> {
 		let step = &state.steps[0];
 		let info = &step.fs_info;
+		let key_manager = &ctx.library_ctx.key_manager;
 
 		// handle overwriting checks, and making sure there's enough available space
 		let output_path = state.init.output_path.clone().map_or_else(
@@ -98,8 +100,7 @@ impl StatefulJob for FileDecryptorJob {
 					let index = header.find_key_index(password_bytes.clone()).await?;
 
 					// inherit the encryption algorithm from the keyslot
-					ctx.library_ctx
-						.key_manager
+					key_manager
 						.add_to_keystore(
 							Password::new(password),
 							header.algorithm,
@@ -118,7 +119,18 @@ impl StatefulJob for FileDecryptorJob {
 				)));
 			}
 		} else {
-			let keys = ctx.library_ctx.key_manager.enumerate_hashed_keys();
+			if state.init.mount_associated_key {
+				for key in key_manager.dump_keystore().iter().filter(|x| {
+					header
+						.keyslots
+						.iter()
+						.any(|k| k.content_salt == x.content_salt)
+				}) {
+					key_manager.mount(key.uuid).await.ok();
+				}
+			}
+
+			let keys = key_manager.enumerate_hashed_keys();
 
 			header.decrypt_master_key_from_prehashed(keys).await?
 		};
diff --git a/packages/client/src/core.ts b/packages/client/src/core.ts
index 7828b72e5..d80d82d4c 100644
--- a/packages/client/src/core.ts
+++ b/packages/client/src/core.ts
@@ -107,7 +107,7 @@ export interface FileCopierJobInit { source_location_id: number, source_path_id:
 
 export interface FileCutterJobInit { source_location_id: number, source_path_id: number, target_location_id: number, target_path: string }
 
-export interface FileDecryptorJobInit { location_id: number, path_id: number, output_path: string | null, password: string | null, save_to_library: boolean | null }
+export interface FileDecryptorJobInit { location_id: number, path_id: number, mount_associated_key: boolean, output_path: string | null, password: string | null, save_to_library: boolean | null }
 
 export interface FileDeleterJobInit { location_id: number, path_id: number }
 
diff --git a/packages/interface/src/components/dialog/DecryptFileDialog.tsx b/packages/interface/src/components/dialog/DecryptFileDialog.tsx
index c58ff09af..3b4b98e50 100644
--- a/packages/interface/src/components/dialog/DecryptFileDialog.tsx
+++ b/packages/interface/src/components/dialog/DecryptFileDialog.tsx
@@ -16,6 +16,7 @@ interface DecryptDialogProps extends UseDialogProps {
 const schema = z.object({
 	type: z.union([z.literal('password'), z.literal('key')]),
 	outputPath: z.string(),
+	mountAssociatedKey: z.boolean(),
 	password: z.string(),
 	saveToKeyManager: z.boolean()
 });
@@ -63,7 +64,8 @@ export const DecryptFileDialog = (props: DecryptDialogProps) => {
 			type: hasMountedKeys ? 'key' : 'password',
 			saveToKeyManager: true,
 			outputPath: '',
-			password: ''
+			password: '',
+			mountAssociatedKey: true
 		},
 		schema
 	});
@@ -73,6 +75,7 @@ export const DecryptFileDialog = (props: DecryptDialogProps) => {
 			location_id: props.location_id,
 			path_id: props.path_id,
 			output_path: data.outputPath !== '' ? data.outputPath : null,
+			mount_associated_key: data.mountAssociatedKey,
 			password: data.type === 'password' ? data.password : null,
 			save_to_library: data.type === 'password' ? data.saveToKeyManager : null
 		})
@@ -117,6 +120,24 @@ export const DecryptFileDialog = (props: DecryptDialogProps) => {
 				</div>
 			</RadioGroup>
 
+			{form.watch('type') === 'key' && (
+				<div className="relative flex flex-grow mt-3 mb-2">
+					<div className="space-x-2">
+						<Switch
+							className="bg-app-selected"
+							size="sm"
+							name=""
+							checked={form.watch('mountAssociatedKey')}
+							onCheckedChange={(e) => form.setValue('mountAssociatedKey', e)}
+						/>
+					</div>
+					<span className="ml-3 text-xs font-medium mt-0.5">Automatically mount key</span>
+					<Tooltip label="The key linked with the file will be automatically mounted">
+						<Info className="w-4 h-4 ml-1.5 text-ink-faint mt-0.5" />
+					</Tooltip>
+				</div>
+			)}
+
 			{form.watch('type') === 'password' && (
 				<>
 					<div className="relative flex flex-grow mt-3 mb-2">