import { Message } from '@proton/shared/lib/interfaces/mail/Message' ; import { Api , LabelCount } from '@proton/shared/lib/interfaces' ; import { EVENT_ACTIONS } from '@proton/shared/lib/constants' ; import { range } from '@proton/shared/lib/helpers/array' ; import isTruthy from '@proton/shared/lib/helpers/isTruthy' ; import runInQueue from '@proton/shared/lib/helpers/runInQueue' ; import { IDBPDatabase } from 'idb' ; import { MessageEvent } from '../../models/event' ; import { GetMessageKeys } from '../../hooks/message/useGetMessageKeys' ; import { ESMessage , EncryptedSearchDB , ESBaseMessage , ESCache , NormalisedSearchParams , StoredCiphertext , } from '../../models/encryptedSearch' ; import { AesKeyGenParams , ES_MAX_CONCURRENT , ES_MAX_PAGES_PER_BATCH , ES_MAX_PARALLEL_MESSAGES , PAGE_SIZE , } from '../../constants' ; import { getNumMessagesDB , getOldestMessage , updateSizeIDB , getTotalMessages , refreshOpenpgp , openESDB , esSentryReport , getES , } from './esUtils' ; import { applySearch , normaliseSearchParams , uncachedSearch } from './esSearch' ; import { queryEvents , queryMessagesMetadata } from './esAPI' ; import { encryptToDB , fetchMessage , prepareMessageMetadata } from './esBuild' ; import { sizeOfCachedMessage , removeFromESCache , addToESCache , replaceInESCache } from './esCache' ; /**
 * Check whether the DB is limited, either after indexing or if it became so
 * after an update
 */ export const checkIsDBLimited = async ( userID , messageCounts ) => { const count = await getNumMessagesDB ( userID ) ; const totalMessages = await getTotalMessages ( messageCounts ) ; return count < totalMessages ; } ; /**
 * Decrypt encrypted object from IndexedDB
 */ export const decryptFromDB = async ( storedCiphertext , indexKey ) : Promise < ESMessage > => { const textDecoder = new TextDecoder ( ) ; const { aesGcmCiphertext } = storedCiphertext ; const decryptedMessage = await crypto . subtle . decrypt ( { iv : aesGcmCiphertext . iv , name : AesKeyGenParams . name } , indexKey , aesGcmCiphertext . ciphertext ) ; return JSON . parse ( textDecoder . decode ( new Uint8Array ( decryptedMessage ) ) ) ; } ; /**
 * Get message from IndexedDB
 */ export const getMessageFromDB = async ( ID , indexKey , esDB ) => { const storedCiphertext = await esDB . get ( 'messages' , ID ) ; if ( ! storedCiphertext ) { return ; } return decryptFromDB ( storedCiphertext , indexKey ) ; } ; /**
 * Stores a message to IndexedDB. If there is not enough space, older messages are
 * removed in favour of new ones
 */ export const storeToDB = async ( newCiphertextToStore , esDB ) => { const retryStoring = true ; while ( retryStoring ) { try { await esDB . put ( 'messages' , newCiphertextToStore ) ; return true ; } catch ( error ) { esSentryReport ( 'storeToDB: put failed' , { error } ) ; if ( error . name === 'QuotaExceededError' ) { const oldestMessage = await getOldestMessage ( esDB ) ; if ( ! oldestMessage ) { return false ; } const oldestTime = [ oldestMessage . Time , oldestMessage . Order ] ; const currentTime = [ newCiphertextToStore . Time , newCiphertextToStore . Order ] ; if ( currentTime < oldestTime ) { return true ; } await esDB . delete ( 'messages' , oldestMessage . ID ) ; } else { return false ; } } } } ; /**
 * Check whether two ESBaseMessage objects are equal
 */ export const compareESBaseMessages = ( message1 , message2 ) => { let key1 ; for ( key1 in message1 ) { if ( JSON . stringify ( message1 [ key1 ] ) !== JSON . stringify ( message2 [ key1 ] ) ) { return false ; } } let key2 ; for ( key2 in message2 ) { if ( JSON . stringify ( message2 [ key2 ] ) !== JSON . stringify ( message1 [ key2 ] ) ) { return false ; } } return true ; } ; /**
 * Remove messages from and add messages to IDB
 */ export const executeIDBOperations = async ( esDB , messagesToRemove , messagesToAdd ) => { const tx = esDB . transaction ( 'messages' , 'readwrite' ) ; if ( messagesToRemove . length ) { for ( const ID of messagesToRemove ) { void tx . store . delete ( ID ) ; } } try { if ( messagesToAdd . length ) { for ( const ciphertext of messagesToAdd ) { void tx . store . put ( ciphertext ) ; } } await tx . done ; } catch ( error ) { if ( error . name === 'QuotaExceededError' ) { for ( const ciphertext of messagesToAdd ) { if ( ! ( await storeToDB ( ciphertext , esDB ) ) ) { throw new Error ( 'Sync of some messages failed' ) ; } } } else { throw error ; } } } ; /**
 * Synchronise IDB (and optionally cache and search results) with new message events
 */ export const syncMessageEvents = async ( Messages , userID , esCacheRef , permanentResults , isSearch , api , getMessageKeys , indexKey , normalisedSearchParams , recordProgressLocal ? ) => { const esDB = await openESDB ( userID ) ; let searchChanged = false ; const updatePermanentResults = ( { resultIndex , messageToCache , } ) => { if ( messageToCache ) { if ( resultIndex ) { permanentResults . splice ( resultIndex , 1 , messageToCache ) ; } else { permanentResults . push ( messageToCache ) ; } } else { permanentResults . splice ( resultIndex as number , 1 ) ; } searchChanged = true ; } ; const messagesToRemove = [ ] ; const messagesToAdd = [ ] ; for ( let batch = 0 ; batch < Messages . length ; batch += ES_MAX_PARALLEL_MESSAGES ) { const messageEventsBatch = Messages . slice ( batch , Math . min ( Messages . length , batch + ES_MAX_PARALLEL_MESSAGES ) ) ; const prefetchedMessages = ( await Promise . all ( messageEventsBatch . map ( async ( messageEvent ) => { const { ID , Action } = messageEvent ; if ( Action === EVENT_ACTIONS . CREATE || Action === EVENT_ACTIONS . UPDATE_DRAFT ) { return fetchMessage ( ID , api , getMessageKeys ) ; } } ) ) ) . filter ( isTruthy ) ; for ( const messageEvent of messageEventsBatch ) { const { ID , Action , Message } = messageEvent ; if ( Action === EVENT_ACTIONS . DELETE ) { messagesToRemove . push ( ID ) ; const size = removeFromESCache ( ID , esCacheRef ) || 0 ; updateSizeIDB ( userID , - size ) ; const resultIndex = permanentResults . findIndex ( ( message ) => message . ID === ID ) ; if ( isSearch && resultIndex !== - 1 ) { updatePermanentResults ( { resultIndex } ) ; } } if ( Action === EVENT_ACTIONS . CREATE ) { if ( ! Message ) { continue ; } const messageToCache = prefetchedMessages . find ( ( prefetchedMessage ) => prefetchedMessage . ID === ID ) ; if ( ! messageToCache ) { continue ; } const newCiphertextToStore = await encryptToDB ( messageToCache , indexKey ) ; messagesToAdd . push ( newCiphertextToStore ) ; const size = sizeOfCachedMessage ( messageToCache ) ; updateSizeIDB ( userID , size ) ; addToESCache ( messageToCache , esCacheRef , size ) ; if ( isSearch && applySearch ( normalisedSearchParams , messageToCache ) ) { updatePermanentResults ( { messageToCache } ) ; } } if ( Action === EVENT_ACTIONS . UPDATE_DRAFT || Action === EVENT_ACTIONS . UPDATE_FLAGS ) { if ( ! Message ) { continue ; } const oldMessage = await getMessageFromDB ( ID , indexKey , esDB ) ; if ( ! oldMessage ) { continue ; } let newMessageToCache ; if ( Action === EVENT_ACTIONS . UPDATE_DRAFT ) { const fetchedMessageToCache = prefetchedMessages . find ( ( prefetchedMessage ) => prefetchedMessage . ID === ID ) ; if ( ! fetchedMessageToCache ) { continue ; } newMessageToCache = fetchedMessageToCache ; } if ( Action === EVENT_ACTIONS . UPDATE_FLAGS ) { const { LabelIDsRemoved , LabelIDsAdded , ... otherChanges } = Message ; let { LabelIDs } = oldMessage ; if ( LabelIDsRemoved ) { LabelIDs = LabelIDs . filter ( ( labelID ) => ! LabelIDsRemoved . includes ( labelID ) ) ; } if ( LabelIDsAdded ) { LabelIDs = LabelIDs . concat ( LabelIDsAdded ) ; } newMessageToCache = { ... oldMessage , ... otherChanges , LabelIDs , } ; } if ( ! newMessageToCache ) { throw new Error ( 'Plaintext to store is undefined' ) ; } const newCiphertextToStore = await encryptToDB ( newMessageToCache , indexKey ) ; messagesToAdd . push ( newCiphertextToStore ) ; const sizeDelta = sizeOfCachedMessage ( newMessageToCache ) - sizeOfCachedMessage ( oldMessage ) ; updateSizeIDB ( userID , sizeDelta ) ; replaceInESCache ( newMessageToCache , esCacheRef , Action === EVENT_ACTIONS . UPDATE_DRAFT , sizeDelta ) ; if ( isSearch ) { if ( applySearch ( normalisedSearchParams , oldMessage ) ) { const resultIndex = permanentResults . findIndex ( ( message ) => message . ID === ID ) ; if ( applySearch ( normalisedSearchParams , newMessageToCache ) ) { updatePermanentResults ( { resultIndex , messageToCache : newMessageToCache } ) ; } else { updatePermanentResults ( { resultIndex } ) ; } } else if ( applySearch ( normalisedSearchParams , newMessageToCache ) ) { updatePermanentResults ( { messageToCache : newMessageToCache } ) ; } } } if ( recordProgressLocal ) { recordProgressLocal ( ) ; } } } await executeIDBOperations ( esDB , messagesToRemove , messagesToAdd ) ; esDB . close ( ) ; return searchChanged ; } ; /**
 * When an old key is activated, try to correct any previous decryption errors
 */ export const correctDecryptionErrors = async ( userID , indexKey , api , getMessageKeys , recordProgress , esCacheRef ? ) => { const { resultsArray : searchResults } = await uncachedSearch ( userID , indexKey , { ... normaliseSearchParams ( { } , '5' ) , decryptionError : true } , { } ) ; if ( ! searchResults . length ) { return false ; } recordProgress ( 0 , searchResults . length ) ; let counter = 0 ; const messagesToAdd = ( await Promise . all ( searchResults . map ( async ( message ) => { const newMessage = await fetchMessage ( message . ID , api , getMessageKeys ) ; if ( ! newMessage || newMessage . decryptionError ) { return ; } const newCiphertextToStore = await encryptToDB ( newMessage , indexKey ) ; const size = sizeOfCachedMessage ( newMessage ) ; updateSizeIDB ( userID , size ) ; if ( esCacheRef ) { addToESCache ( newMessage , esCacheRef , size ) ; } recordProgress ( ++ counter , searchResults . length ) ; return newCiphertextToStore ; } ) ) ) . filter ( isTruthy ) ; const newMessagesFound = messagesToAdd . length ; if ( newMessagesFound ) { const esDB = await openESDB ( userID ) ; await executeIDBOperations ( esDB , [ ] , messagesToAdd ) ; esDB . close ( ) ; } return newMessagesFound ; } ; /**
 * Fetch, prepare and store the specified messages to IndexedDB
 */ const syncMessagesBatch = async ( userID , messagesMetadata , esDB , indexKey , api , getMessageKeys , recordLocalProgress ) => { let counter = 0 ; const esIteratee = async ( message ) => { const messageToCache = ( await fetchMessage ( message . ID , api , getMessageKeys , undefined , message ) ) ! ; const size = sizeOfCachedMessage ( messageToCache ) ; updateSizeIDB ( userID , size ) ; recordLocalProgress ( ++ counter ) ; return encryptToDB ( messageToCache , indexKey ) ; } ; const ciphertexts = await runInQueue < StoredCiphertext > ( messagesMetadata . map ( ( message ) => ( ) => esIteratee ( message ) ) , ES_MAX_CONCURRENT ) ; const tx = esDB . transaction ( 'messages' , 'readwrite' ) ; await Promise . all ( ciphertexts . map ( async ( ciphertext ) => tx . store . put ( ciphertext ) ) ) ; await tx . done ; } ; /**
 * Refresh index when we get a Refresh=1 from queryEvents
 */ export const refreshIndex = async ( userID , api , indexKey , getMessageKeys , recordProgress , messageCounts ) => { const eventSinceRefresh = await queryEvents ( api ) ; if ( ! eventSinceRefresh || ! eventSinceRefresh . EventID ) { throw new Error ( 'Event fetching failed' ) ; } await correctDecryptionErrors ( userID , indexKey , api , getMessageKeys , recordProgress ) ; recordProgress ( 0 , 0 ) ; const esDB = await openESDB ( userID ) ; const Total = await getTotalMessages ( messageCounts ) ; const numPages = Math . ceil ( Total / PAGE_SIZE ) ; const numBatches = Math . ceil ( numPages / ES_MAX_PAGES_PER_BATCH ) ; let numMessages = 0 ; recordProgress ( 0 , Total ) ; let drafts = [ ] ; try { ( { resultsArray : drafts } = await uncachedSearch ( userID , indexKey , normaliseSearchParams ( { } , '8' ) , { } ) ) ; } catch ( error ) { throw new Error ( 'Drafts fetching failed' ) ; } const indexedIDs = new Set ( await esDB . getAllKeys ( 'messages' ) ) ; for ( let startPageBatch = 0 ; startPageBatch < numBatches ; startPageBatch ++ ) { const startPage = startPageBatch * ES_MAX_PAGES_PER_BATCH ; const endPage = Math . min ( ( startPageBatch + 1 ) * ES_MAX_PAGES_PER_BATCH , numPages ) ; const fetchedMetadata = [ ] ; const pagesRange = range ( startPage , endPage ) ; const error = await Promise . all ( pagesRange . map ( async ( Page ) => { const resultMetadata = await queryMessagesMetadata ( api , { Page , Limit : PAGE_SIZE } ) ; if ( ! resultMetadata ) { throw new Error ( ) ; } fetchedMetadata . push ( ... resultMetadata . Messages ) ; } ) ) . then ( ( ) => false ) . catch ( ( ) => true ) ; if ( error ) { throw new Error ( 'Metadata fetching failed' ) ; } const messagesToAdd = [ ] ; const messagesToFetch = [ ] ; for ( const message of fetchedMetadata ) { const { ID } = message ; if ( indexedIDs . has ( ID ) ) { const draft = drafts . find ( ( { ID } ) => message . ID === ID ) ; if ( draft ) { if ( draft . Time !== message . Time ) { messagesToFetch . push ( message ) ; } indexedIDs . delete ( ID ) ; continue ; } const oldMessage = await getMessageFromDB ( ID , indexKey , esDB ) ; if ( ! oldMessage ) { throw new Error ( 'Old message is undefined' ) ; } const { decryptionError , decryptedBody , decryptedSubject } = oldMessage ; const oldBaseMessage = prepareMessageMetadata ( oldMessage ) ; const newBaseMessage = prepareMessageMetadata ( message ) ; if ( ! compareESBaseMessages ( oldBaseMessage , newBaseMessage ) ) { const newMessageToCache = { decryptionError , decryptedBody , decryptedSubject , ... newBaseMessage , } ; const newCiphertextToStore = await encryptToDB ( newMessageToCache , indexKey ) ; const sizeDelta = sizeOfCachedMessage ( newMessageToCache ) - sizeOfCachedMessage ( oldMessage ) ; updateSizeIDB ( userID , sizeDelta ) ; messagesToAdd . push ( newCiphertextToStore ) ; } indexedIDs . delete ( ID ) ; recordProgress ( numMessages ++ , Total ) ; } else { messagesToFetch . push ( message ) ; } } if ( messagesToAdd . length ) { await executeIDBOperations ( esDB , [ ] , messagesToAdd ) ; } if ( messagesToFetch . length ) { const recordLocalProgress = ( numMessages ) => ( progressLocal ) => recordProgress ( numMessages + progressLocal , Total ) ; await syncMessagesBatch ( userID , messagesToFetch , esDB , indexKey , api , getMessageKeys , recordLocalProgress ( numMessages ) ) ; await refreshOpenpgp ( ) ; } } const messagesToRemove = [ ... indexedIDs ] ; await executeIDBOperations ( esDB , messagesToRemove , [ ] ) ; esDB . close ( ) ; const newMessageEvent = await queryEvents ( api , eventSinceRefresh . EventID ) ; if ( ! newMessageEvent || ! newMessageEvent . EventID ) { throw new Error ( 'Failed to fetch events after refresh' ) ; } return newMessageEvent ; } ; /**
 * Return the (possibly partial) list of events since the one stored in localStorage
 */ export const getEventFromLS = async ( userID , api ) => { const storedEventID = getES . Event ( userID ) ; if ( ! storedEventID ) { return ; } return queryEvents ( api , storedEventID ) ; } ;