@protontech/drive-sdk 0.1.2 → 0.2.0

package/src/internal/download/fileDownloader.test.ts

@@ -1,9 +1,9 @@
-import { Revision } from '../../interface';
+import { APIHTTPError, HTTPErrorCode } from '../apiService';
+import { DecryptedRevision } from '../nodes';
 import { FileDownloader } from './fileDownloader';
 import { DownloadTelemetry } from './telemetry';
 import { DownloadAPIService } from './apiService';
 import { DownloadCryptoService } from './cryptoService';
-import { APIHTTPError, HTTPErrorCode } from '../apiService';
 
 function mockBlockDownload(_: string, token: string, onProgress: (downloadedBytes: number) => void) {
     const index = parseInt(token.slice(5, 6));
@@ -21,7 +21,7 @@ describe('FileDownloader', () => {
     let apiService: DownloadAPIService;
     let cryptoService: DownloadCryptoService;
     let nodeKey: { key: object; contentKeyPacketSessionKey: string };
-    let revision: Revision;
+    let revision: DecryptedRevision;
 
     beforeEach(() => {
         // @ts-expect-error No need to implement all methods for mocking
@@ -74,7 +74,8 @@ describe('FileDownloader', () => {
         revision = {
             uid: 'revisionUid',
             claimedSize: 1024,
-        } as Revision;
+            claimedBlockSizes: [16, 16, 16, 16],
+        } as DecryptedRevision;
     });
 
     describe('writeToStream', () => {
@@ -262,7 +263,7 @@ describe('FileDownloader', () => {
             await verifyOnProgress([1, 2, 3]);
         });
 
-        it('should handle failure when veryfing block', async () => {
+        it('should handle failure when verifying block', async () => {
             cryptoService.verifyBlockIntegrity = jest.fn().mockImplementation(async function () {
                 throw new Error('Failed to verify block');
             });
@@ -270,7 +271,7 @@ describe('FileDownloader', () => {
             await verifyFailure('Failed to verify block', undefined);
         });
 
-        it('should handle one time-off failure when veryfing block', async () => {
+        it('should handle one time-off failure when verifying block', async () => {
             let count = 0;
             cryptoService.verifyBlockIntegrity = jest.fn().mockImplementation(async function () {
                 if (count === 0) {
@@ -335,7 +336,7 @@ describe('FileDownloader', () => {
             await verifyOnProgress([1, 2, 3]);
         });
 
-        it('should handle failure when veryfing manifest', async () => {
+        it('should handle failure when verifying manifest', async () => {
             cryptoService.verifyManifest = jest.fn().mockImplementation(async function () {
                 throw new Error('Failed to verify manifest');
             });
@@ -394,4 +395,96 @@ describe('FileDownloader', () => {
             expect(onFinish).toHaveBeenCalledTimes(1);
         });
     });
+
+    describe('getSeekableStream', () => {
+        let onFinish: () => void;
+        let downloader: FileDownloader;
+
+        beforeEach(() => {
+            apiService.downloadBlock = jest.fn().mockImplementation(async function (_, token) {
+                const index = parseInt(token.slice(5, 6)) - 1;
+                const data = new Uint8Array(16);
+                for (let i = 0; i < data.length; i++) {
+                    data[i] = index * 16 + i;
+                }
+                return data;
+            });
+
+            onFinish = jest.fn();
+
+            downloader = new FileDownloader(
+                telemetry,
+                apiService,
+                cryptoService,
+                nodeKey as any,
+                revision,
+                undefined,
+                onFinish,
+            );
+        });
+
+        it('should read the stream', async () => {
+            const stream = downloader.getSeekableStream();
+
+            const data = await stream.read(32);
+            expect(data.value).toEqual(
+                new Uint8Array([
+                    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26,
+                    27, 28, 29, 30, 31,
+                ]),
+            );
+            expect(data.done).toEqual(false);
+
+            const data2 = await stream.read(32);
+            expect(data2.value).toEqual(
+                new Uint8Array([
+                    32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56,
+                    57, 58, 59, 60, 61, 62, 63,
+                ]),
+            );
+            expect(data2.done).toEqual(false);
+
+            const data3 = await stream.read(32);
+            expect(data3.value).toEqual(new Uint8Array([]));
+            expect(data3.done).toEqual(true);
+
+            expect(cryptoService.decryptBlock).toHaveBeenCalledTimes(4);
+            expect(cryptoService.decryptBlock).toHaveBeenCalledWith(expect.anything(), {
+                key: 'privateKey',
+                contentKeyPacketSessionKey: 'contentSessionKey',
+                verificationKeys: 'verificationKeys',
+            });
+        });
+
+        it('should read the stream with seeking', async () => {
+            const stream = downloader.getSeekableStream();
+
+            const data1 = await stream.read(5);
+            expect(data1.value).toEqual(new Uint8Array([0, 1, 2, 3, 4]));
+            expect(data1.done).toEqual(false);
+            expect(cryptoService.decryptBlock).toHaveBeenCalledTimes(1);
+
+            await stream.seek(10);
+
+            // Seek withing first block, so no new block is downloaded.
+            const data2 = await stream.read(5);
+            expect(data2.value).toEqual(new Uint8Array([10, 11, 12, 13, 14]));
+            expect(data2.done).toEqual(false);
+            expect(cryptoService.decryptBlock).toHaveBeenCalledTimes(1);
+
+            // Seek and read from second and third blocks.
+            await stream.seek(30);
+
+            const data3 = await stream.read(5);
+            expect(data3.value).toEqual(new Uint8Array([30, 31, 32, 33, 34]));
+            expect(data3.done).toEqual(false);
+            expect(cryptoService.decryptBlock).toHaveBeenCalledTimes(3);
+
+            expect(cryptoService.decryptBlock).toHaveBeenCalledWith(expect.anything(), {
+                key: 'privateKey',
+                contentKeyPacketSessionKey: 'contentSessionKey',
+                verificationKeys: 'verificationKeys',
+            });
+        });
+    });
 });

package/src/internal/download/fileDownloader.ts

@@ -1,11 +1,14 @@
 import { PrivateKey, SessionKey, base64StringToUint8Array } from '../../crypto';
-import { Logger, Revision } from '../../interface';
+import { Logger } from '../../interface';
 import { LoggerWithPrefix } from '../../telemetry';
 import { APIHTTPError, HTTPErrorCode } from '../apiService';
+import { DecryptedRevision } from '../nodes';
 import { DownloadAPIService } from './apiService';
+import { getBlockIndex } from './blockIndex';
 import { DownloadController } from './controller';
 import { DownloadCryptoService } from './cryptoService';
 import { BlockMetadata, RevisionKeys } from './interface';
+import { BufferedSeekableStream } from './seekableStream';
 import { DownloadTelemetry } from './telemetry';
 
 /**
@@ -33,7 +36,7 @@ export class FileDownloader {
         private apiService: DownloadAPIService,
         private cryptoService: DownloadCryptoService,
         private nodeKey: { key: PrivateKey; contentKeyPacketSessionKey: SessionKey },
-        private revision: Revision,
+        private revision: DecryptedRevision,
         private signal?: AbortSignal,
         private onFinish?: () => void,
     ) {
@@ -52,6 +55,89 @@ export class FileDownloader {
         return this.revision.claimedSize;
     }
 
+    getSeekableStream(): BufferedSeekableStream {
+        let position = 0;
+        let cryptoKeys: RevisionKeys;
+
+        const logger = new LoggerWithPrefix(this.logger, `seekable stream`);
+
+        const claimedBlockSizes = this.revision.claimedBlockSizes;
+        if (!claimedBlockSizes) {
+            // Old nodes will not have claimed block sizes. One option is to
+            // use default block size, but old clients didn't use the same
+            // size (4 MiB vs 4 MB, for example).
+            // Ideally, we should throw error that client can easily handle,
+            // at the same time, new nodes shouldn't have this issue.
+            // For now, we throw general error that client must handle as any
+            // error from download - do not support seeking and ask user to
+            // download the whole file instead.
+            // In the future, we might either change this error, or have some
+            // clever way to detect block sizes from the first block and work
+            // around this issue.
+            throw new Error('Revision does not have defined claimed block sizes');
+        }
+
+        const stream = new BufferedSeekableStream({
+            start: async () => {
+                logger.debug(`Starting`);
+                cryptoKeys = await this.cryptoService.getRevisionKeys(this.nodeKey, this.revision);
+            },
+            pull: async (controller) => {
+                logger.debug(`Pulling at position ${position}`);
+
+                const result = await this.downloadDataFromPosition(claimedBlockSizes, position, cryptoKeys);
+                if (result instanceof Error) {
+                    logger.error('Download failed', result);
+                    controller.error(result);
+                    return;
+                }
+                if (!result) {
+                    logger.debug(`Download finished at position ${position}`);
+                    controller.close();
+                    return;
+                }
+                controller.enqueue(result);
+                position += result.length;
+            },
+            cancel: (reason?: unknown) => {
+                logger.info(`Cancelled: ${reason}`);
+                this.onFinish?.();
+            },
+            seek: async (newPosition) => {
+                logger.info(`Seeking to position ${newPosition}`);
+                position = newPosition;
+            },
+        });
+        return stream;
+    }
+
+    private async downloadDataFromPosition(
+        claimedBlockSizes: number[],
+        position: number,
+        cryptoKeys: RevisionKeys,
+    ): Promise<Uint8Array | Error | undefined> {
+        const { value, done } = getBlockIndex(claimedBlockSizes, position);
+        if (done) {
+            return;
+        }
+
+        this.logger.info(`Downloading data from block ${value.blockIndex} at offset ${value.blockOffset}`);
+
+        try {
+            const { blockIndex, blockOffset } = value;
+            const blockMetadata = await this.apiService.getRevisionBlockToken(
+                this.revision.uid,
+                blockIndex,
+                this.signal,
+            );
+
+            const blockData = await this.downloadBlockData(blockMetadata, true, cryptoKeys);
+            return blockData.slice(blockOffset);
+        } catch (error: unknown) {
+            return error instanceof Error ? error : new Error(`Unknown error: ${error}`);
+        }
+    }
+
     writeToStream(stream: WritableStream, onProgress?: (downloadedBytes: number) => void): DownloadController {
         if (this.controller.promise) {
             throw new Error(`Download already started`);
@@ -71,7 +157,7 @@ export class FileDownloader {
 
     private async downloadToStream(
         stream: WritableStream,
-        onProgress?: (writtenBytes: number) => void,
+        onProgress?: (downloadedBytes: number) => void,
         ignoreIntegrityErrors = false,
     ): Promise<void> {
         const writer = stream.getWriter();
@@ -113,12 +199,16 @@ export class FileDownloader {
                 this.ongoingDownloads.set(blockMetadata.index, { downloadPromise });
 
                 await this.waitForDownloadCapacity();
-                await this.flushCompletedBlocks(writer);
+                await this.flushCompletedBlocks(async (chunk) => {
+                    await writer.write(chunk);
+                });
             }
 
             this.logger.debug(`All blocks downloading, waiting for them to finish`);
             await Promise.all(this.downloadPromises);
-            await this.flushCompletedBlocks(writer);
+            await this.flushCompletedBlocks(async (chunk) => {
+                await writer.write(chunk);
+            });
 
             if (this.ongoingDownloads.size > 0) {
                 this.logger.error(`Some blocks were not downloaded: ${this.ongoingDownloads.keys()}`);
@@ -158,6 +248,16 @@ export class FileDownloader {
         cryptoKeys: RevisionKeys,
         onProgress: (downloadedBytes: number) => void,
     ) {
+        const blockData = await this.downloadBlockData(blockMetadata, ignoreIntegrityErrors, cryptoKeys, onProgress);
+        this.ongoingDownloads.get(blockMetadata.index)!.decryptedBufferedBlock = blockData;
+    }
+
+    private async downloadBlockData(
+        blockMetadata: BlockMetadata,
+        ignoreIntegrityErrors: boolean,
+        cryptoKeys: RevisionKeys,
+        onProgress?: (downloadedBytes: number) => void,
+    ): Promise<Uint8Array> {
         const logger = new LoggerWithPrefix(this.logger, `block ${blockMetadata.index}`);
         logger.info(`Download started`);
 
@@ -219,8 +319,8 @@ export class FileDownloader {
             }
         }
 
-        this.ongoingDownloads.get(blockMetadata.index)!.decryptedBufferedBlock = decryptedBlock;
         logger.info(`Downloaded`);
+        return decryptedBlock;
     }
 
     private async waitForDownloadCapacity() {
@@ -251,16 +351,16 @@ export class FileDownloader {
         }
     }
 
-    private async flushCompletedBlocks(writer: WritableStreamDefaultWriter<Uint8Array>) {
+    private async flushCompletedBlocks(write: (chunk: Uint8Array) => void | Promise<void>) {
         this.logger.debug(`Flushing completed blocks`);
         while (this.isNextBlockDownloaded) {
             const decryptedBlock = this.ongoingDownloads.get(this.nextBlockIndex)!.decryptedBufferedBlock!;
             this.logger.info(`Flushing completed block ${this.nextBlockIndex}`);
             try {
-                await writer.write(decryptedBlock);
+                await write(decryptedBlock);
             } catch (error) {
                 this.logger.error(`Failed to write block, retrying once`, error);
-                await writer.write(decryptedBlock);
+                await write(decryptedBlock);
             }
             this.ongoingDownloads.delete(this.nextBlockIndex);
             this.nextBlockIndex++;

package/src/internal/download/interface.ts

@@ -1,6 +1,6 @@
 import { PrivateKey, PublicKey, SessionKey } from '../../crypto';
-import { NodeType, Result, Revision, MissingNode, MetricVolumeType } from '../../interface';
-import { DecryptedNode } from '../nodes';
+import { NodeType, Result, MissingNode, MetricVolumeType } from '../../interface';
+import { DecryptedNode, DecryptedRevision } from '../nodes';
 
 export type BlockMetadata = {
     index: number;
@@ -29,9 +29,9 @@ export interface NodesService {
 export interface NodesServiceNode {
     uid: string;
     type: NodeType;
-    activeRevision?: Result<Revision, Error>;
+    activeRevision?: Result<DecryptedRevision, Error>;
 }
 
 export interface RevisionsService {
-    getRevision(nodeRevisionUid: string): Promise<Revision>;
+    getRevision(nodeRevisionUid: string): Promise<DecryptedRevision>;
 }

package/src/internal/download/seekableStream.test.ts

@@ -0,0 +1,187 @@
+import { SeekableReadableStream, BufferedSeekableStream } from './seekableStream';
+
+describe('SeekableReadableStream', () => {
+    it('should call the seek callback when seek is called', async () => {
+        const mockSeek = jest.fn().mockResolvedValue(undefined);
+        const mockStart = jest.fn();
+
+        const stream = new SeekableReadableStream({
+            start: mockStart,
+            seek: mockSeek,
+        });
+
+        await stream.seek(100);
+
+        expect(mockSeek).toHaveBeenCalledWith(100);
+        expect(mockSeek).toHaveBeenCalledTimes(1);
+    });
+
+    it('should handle synchronous seek callback', async () => {
+        const mockSeek = jest.fn().mockReturnValue(undefined);
+
+        const stream = new SeekableReadableStream({
+            seek: mockSeek,
+        });
+
+        await stream.seek(250);
+
+        expect(mockSeek).toHaveBeenCalledWith(250);
+    });
+});
+
+describe('BufferedSeekableStream', () => {
+    let startWithCloseMock: jest.Mock;
+    let pullMock: jest.Mock;
+
+    const data1 = new Uint8Array([1, 2, 3, 4, 5]);
+    const data2 = new Uint8Array([6, 7, 8, 9, 10]);
+
+    beforeEach(() => {
+        startWithCloseMock = jest.fn().mockImplementation((controller) => {
+            controller.enqueue(data1);
+            controller.close();
+        });
+
+        let readIndex = 0;
+        pullMock = jest.fn().mockImplementation((controller) => {
+            if (readIndex === 0) {
+                controller.enqueue(data1);
+            } else if (readIndex === 1) {
+                controller.enqueue(data2);
+            } else {
+                controller.close();
+            }
+            readIndex++;
+        });
+    });
+
+    it('should throw error if highWaterMark is not 0', () => {
+        expect(() => {
+            new BufferedSeekableStream({ seek: jest.fn() }, { highWaterMark: 1 });
+        }).toThrow('highWaterMark must be 0');
+    });
+
+    it('should throw error when reading invalid number of bytes', async () => {
+        const stream = new BufferedSeekableStream({
+            seek: jest.fn(),
+        });
+
+        await expect(stream.read(0)).rejects.toThrow('Invalid number of bytes to read');
+        await expect(stream.read(-1)).rejects.toThrow('Invalid number of bytes to read');
+    });
+
+    it('should read exact number of bytes when underlying source provides exact amount', async () => {
+        const stream = new BufferedSeekableStream({
+            start: startWithCloseMock,
+            seek: jest.fn(),
+        });
+
+        const result = await stream.read(5);
+
+        expect(result).toEqual({ value: data1, done: false });
+    });
+
+    it('should buffer extra bytes when underlying source provides more than requested', async () => {
+        const stream = new BufferedSeekableStream({
+            pull: pullMock,
+            seek: jest.fn(),
+        });
+
+        const result1 = await stream.read(3);
+        expect(result1).toEqual({ value: new Uint8Array([1, 2, 3]), done: false });
+        expect(pullMock).toHaveBeenCalledTimes(1);
+
+        const result2 = await stream.read(2);
+        expect(result2).toEqual({ value: new Uint8Array([4, 5]), done: false });
+        expect(pullMock).toHaveBeenCalledTimes(1);
+    });
+
+    it('should use buffered data and read more when buffer is not enough for next read', async () => {
+        const stream = new BufferedSeekableStream({
+            pull: pullMock,
+            seek: jest.fn(),
+        });
+
+        const result1 = await stream.read(3);
+        expect(result1).toEqual({ value: new Uint8Array([1, 2, 3]), done: false });
+        expect(pullMock).toHaveBeenCalledTimes(1);
+
+        const result2 = await stream.read(5);
+        expect(result2).toEqual({ value: new Uint8Array([4, 5, 6, 7, 8]), done: false });
+        expect(pullMock).toHaveBeenCalledTimes(2);
+    });
+
+    it('should handle end of file gracefully when not enough data available', async () => {
+        const stream = new BufferedSeekableStream({
+            start: startWithCloseMock,
+            seek: jest.fn(),
+        });
+
+        const result = await stream.read(10);
+        expect(result).toEqual({ value: data1, done: true });
+    });
+
+    it('should clear buffer when seeking back', async () => {
+        const stream = new BufferedSeekableStream({
+            pull: pullMock,
+            seek: jest.fn(),
+        });
+
+        const result1 = await stream.read(2);
+        expect(result1).toEqual({ value: new Uint8Array([1, 2]), done: false });
+
+        await stream.seek(0);
+
+        const result2 = await stream.read(3);
+        expect(result2).toEqual({ value: new Uint8Array([6, 7, 8]), done: false });
+    });
+
+    it('should clear buffer when seeking past buffer end', async () => {
+        const stream = new BufferedSeekableStream({
+            pull: pullMock,
+            seek: jest.fn(),
+        });
+
+        const result1 = await stream.read(2);
+        expect(result1).toEqual({ value: new Uint8Array([1, 2]), done: false });
+
+        await stream.seek(100);
+
+        const result2 = await stream.read(3);
+        expect(result2).toEqual({ value: new Uint8Array([6, 7, 8]), done: false });
+    });
+
+    it('should update buffer correctly when seeking within buffer range', async () => {
+        const stream = new BufferedSeekableStream({
+            pull: pullMock,
+            seek: jest.fn(),
+        });
+
+        const result1 = await stream.read(1);
+        expect(result1).toEqual({ value: new Uint8Array([1]), done: false });
+
+        await stream.seek(3);
+
+        const result2 = await stream.read(3);
+        expect(result2).toEqual({ value: new Uint8Array([4, 5, 6]), done: false });
+    });
+
+    it('should handle multiple read operations correctly', async () => {
+        const stream = new BufferedSeekableStream({
+            pull: pullMock,
+            seek: jest.fn(),
+        });
+
+        const result1 = await stream.read(2);
+        expect(result1).toEqual({ value: new Uint8Array([1, 2]), done: false });
+
+        const result2 = await stream.read(4);
+        expect(result2).toEqual({ value: new Uint8Array([3, 4, 5, 6]), done: false });
+
+        const result3 = await stream.read(3);
+        expect(result3).toEqual({ value: new Uint8Array([7, 8, 9]), done: false });
+
+        const result4 = await stream.read(2);
+        expect(result4).toEqual({ value: new Uint8Array([10]), done: true });
+    });
+});