@protontech/drive-sdk 0.2.1 → 0.3.1

package/src/internal/nodes/cryptoService.ts

@@ -61,6 +61,8 @@ export class NodesCryptoService {
         node: EncryptedNode,
         parentKey: PrivateKey,
     ): Promise<{ node: DecryptedUnparsedNode; keys?: DecryptedNodeKeys }> {
+        const start = Date.now();
+
         const commonNodeMetadata = {
             ...node,
             encryptedCrypto: undefined,
@@ -89,16 +91,17 @@ export class NodesCryptoService {
                 : nodeParentKeys;
         }
 
-        const { name, author: nameAuthor } = await this.decryptName(node, parentKey, nameVerificationKeys);
-
-        let membership;
-        if (node.membership) {
-            membership = await this.decryptMembership(node);
-        }
+        // Start promises early, but await them only when required to do
+        // as much work as possible in parallel.
+        const [membershipPromise, namePromise, keyPromise] = [
+            node.membership ? this.decryptMembership(node) : undefined,
+            this.decryptName(node, parentKey, nameVerificationKeys),
+            this.decryptKey(node, parentKey, keyVerificationKeys),
+        ];
 
         let passphrase, key, passphraseSessionKey, keyAuthor;
         try {
-            const keyResult = await this.decryptKey(node, parentKey, keyVerificationKeys);
+            const keyResult = await keyPromise;
             passphrase = keyResult.passphrase;
             key = keyResult.key;
             passphraseSessionKey = keyResult.passphraseSessionKey;
@@ -107,12 +110,17 @@ export class NodesCryptoService {
             void this.reportDecryptionError(node, 'nodeKey', error);
             const message = getErrorMessage(error);
             const errorMessage = c('Error').t`Failed to decrypt node key: ${message}`;
+            const { name, author: nameAuthor } = await namePromise;
+            const membership = await membershipPromise;
             return {
                 node: {
                     ...commonNodeMetadata,
                     name,
                     keyAuthor: resultError({
-                        claimedAuthor: node.encryptedCrypto.signatureEmail,
+                        claimedAuthor: getClaimedAuthor(
+                            node.encryptedCrypto.signatureEmail,
+                            keyVerificationKeys.length === 0,
+                        ),
                         error: errorMessage,
                     }),
                     nameAuthor,
@@ -131,8 +139,23 @@ export class NodesCryptoService {
         let folder;
         let folderExtendedAttributesAuthor;
         if ('folder' in node.encryptedCrypto) {
+            const folderExtendedAttributesVerificationKeys = node.encryptedCrypto.signatureEmail
+                ? signatureEmailKeys
+                : [key];
+
+            const [hashKeyPromise, folderExtendedAttributesPromise] = [
+                this.decryptHashKey(node, key, signatureEmailKeys),
+                this.decryptExtendedAttributes(
+                    node,
+                    node.encryptedCrypto.folder.armoredExtendedAttributes,
+                    key,
+                    folderExtendedAttributesVerificationKeys,
+                    node.encryptedCrypto.signatureEmail,
+                ),
+            ];
+
             try {
-                const hashKeyResult = await this.decryptHashKey(node, key, signatureEmailKeys);
+                const hashKeyResult = await hashKeyPromise;
                 hashKey = hashKeyResult.hashKey;
                 hashKeyAuthor = hashKeyResult.author;
             } catch (error: unknown) {
@@ -141,16 +164,7 @@ export class NodesCryptoService {
             }
 
             try {
-                const folderExtendedAttributesVerificationKeys = node.encryptedCrypto.signatureEmail
-                    ? signatureEmailKeys
-                    : [key];
-                const extendedAttributesResult = await this.decryptExtendedAttributes(
-                    node,
-                    node.encryptedCrypto.folder.armoredExtendedAttributes,
-                    key,
-                    folderExtendedAttributesVerificationKeys,
-                    node.encryptedCrypto.signatureEmail,
-                );
+                const extendedAttributesResult = await folderExtendedAttributesPromise;
                 folder = {
                     extendedAttributes: extendedAttributesResult.extendedAttributes,
                 };
@@ -165,10 +179,20 @@ export class NodesCryptoService {
         let contentKeyPacketSessionKey;
         let contentKeyPacketAuthor;
         if ('file' in node.encryptedCrypto) {
+            const [activeRevisionPromise, contentKeyPacketSessionKeyPromise] = [
+                this.decryptRevision(node.uid, node.encryptedCrypto.activeRevision, key),
+                this.driveCrypto.decryptAndVerifySessionKey(
+                    node.encryptedCrypto.file.base64ContentKeyPacket,
+                    node.encryptedCrypto.file.armoredContentKeyPacketSignature,
+                    key,
+                    // Content key packet is signed with the node key, but
+                    // in the past some clients signed with the address key.
+                    [key, ...keyVerificationKeys],
+                ),
+            ];
+
             try {
-                activeRevision = resultOk(
-                    await this.decryptRevision(node.uid, node.encryptedCrypto.activeRevision, key),
-                );
+                activeRevision = resultOk(await activeRevisionPromise);
             } catch (error: unknown) {
                 void this.reportDecryptionError(node, 'nodeExtendedAttributes', error);
                 const message = getErrorMessage(error);
@@ -177,18 +201,10 @@ export class NodesCryptoService {
             }
 
             try {
-                const keySessionKeyResult = await this.driveCrypto.decryptAndVerifySessionKey(
-                    node.encryptedCrypto.file.base64ContentKeyPacket,
-                    node.encryptedCrypto.file.armoredContentKeyPacketSignature,
-                    key,
-                    // Content key packet is signed with the node key, but
-                    // in the past some clients signed with the address key.
-                    [key, ...keyVerificationKeys],
-                );
-
+                const keySessionKeyResult = await contentKeyPacketSessionKeyPromise;
                 contentKeyPacketSessionKey = keySessionKeyResult.sessionKey;
                 contentKeyPacketAuthor =
-                    keySessionKeyResult.verified &&
+                    keySessionKeyResult.verified !== undefined &&
                     (await this.handleClaimedAuthor(
                         node,
                         'nodeContentKey',
@@ -230,6 +246,13 @@ export class NodesCryptoService {
             finalKeyAuthor = keyAuthor;
         }
 
+        const { name, author: nameAuthor } = await namePromise;
+        const membership = await membershipPromise;
+
+        const end = Date.now();
+        const duration = end - start;
+        this.logger.debug(`Node ${node.uid} decrypted in ${duration}ms`);
+
         return {
             node: {
                 ...commonNodeMetadata,
@@ -319,7 +342,7 @@ export class NodesCryptoService {
             return {
                 name: resultError(new Error(errorMessage)),
                 author: resultError({
-                    claimedAuthor: nameSignatureEmail,
+                    claimedAuthor: getClaimedAuthor(nameSignatureEmail, verificationKeys.length === 0),
                     error: errorMessage,
                 }),
             };
@@ -635,6 +658,7 @@ export class NodesCryptoService {
         if (this.reportedVerificationErrors.has(node.uid)) {
             return;
         }
+        this.reportedVerificationErrors.add(node.uid);
 
         const fromBefore2024 = node.creationTime < new Date('2024-01-01');
 
@@ -661,7 +685,6 @@ export class NodesCryptoService {
             error: verificationErrors?.map((e) => e.message).join(', '),
             uid: node.uid,
         });
-        this.reportedVerificationErrors.add(node.uid);
     }
 
     private async reportDecryptionError(node: EncryptedNode, field: MetricsDecryptionErrorField, error: unknown) {
@@ -716,3 +739,14 @@ function handleClaimedAuthor(
         error: getVerificationMessage(verified, verificationErrors, signatureType, notAvailableVerificationKeys),
     });
 }
+
+function getClaimedAuthor(
+    claimedAuthor?: string,
+    notAvailableVerificationKeys = false,
+): string | AnonymousUser | undefined {
+    if (!claimedAuthor && notAvailableVerificationKeys) {
+        return null as AnonymousUser;
+    }
+
+    return claimedAuthor;
+}

package/src/internal/nodes/interface.ts

@@ -18,6 +18,8 @@ import {
 interface BaseNode {
     // Internal metadata
     hash?: string; // root node doesn't have any hash
+    // ecnryptedName should not be needed to keep, nameSessionKey should be enough.
+    // We will improve this in the future.
     encryptedName: string;
 
     // Basic node metadata

package/src/internal/nodes/nodesAccess.ts

@@ -24,7 +24,7 @@ const BATCH_LOADING_SIZE = 30;
 // It is a trade-off between performance and memory usage.
 // Higher number means more memory usage, but faster decryption.
 // Lower number means less memory usage, but slower decryption.
-const DECRYPTION_CONCURRENCY = 15;
+const DECRYPTION_CONCURRENCY = 30;
 
 /**
  * Provides access to node metadata.
@@ -234,7 +234,7 @@ export class NodesAccess {
 
         if (errors.length > 0) {
             this.logger.error(`Failed to decrypt ${errors.length} nodes`, errors);
-            throw new ProtonDriveError(c('Error').t`Failed to decrypt some nodes`, { cause: errors });
+            throw new DecryptionError(c('Error').t`Failed to decrypt some nodes`, { cause: errors });
         }
 
         const missingNodeUids = nodeUids.filter((nodeUid) => !returnedNodeUids.includes(nodeUid));

package/src/internal/nodes/nodesManagement.ts

@@ -189,11 +189,8 @@ export class NodesManagement {
     }
 
     async *deleteNodes(nodeUids: string[], signal?: AbortSignal): AsyncGenerator<NodeResult> {
-        const deletedNodeUids = [];
-
         for await (const result of this.apiService.deleteNodes(nodeUids, signal)) {
             if (result.ok) {
-                deletedNodeUids.push(result.uid);
                 await this.nodesAccess.notifyNodeDeleted(result.uid);
             }
             yield result;

package/src/internal/shares/cryptoCache.test.ts

@@ -1,6 +1,7 @@
 import { PrivateKey, SessionKey } from '../../crypto';
 import { MemoryCache } from '../../cache';
 import { CachedCryptoMaterial } from '../../interface';
+import { getMockLogger } from '../../tests/logger';
 import { SharesCryptoCache } from './cryptoCache';
 
 describe('sharesCryptoCache', () => {
@@ -17,7 +18,7 @@ describe('sharesCryptoCache', () => {
 
     beforeEach(() => {
         memoryCache = new MemoryCache();
-        cache = new SharesCryptoCache(memoryCache);
+        cache = new SharesCryptoCache(getMockLogger(), memoryCache);
     });
 
     it('should store and retrieve keys', async () => {
@@ -53,7 +54,7 @@ describe('sharesCryptoCache', () => {
         const keys = { key: generatePrivateKey('privateKey'), passphraseSessionKey: generateSessionKey('sessionKey') };
 
         await cache.setShareKey(shareId, keys);
-        await cache.removeShareKey([shareId]);
+        await cache.removeShareKeys([shareId]);
 
         try {
             await cache.getShareKey(shareId);

package/src/internal/shares/cryptoCache.ts

@@ -1,4 +1,4 @@
-import { ProtonDriveCryptoCache } from '../../interface';
+import { Logger, ProtonDriveCryptoCache } from '../../interface';
 import { DecryptedShareKey } from './interface';
 
 /**
@@ -13,23 +13,42 @@ import { DecryptedShareKey } from './interface';
  * only the root node, thus share cache is not needed.
  */
 export class SharesCryptoCache {
-    constructor(private driveCache: ProtonDriveCryptoCache) {
+    constructor(
+        private logger: Logger,
+        private driveCache: ProtonDriveCryptoCache,
+    ) {
+        this.logger = logger;
         this.driveCache = driveCache;
     }
 
     async setShareKey(shareId: string, key: DecryptedShareKey): Promise<void> {
-        await this.driveCache.setEntity(getCacheUid(shareId), key);
+        await this.driveCache.setEntity(getCacheKey(shareId), {
+            shareKey: key,
+        });
     }
 
     async getShareKey(shareId: string): Promise<DecryptedShareKey> {
-        return this.driveCache.getEntity(getCacheUid(shareId));
+        const nodeKeysData = await this.driveCache.getEntity(getCacheKey(shareId));
+        if (!nodeKeysData.shareKey) {
+            try {
+                await this.removeShareKeys([shareId]);
+            } catch (removingError: unknown) {
+                // The node keys will not be returned, thus SDK will re-fetch
+                // and re-cache it. Setting it again should then fix the problem.
+                this.logger.warn(
+                    `Failed to remove corrupted node keys from the cache: ${removingError instanceof Error ? removingError.message : removingError}`,
+                );
+            }
+            throw new Error(`Failed to deserialize node keys`);
+        }
+        return nodeKeysData.shareKey;
     }
 
-    async removeShareKey(shareIds: string[]): Promise<void> {
-        await this.driveCache.removeEntities(shareIds.map(getCacheUid));
+    async removeShareKeys(shareIds: string[]): Promise<void> {
+        await this.driveCache.removeEntities(shareIds.map(getCacheKey));
     }
 }
 
-function getCacheUid(shareId: string) {
+function getCacheKey(shareId: string) {
     return `shareKey-${shareId}`;
 }

package/src/internal/shares/index.ts

@@ -33,7 +33,7 @@ export function initSharesModule(
 ) {
     const api = new SharesAPIService(apiService);
     const cache = new SharesCache(telemetry.getLogger('shares-cache'), driveEntitiesCache);
-    const cryptoCache = new SharesCryptoCache(driveCryptoCache);
+    const cryptoCache = new SharesCryptoCache(telemetry.getLogger('shares-cache'), driveCryptoCache);
     const cryptoService = new SharesCryptoService(telemetry, crypto, account);
     const sharesManager = new SharesManager(
         telemetry.getLogger('shares'),

package/src/internal/sharing/cryptoService.test.ts

@@ -9,7 +9,7 @@ import {
 } from '../../interface';
 import { getMockTelemetry } from '../../tests/telemetry';
 import { SharesService } from './interface';
-import { SharingCryptoService } from './cryptoService';
+import { PUBLIC_LINK_GENERATED_PASSWORD_LENGTH, SharingCryptoService } from './cryptoService';
 
 describe('SharingCryptoService', () => {
     let telemetry: ProtonDriveTelemetry;
@@ -87,6 +87,27 @@ describe('SharingCryptoService', () => {
             expect(telemetry.recordMetric).not.toHaveBeenCalled();
         });
 
+        it('should decrypt bookmark with custom password', async () => {
+            // First 12 characters are the generated password. Anything beyond is the custom password.
+            driveCrypto.decryptShareUrlPassword = jest.fn().mockResolvedValue('urlPassword1WithCustomPassword');
+
+            const result = await cryptoService.decryptBookmark(encryptedBookmark);
+
+            expect(result).toMatchObject({
+                url: resultOk('https://drive.proton.me/urls/tokenId#urlPassword1'),
+                nodeName: resultOk('nodeName'),
+            });
+            expect(driveCrypto.decryptShareUrlPassword).toHaveBeenCalledWith('encryptedUrlPassword', ['addressKey']);
+            expect(driveCrypto.decryptKeyWithSrpPassword).toHaveBeenCalledWith(
+                'urlPassword1WithCustomPassword',
+                'base64SharePasswordSalt',
+                'armoredKey',
+                'armoredPassphrase',
+            );
+            expect(driveCrypto.decryptNodeName).toHaveBeenCalledWith('encryptedName', 'decryptedKey', []);
+            expect(telemetry.recordMetric).not.toHaveBeenCalled();
+        });
+
         it('should handle undecryptable URL password', async () => {
             const error = new Error('Failed to decrypt URL password');
             driveCrypto.decryptShareUrlPassword = jest.fn().mockRejectedValue(error);

package/src/internal/sharing/cryptoService.ts

@@ -37,6 +37,7 @@ import {
     EncryptedBookmark,
     SharesService,
 } from './interface';
+import { DecryptionError } from '../../errors';
 
 // Version 2 of bcrypt with 2**10 rounds.
 // https://en.wikipedia.org/wiki/Bcrypt#Description
@@ -425,10 +426,11 @@ export class SharingCryptoService {
         // TODO: Signatures are not checked and not specified in the interface.
         // In the future, we will need to add authorship verification.
 
+        let password: string;
         let urlPassword: string;
         let customPassword: Result<string | undefined, Error>;
         try {
-            const password = await this.decryptBookmarkUrlPassword(encryptedBookmark);
+            password = await this.decryptBookmarkUrlPassword(encryptedBookmark);
             const result = splitGeneratedAndCustomPassword(password);
             urlPassword = result.password;
             customPassword = resultOk(result.customPassword);
@@ -446,7 +448,7 @@ export class SharingCryptoService {
 
         let shareKey: PrivateKey;
         try {
-            shareKey = await this.decryptBookmarkKey(encryptedBookmark, urlPassword);
+            shareKey = await this.decryptBookmarkKey(encryptedBookmark, password);
         } catch (originalError: unknown) {
             const error = originalError instanceof Error ? originalError : new Error(c('Error').t`Unknown error`);
             return {
@@ -493,15 +495,15 @@ export class SharingCryptoService {
 
             const message = getErrorMessage(error);
             const errorMessage = c('Error').t`Failed to decrypt bookmark password: ${message}`;
-            throw new Error(errorMessage);
+            throw new DecryptionError(errorMessage, { cause: error });
         }
     }
 
-    private async decryptBookmarkKey(encryptedBookmark: EncryptedBookmark, urlPassword: string): Promise<PrivateKey> {
+    private async decryptBookmarkKey(encryptedBookmark: EncryptedBookmark, password: string): Promise<PrivateKey> {
         try {
             // Use the password to decrypt the share key.
             const { key: shareKey } = await this.driveCrypto.decryptKeyWithSrpPassword(
-                urlPassword,
+                password,
                 encryptedBookmark.url.base64SharePasswordSalt,
                 encryptedBookmark.share.armoredKey,
                 encryptedBookmark.share.armoredPassphrase,
@@ -519,7 +521,7 @@ export class SharingCryptoService {
 
             const message = getErrorMessage(error);
             const errorMessage = c('Error').t`Failed to decrypt bookmark key: ${message}`;
-            throw new Error(errorMessage);
+            throw new DecryptionError(errorMessage, { cause: error });
         }
     }
 

package/src/internal/sharing/index.ts

@@ -32,6 +32,7 @@ export function initSharingModule(
     const sharingManagement = new SharingManagement(
         telemetry.getLogger('sharing'),
         api,
+        cache,
         cryptoService,
         account,
         sharesService,

package/src/internal/sharing/interface.ts

@@ -145,10 +145,6 @@ export interface SharesService {
     getMyFilesIDs(): Promise<{ volumeId: string }>;
     loadEncryptedShare(shareId: string): Promise<EncryptedShare>;
     getMyFilesShareMemberEmailKey(): Promise<{
-        addressId: string;
-        addressKey: PrivateKey;
-    }>;
-    getContextShareMemberEmailKey(shareId: string): Promise<{
         email: string;
         addressId: string;
         addressKey: PrivateKey;

package/src/internal/sharing/sharingAccess.test.ts

@@ -3,7 +3,7 @@ import { SharingAPIService } from './apiService';
 import { SharingCache } from './cache';
 import { SharingCryptoService } from './cryptoService';
 import { SharesService, NodesService } from './interface';
-import { SharingAccess } from './sharingAccess';
+import { SharingAccess, BATCH_LOADING_SIZE } from './sharingAccess';
 
 describe('SharingAccess', () => {
     let apiService: SharingAPIService;
@@ -14,7 +14,7 @@ describe('SharingAccess', () => {
 
     let sharingAccess: SharingAccess;
 
-    const nodeUids = Array.from({ length: 15 }, (_, i) => `nodeUid${i}`);
+    const nodeUids = Array.from({ length: BATCH_LOADING_SIZE + 5 }, (_, i) => `nodeUid${i}`);
     const nodes = nodeUids.map((nodeUid) => ({ nodeUid }));
     const nodeUidsIterator = async function* () {
         for (const nodeUid of nodeUids) {
@@ -84,7 +84,7 @@ describe('SharingAccess', () => {
 
             expect(result).toEqual(nodes);
             expect(apiService.iterateSharedNodeUids).toHaveBeenCalledWith('volumeId', undefined);
-            expect(nodesService.iterateNodes).toHaveBeenCalledTimes(2); // 15 / 10 per batch
+            expect(nodesService.iterateNodes).toHaveBeenCalledTimes(2); // Mocked is a bit more over one batch
             expect(cache.setSharedByMeNodeUids).toHaveBeenCalledWith(nodeUids);
         });
     });
@@ -107,7 +107,7 @@ describe('SharingAccess', () => {
 
             expect(result).toEqual(nodes);
             expect(apiService.iterateSharedWithMeNodeUids).toHaveBeenCalledWith(undefined);
-            expect(nodesService.iterateNodes).toHaveBeenCalledTimes(2); // 15 / 10 per batch
+            expect(nodesService.iterateNodes).toHaveBeenCalledTimes(2); // Mocked is a bit more over one batch
             expect(cache.setSharedWithMeNodeUids).toHaveBeenCalledWith(nodeUids);
         });
     });

package/src/internal/sharing/sharingAccess.ts

@@ -9,6 +9,10 @@ import { SharingCache } from './cache';
 import { SharingCryptoService } from './cryptoService';
 import { SharesService, NodesService } from './interface';
 
+// This is the number of nodes that are loaded in parallel.
+// It is a trade-off between initial wait time and overhead of API calls.
+export const BATCH_LOADING_SIZE = 30;
+
 /**
  * Provides high-level actions for access shared nodes.
  *
@@ -66,6 +70,7 @@ export class SharingAccess {
     ): AsyncGenerator<DecryptedNode> {
         const batchLoading = new BatchLoading<string, DecryptedNode>({
             iterateItems: (nodeUids) => this.iterateNodesAndIgnoreMissingOnes(nodeUids, signal),
+            batchSize: BATCH_LOADING_SIZE,
         });
         for (const nodeUid of nodeUids) {
             yield* batchLoading.load(nodeUid);
@@ -81,6 +86,7 @@ export class SharingAccess {
         const loadedNodeUids = [];
         const batchLoading = new BatchLoading<string, DecryptedNode>({
             iterateItems: (nodeUids) => this.iterateNodesAndIgnoreMissingOnes(nodeUids, signal),
+            batchSize: BATCH_LOADING_SIZE,
         });
         for await (const nodeUid of nodeUidsIterator) {
             loadedNodeUids.push(nodeUid);