Skip to content

Commit

Permalink
Request multiple blobs of the same instance in one request
Browse files Browse the repository at this point in the history 
#8069

Co-authored-by: ivk <[email protected]>
  • Loading branch information
@BijinDev @charlag @paw-hub
2 people authored and paw-hub committed Dec 20, 2024
1 parent 4d7086c commit 5c02afc
Show file tree
Hide file tree
Showing 2 changed files with 239 additions and 10 deletions.
107 changes: 102 additions & 5 deletions src/common/api/worker/facades/lazy/BlobFacade.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,19 @@
import { addParamsToUrl, isSuspensionResponse, RestClient, SuspensionBehavior } from "../../rest/RestClient.js"
import { CryptoFacade } from "../../crypto/CryptoFacade.js"
import { clear, concat, isEmpty, neverNull, promiseMap, splitUint8ArrayInChunks, uint8ArrayToBase64, uint8ArrayToString } from "@tutao/tutanota-utils"
import {
assertNonNull,
base64ToBase64Ext,
clear,
concat,
getFirstOrThrow,
groupBy,
isEmpty,
neverNull,
promiseMap,
splitUint8ArrayInChunks,
uint8ArrayToBase64,
uint8ArrayToString,
} from "@tutao/tutanota-utils"
import { ArchiveDataType, MAX_BLOB_SIZE_BYTES } from "../../../common/TutanotaConstants.js"

import { HttpMethod, MediaType, resolveTypeReference } from "../../../common/EntityFunctions.js"
Expand All @@ -16,7 +29,7 @@ import { FileReference } from "../../../common/utils/FileUtils.js"
import { handleRestError } from "../../../common/error/RestError.js"
import { ProgrammingError } from "../../../common/error/ProgrammingError.js"
import { IServiceExecutor } from "../../../common/ServiceRequest.js"
import { BlobGetInTypeRef, BlobPostOut, BlobPostOutTypeRef, BlobServerAccessInfo, createBlobGetIn } from "../../../entities/storage/TypeRefs.js"
import { BlobGetInTypeRef, BlobPostOut, BlobPostOutTypeRef, BlobServerAccessInfo, createBlobGetIn, createBlobId } from "../../../entities/storage/TypeRefs.js"
import { AuthDataProvider } from "../UserFacade.js"
import { doBlobRequestWithRetry, tryServers } from "../../rest/EntityRestClient.js"
import { BlobAccessTokenFacade } from "../BlobAccessTokenFacade.js"
Expand Down Expand Up @@ -114,14 +127,29 @@ export class BlobFacade {
blobLoadOptions: BlobLoadOptions = {},
): Promise<Uint8Array> {
const sessionKey = await this.resolveSessionKey(referencingInstance.entity)
// Currently assumes that all the blobs of the instance are in the same archive.
// If this changes we need to group by archive and do request for each archive and then concatenate all the chunks.
const doBlobRequest = async () => {
const blobServerAccessInfo = await this.blobAccessTokenFacade.requestReadTokenBlobs(archiveDataType, referencingInstance, blobLoadOptions)
return promiseMap(referencingInstance.blobs, (blob) => this.downloadAndDecryptChunk(blob, blobServerAccessInfo, sessionKey, blobLoadOptions))
return this.downloadAndDecryptMultipleBlobsOfArchive(referencingInstance.blobs, blobServerAccessInfo, sessionKey, blobLoadOptions)
}
const doEvictToken = () => this.blobAccessTokenFacade.evictReadBlobsToken(referencingInstance)

const blobData = await doBlobRequestWithRetry(doBlobRequest, doEvictToken)
return concat(...blobData)
const blobChunks = await doBlobRequestWithRetry(doBlobRequest, doEvictToken)
return this.concatenateBlobChunks(referencingInstance, blobChunks)
}

private concatenateBlobChunks(referencingInstance: BlobReferencingInstance, blobChunks: Map<Id, Uint8Array>) {
const resultSize = Array.from(blobChunks.values()).reduce((sum, blob) => blob.length + sum, 0)
const resultBuffer = new Uint8Array(resultSize)
let offset = 0
for (const blob of referencingInstance.blobs) {
const data = blobChunks.get(blob.blobId)
assertNonNull(data, `Server did not return blob for id : ${blob.blobId}`)
resultBuffer.set(data, offset)
offset += data.length
}
return resultBuffer
}

/**
Expand Down Expand Up @@ -355,6 +383,52 @@ export class BlobFacade {
)
}

private async downloadAndDecryptMultipleBlobsOfArchive(
blobs: readonly Blob[],
blobServerAccessInfo: BlobServerAccessInfo,
sessionKey: AesKey,
blobLoadOptions: BlobLoadOptions,
): Promise<Map<Id, Uint8Array>> {
if (isEmpty(blobs)) {
throw new ProgrammingError("Blobs are empty")
}
const archiveId = getFirstOrThrow(blobs).archiveId
if (blobs.some((blob) => blob.archiveId !== archiveId)) {
throw new ProgrammingError("Must only request blobs of the same archive together")
}
const getData = createBlobGetIn({
archiveId,
blobId: null,
blobIds: blobs.map(({ blobId }) => createBlobId({ blobId: blobId })),
})
const BlobGetInTypeModel = await resolveTypeReference(BlobGetInTypeRef)
const literalGetData = await this.instanceMapper.encryptAndMapToLiteral(BlobGetInTypeModel, getData, null)
const body = JSON.stringify(literalGetData)
const queryParams = await this.blobAccessTokenFacade.createQueryParams(blobServerAccessInfo, {}, BlobGetInTypeRef)
const concatBinaryData = await tryServers(
blobServerAccessInfo.servers,
async (serverUrl) => {
return await this.restClient.request(BLOB_SERVICE_REST_PATH, HttpMethod.GET, {
queryParams: queryParams,
body,
responseType: MediaType.Binary,
baseUrl: serverUrl,
noCORS: true,
headers: blobLoadOptions.extraHeaders,
suspensionBehavior: blobLoadOptions.suspensionBehavior,
})
},
`can't download from server `,
)
const mapWithEncryptedBlobs = parseMultipleBlobsResponse(concatBinaryData)
const mapWithDecryptedBlobs = new Map<Id, Uint8Array>()
for (const [blobId, encryptedData] of mapWithEncryptedBlobs) {
const decryptedData = aesDecrypt(sessionKey, encryptedData)
mapWithDecryptedBlobs.set(blobId, decryptedData)
}
return mapWithDecryptedBlobs
}

private async downloadAndDecryptChunkNative(blob: Blob, blobServerAccessInfo: BlobServerAccessInfo, sessionKey: AesKey): Promise<FileUri> {
const { archiveId, blobId } = blob
const getData = createBlobGetIn({
Expand Down Expand Up @@ -409,3 +483,26 @@ export class BlobFacade {
}
}
}

/**
* Deserializes a list of BlobWrappers that are in the following binary format
* element [ blobId ] [ blobHash ] [blobSize] [blob] [ . . . ] [ blobNId ] [ blobNHash ] [blobNSize] [blobN]
* bytes 9 6 4 blobSize 9 6 4 blobSize
*
* @return a map from blobId to the binary data
*/
export function parseMultipleBlobsResponse(concatBinaryData: Uint8Array): Map<Id, Uint8Array> {
let offset = 0
const dataView = new DataView(concatBinaryData.buffer)
const result = new Map<Id, Uint8Array>()
while (offset < concatBinaryData.length) {
const blobIdBytes = concatBinaryData.slice(offset, offset + 9)
const blobId = base64ToBase64Ext(uint8ArrayToBase64(blobIdBytes))

const blobSize = dataView.getInt32(offset + 15)
const contents = concatBinaryData.slice(offset + 19, offset + 19 + blobSize)
result.set(blobId, contents)
offset += 9 + 6 + 4 + blobSize
}
return result
}
142 changes: 137 additions & 5 deletions test/tests/api/worker/facades/BlobFacadeTest.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import o from "@tutao/otest"
import { BLOB_SERVICE_REST_PATH, BlobFacade } from "../../../../../src/common/api/worker/facades/lazy/BlobFacade.js"
import { BLOB_SERVICE_REST_PATH, BlobFacade, parseMultipleBlobsResponse } from "../../../../../src/common/api/worker/facades/lazy/BlobFacade.js"
import { RestClient } from "../../../../../src/common/api/worker/rest/RestClient.js"
import { SuspensionHandler } from "../../../../../src/common/api/worker/SuspensionHandler.js"
import { NativeFileApp } from "../../../../../src/common/native/common/FileApp.js"
Expand All @@ -12,7 +12,7 @@ import { ServiceExecutor } from "../../../../../src/common/api/worker/rest/Servi
import { instance, matchers, object, verify, when } from "testdouble"
import { HttpMethod } from "../../../../../src/common/api/common/EntityFunctions.js"
import { aes256RandomKey, aesDecrypt, aesEncrypt, generateIV } from "@tutao/tutanota-crypto"
import { arrayEquals, neverNull, stringToUtf8Uint8Array } from "@tutao/tutanota-utils"
import { arrayEquals, base64ExtToBase64, base64ToUint8Array, concat, neverNull, stringToUtf8Uint8Array } from "@tutao/tutanota-utils"
import { Mode } from "../../../../../src/common/api/common/Env.js"
import { CryptoFacade } from "../../../../../src/common/api/worker/crypto/CryptoFacade.js"
import { FileReference } from "../../../../../src/common/api/common/utils/FileUtils.js"
Expand Down Expand Up @@ -161,7 +161,8 @@ o.spec("BlobFacade test", function () {
o("downloadAndDecrypt", async function () {
const sessionKey = aes256RandomKey()
const blobData = new Uint8Array([1, 2, 3])
file.blobs.push(createTestEntity(BlobTypeRef))
const blobId = "--------0s--"
file.blobs.push(createTestEntity(BlobTypeRef, { blobId, size: String(65) }))
const encryptedBlobData = aesEncrypt(sessionKey, blobData, generateIV(), true, true)

let blobAccessInfo = createTestEntity(BlobServerAccessInfoTypeRef, {
Expand All @@ -176,18 +177,81 @@ o.spec("BlobFacade test", function () {
when(cryptoFacadeMock.resolveSessionKeyForInstance(file)).thenResolve(sessionKey)
const requestBody = { "request-body": true }
when(instanceMapperMock.encryptAndMapToLiteral(anything(), anything(), anything())).thenResolve(requestBody)
when(restClientMock.request(BLOB_SERVICE_REST_PATH, HttpMethod.GET, anything())).thenResolve(encryptedBlobData)
// data size is 65 (16 data block, 16 iv, 32 hmac, 1 byte for mac marking)
const blobSizeBinary = new Uint8Array([0, 0, 0, 65])
const blobResponse = concat(
// blob id
base64ToUint8Array(base64ExtToBase64(blobId)),
// blob hash
new Uint8Array([1, 2, 3, 4, 5, 6]),
// blob size
blobSizeBinary,
// blob data
encryptedBlobData,
)
when(restClientMock.request(BLOB_SERVICE_REST_PATH, HttpMethod.GET, anything())).thenResolve(blobResponse)

const decryptedData = await blobFacade.downloadAndDecrypt(archiveDataType, wrapTutanotaFile(file))

o(arrayEquals(decryptedData, blobData)).equals(true)("decrypted data")
o(decryptedData).deepEquals(blobData)("decrypted data is equal")
const optionsCaptor = captor()
verify(restClientMock.request(BLOB_SERVICE_REST_PATH, HttpMethod.GET, optionsCaptor.capture()))
o(optionsCaptor.value.baseUrl).equals("someBaseUrl")
o(optionsCaptor.value.queryParams.blobAccessToken).deepEquals(blobAccessInfo.blobAccessToken)
o(optionsCaptor.value.body).deepEquals(JSON.stringify(requestBody))
})

o("downloadAndDecrypt multiple", async function () {
const sessionKey = aes256RandomKey()
const blobData1 = new Uint8Array([1, 2, 3])
const blobId1 = "--------0s-1"
file.blobs.push(createTestEntity(BlobTypeRef, { blobId: blobId1, size: String(65) }))
const encryptedBlobData1 = aesEncrypt(sessionKey, blobData1, generateIV(), true, true)

const blobData2 = new Uint8Array([4, 5, 6, 7, 8, 9])
const blobId2 = "--------0s-2"
file.blobs.push(createTestEntity(BlobTypeRef, { blobId: blobId2, size: String(65) }))
const encryptedBlobData2 = aesEncrypt(sessionKey, blobData2, generateIV(), true, true)

const blobAccessInfo = createTestEntity(BlobServerAccessInfoTypeRef, {
blobAccessToken: "123",
servers: [createTestEntity(BlobServerUrlTypeRef, { url: "someBaseUrl" })],
})
when(blobAccessTokenFacade.requestReadTokenBlobs(anything(), anything(), matchers.anything())).thenResolve(blobAccessInfo)
when(blobAccessTokenFacade.createQueryParams(blobAccessInfo, anything(), anything())).thenResolve({
baseUrl: "someBaseUrl",
blobAccessToken: blobAccessInfo.blobAccessToken,
})
when(cryptoFacadeMock.resolveSessionKeyForInstance(file)).thenResolve(sessionKey)
const requestBody = { "request-body": true }
when(instanceMapperMock.encryptAndMapToLiteral(anything(), anything(), anything())).thenResolve(requestBody)
// data size is 65 (16 data block, 16 iv, 32 hmac, 1 byte for mac marking)
const blobSizeBinary = new Uint8Array([0, 0, 0, 65])
const blobResponse = concat(
// blob id
base64ToUint8Array(base64ExtToBase64(blobId1)),
// blob hash
new Uint8Array([1, 2, 3, 4, 5, 6]),
// blob size
blobSizeBinary,
// blob data
encryptedBlobData1,
// blob id
base64ToUint8Array(base64ExtToBase64(blobId2)),
// blob hash
new Uint8Array([6, 5, 4, 3, 2, 1]),
// blob size
blobSizeBinary,
// blob data
encryptedBlobData2,
)
when(restClientMock.request(BLOB_SERVICE_REST_PATH, HttpMethod.GET, anything())).thenResolve(blobResponse)

const decryptedData = await blobFacade.downloadAndDecrypt(archiveDataType, wrapTutanotaFile(file))

o(decryptedData).deepEquals(concat(blobData1, blobData2))("decrypted data is equal")
})

o("downloadAndDecryptNative", async function () {
const sessionKey = aes256RandomKey()

Expand Down Expand Up @@ -266,6 +330,74 @@ o.spec("BlobFacade test", function () {
verify(fileAppMock.deleteFile(decryptedChunkUri))
})
})

o.spec("parseMultipleBlobsResponse", function () {
o.test("parses two blobs", function () {
// Blob id OETv4XP----0 hash [3, -112, 88, -58, -14, -64] bytes [1, 2, 3]
// Blob id OETv4XS----0 hash [113, -110, 56, 92, 60, 6] bytes [1, 2, 3, 4, 5, 6]
const binaryData = new Int8Array([
// blob id 1 [0-8]
100, -9, -69, 22, 38, -128, 0, 0, 1,
// blob hash 1 [9-14]
3, -112, 88, -58, -14, -64,
// blob size 1 [15-18]
0, 0, 0, 3,
// blob data 1 [19-21]
1, 2, 3,
// blob id 2
100, -9, -69, 22, 39, 64, 0, 0, 1,
// blob hash 2
113, -110, 56, 92, 60, 6,
// blob size 2
0, 0, 0, 6,
// blob data 2
1, 2, 3, 4, 5, 6,
])

const result = parseMultipleBlobsResponse(new Uint8Array(binaryData))
o(result).deepEquals(
new Map([
["OETv4XP----0", new Uint8Array([1, 2, 3])],
["OETv4XS----0", new Uint8Array([1, 2, 3, 4, 5, 6])],
]),
)
})

o.test("parses one blob", function () {
// Blob id OETv4XP----0 hash [3, -112, 88, -58, -14, -64] bytes [1, 2, 3]
const binaryData = new Int8Array([
// blob id 1 [0-8]
100, -9, -69, 22, 38, -128, 0, 0, 1,
// blob hash 1 [9-14]
3, -112, 88, -58, -14, -64,
// blob size 1 [15-18]
0, 0, 0, 3,
// blob data 1 [19-21]
1, 2, 3,
])

const result = parseMultipleBlobsResponse(new Uint8Array(binaryData))
o(result).deepEquals(new Map([["OETv4XP----0", new Uint8Array([1, 2, 3])]]))
})

o.test("parses blob with big size", function () {
// Blob id OETv4XP----0 hash [3, -112, 88, -58, -14, -64] bytes [1, 2, 3]
const blobDataNumbers = Array(384).fill(1)
const binaryData = new Int8Array(
[
// blob id 1 [0-8]
100, -9, -69, 22, 38, -128, 0, 0, 1,
// blob hash 1 [9-14]
3, -112, 88, -58, -14, -64,
// blob size 1 [15-18] 384
0, 0, 1, 128,
].concat(blobDataNumbers),
)

const result = parseMultipleBlobsResponse(new Uint8Array(binaryData))
o(result).deepEquals(new Map([["OETv4XP----0", new Uint8Array(blobDataNumbers)]]))
})
})
})

function wrapTutanotaFile(tutanotaFile: TutanotaFile): BlobReferencingInstance {
Expand Down

0 comments on commit 5c02afc

Please sign in to comment.