Improved FILE_REFERENCE_EXPIRED handling

src/helpers/cancellablePromise.ts

@@ -5,7 +5,7 @@
  */
 
 export interface CancellablePromise<T> extends Promise<T> {
-  resolve?: (...args: any[]) => void,
+  resolve?: (value: T) => void,
   reject?: (...args: any[]) => void,
   cancel?: () => void,
 

src/layer.d.ts

@@ -342,7 +342,8 @@ export namespace InputFileLocation {
 		id: string,
 		access_hash: string,
 		file_reference: Uint8Array | number[],
-		thumb_size: string
+		thumb_size: string,
+		checkedReference?: boolean
 	};
 
 	export type inputSecureFileLocation = {

src/lib/appManagers/appDownloadManager.ts

@@ -75,8 +75,8 @@ export class AppDownloadManager {
     });
   }
 
-  private getNewDeferred(fileName: string) {
-    const deferred = deferredPromise<Blob>();
+  private getNewDeferred<T>(fileName: string) {
+    const deferred = deferredPromise<T>();
 
     deferred.cancel = () => {
       //try {
@@ -101,7 +101,7 @@ export class AppDownloadManager {
       this.clearDownload(fileName);
     });
 
-    return this.downloads[fileName] = deferred;
+    return this.downloads[fileName] = deferred as any;
   }
 
   private clearDownload(fileName: string) {
@@ -109,7 +109,7 @@ export class AppDownloadManager {
   }
 
   public fakeDownload(fileName: string, value: Blob | string) {
-    const deferred = this.getNewDeferred(fileName);
+    const deferred = this.getNewDeferred<Blob>(fileName);
     if(typeof(value) === 'string') {
       fetch(value)
       .then(response => response.blob())
@@ -125,28 +125,10 @@ export class AppDownloadManager {
     const fileName = getFileNameByLocation(options.location, {fileName: options.fileName});
     if(this.downloads.hasOwnProperty(fileName)) return this.downloads[fileName];
 
-    const deferred = this.getNewDeferred(fileName);
+    const deferred = this.getNewDeferred<Blob>(fileName);
 
     const onError = (err: ApiError) => {
-      switch(err.type) {
-        case 'FILE_REFERENCE_EXPIRED': {
-          // @ts-ignore
-          const bytes: ReferenceBytes = options?.location?.file_reference;
-          if(bytes) {
-            referenceDatabase.refreshReference(bytes).then(tryDownload);
-            /* referenceDatabase.refreshReference(bytes).then(() => {
-              console.log('FILE_REFERENCE_EXPIRED: refreshed reference', bytes);
-            }); */
-            break;
-          } else {
-            console.warn('FILE_REFERENCE_EXPIRED: no context for bytes:', bytes);
-          }
-        }
-
-        default:
-          deferred.reject(err);
-          break;
-      }
+      deferred.reject(err);
     };
 
     const tryDownload = (): Promise<unknown> => {
@@ -198,7 +180,7 @@ export class AppDownloadManager {
       }
     }
 
-    const deferred = this.getNewDeferred(fileName);
+    const deferred = this.getNewDeferred<InputFile>(fileName);
     apiManager.uploadFile({file, fileName}).then(deferred.resolve, deferred.reject);
 
     deferred.finally(() => {

src/lib/cacheStorage.ts

@@ -50,6 +50,7 @@ export default class CacheStorageController {
   }
 
   public save(entryName: string, response: Response) {
+    // return new Promise((resolve) => {}); // DEBUG
     return this.timeoutOperation((cache) => cache.put('/' + entryName, response));
   }
 

src/lib/mtproto/apiFileManager.ts

@@ -9,6 +9,7 @@
  * https://github.com/zhukov/webogram/blob/master/LICENSE
  */
 
+import type { ReferenceBytes } from "./referenceDatabase";
 import { MOUNT_CLASS_TO } from "../../config/debug";
 import Modes from "../../config/modes";
 import { readBlobAsArrayBuffer } from "../../helpers/blob";
@@ -17,18 +18,20 @@ import { notifyAll, notifySomeone } from "../../helpers/context";
 import { getFileNameByLocation } from "../../helpers/fileName";
 import { nextRandomInt } from "../../helpers/random";
 import { InputFile, InputFileLocation, UploadFile } from "../../layer";
-import { DcId } from "../../types";
+import { DcId, WorkerTaskVoidTemplate } from "../../types";
 import CacheStorageController from "../cacheStorage";
 import cryptoWorker from "../crypto/cryptoworker";
 import FileManager from "../filemanager";
 import { logger, LogTypes } from "../logger";
 import apiManager from "./apiManager";
 import { isWebpSupported } from "./mtproto.worker";
+import { bytesToHex } from "../../helpers/bytes";
+import assumeType from "../../helpers/assumeType";
 
 type Delayed = {
   offset: number, 
-  writeFilePromise: CancellablePromise<unknown>, 
-  writeFileDeferred: CancellablePromise<unknown>
+  writeFilePromise: CancellablePromise<void>, 
+  writeFileDeferred: CancellablePromise<void>
 };
 
 export type DownloadOptions = {
@@ -44,6 +47,17 @@ export type DownloadOptions = {
 
 type MyUploadFile = UploadFile.uploadFile;
 
+export interface RefreshReferenceTask extends WorkerTaskVoidTemplate {
+  type: 'refreshReference',
+  payload: ReferenceBytes,
+};
+
+export interface RefreshReferenceTaskResponse extends WorkerTaskVoidTemplate {
+  type: 'refreshReference',
+  payload: ReferenceBytes,
+  originalPayload: ReferenceBytes
+};
+
 const MAX_FILE_SAVE_SIZE = 20e6;
 
 export class ApiFileManager {
@@ -72,12 +86,24 @@ export class ApiFileManager {
   private downloadActives: {[dcId: string]: number} = {};
 
   public webpConvertPromises: {[fileName: string]: CancellablePromise<Uint8Array>} = {};
+  public refreshReferencePromises: {[referenceHex: string]: CancellablePromise<ReferenceBytes>} = {};
 
   private log: ReturnType<typeof logger> = logger('AFM', LogTypes.Error | LogTypes.Log);
   private tempId = 0;
   private queueId = 0;
   private debug = Modes.debug;
 
+  constructor() {
+    setInterval(() => { // clear old promises
+      for(const hex in this.refreshReferencePromises) {
+        const deferred = this.refreshReferencePromises[hex];
+        if(deferred.isFulfilled || deferred.isRejected) {
+          delete this.refreshReferencePromises[hex];
+        }
+      }
+    }, 1800e3);
+  }
+
   private downloadRequest(dcId: 'upload', id: number, cb: () => Promise<void>, activeDelta: number, queueId?: number): Promise<void>;
   private downloadRequest(dcId: number, id: number, cb: () => Promise<MyUploadFile>, activeDelta: number, queueId?: number): Promise<MyUploadFile>;
   private downloadRequest(dcId: number | string, id: number, cb: () => Promise<MyUploadFile | void>, activeDelta: number, queueId: number = 0) {
@@ -161,14 +187,36 @@ export class ApiFileManager {
     return this.downloadRequest(dcId, id, async() => {
       checkCancel && checkCancel();
 
-      return apiManager.invokeApi('upload.getFile', {
-        location,
-        offset,
-        limit
-      } as any, {
-        dcId,
-        fileDownload: true
-      }) as Promise<MyUploadFile>;
+      const invoke = (): Promise<MyUploadFile> => {
+        const promise = apiManager.invokeApi('upload.getFile', {
+          location,
+          offset,
+          limit
+        } as any, {
+          dcId,
+          fileDownload: true
+        }) as Promise<MyUploadFile>;
+
+        return promise.catch((err) => {
+          if(err.type === 'FILE_REFERENCE_EXPIRED') {
+            return this.refreshReference(location).then(() => invoke());
+          }
+
+          throw err;
+        });
+      };
+
+      assumeType<InputFileLocation.inputDocumentFileLocation>(location);
+      const reference = location.file_reference;
+      if(reference && !location.checkedReference) { // check stream's location because it's new every call
+        location.checkedReference = true;
+        const hex = bytesToHex(reference);
+        if(this.refreshReferencePromises[hex]) {
+          return this.refreshReference(location).then(() => invoke());
+        }
+      }
+
+      return invoke();
     }, this.getDelta(limit), queueId);
   }
 
@@ -205,6 +253,29 @@ export class ApiFileManager {
     return this.webpConvertPromises[fileName] = convertPromise;
   };
 
+  private refreshReference(inputFileLocation: InputFileLocation) {
+    const reference = (inputFileLocation as InputFileLocation.inputDocumentFileLocation).file_reference;
+    const hex = bytesToHex(reference);
+    let promise = this.refreshReferencePromises[hex];
+    const havePromise = !!promise;
+
+    if(!havePromise) {
+      promise = deferredPromise<ReferenceBytes>();
+    }
+    
+    promise.then(reference => {
+      (inputFileLocation as InputFileLocation.inputDocumentFileLocation).file_reference = reference;
+    });
+
+    if(havePromise) {
+      return promise;
+    }
+
+    const task = {type: 'refreshReference', payload: reference};
+    notifySomeone(task);
+    return this.refreshReferencePromises[hex] = promise;
+  }
+
   public downloadFile(options: DownloadOptions): CancellablePromise<Blob> {
     if(!FileManager.isAvailable()) {
       return Promise.reject({type: 'BROWSER_BLOB_NOT_SUPPORTED'});
@@ -293,8 +364,8 @@ export class ApiFileManager {
         const limit = options.limitPart || this.getLimitPart(size);
         let offset: number;
         let startOffset = 0;
-        let writeFilePromise: CancellablePromise<unknown> = Promise.resolve(),
-          writeFileDeferred: CancellablePromise<unknown>;
+        let writeFilePromise: CancellablePromise<void> = Promise.resolve(),
+          writeFileDeferred: CancellablePromise<void>;
         //const maxRequests = 13107200 / limit; // * 100 Mb speed
         const maxRequests = Infinity;
 

src/lib/mtproto/mtproto.worker.ts

@@ -10,7 +10,7 @@ import '../polyfill';
 import apiManager from "./apiManager";
 import cryptoWorker from "../crypto/cryptoworker";
 import networkerFactory from "./networkerFactory";
-import apiFileManager from './apiFileManager';
+import apiFileManager, { RefreshReferenceTaskResponse } from './apiFileManager';
 import type { RequestFilePartTask, RequestFilePartTaskResponse } from '../serviceWorker/index.service';
 import { ctx } from '../../helpers/userAgent';
 import { notifyAll } from '../../helpers/context';
@@ -21,6 +21,7 @@ import { LocalStorageProxyTask } from '../localStorage';
 import { WebpConvertTask } from '../webp/webpWorkerController';
 import { socketsProxied } from './transports/socketProxied';
 import { ToggleStorageTask } from './mtprotoworker';
+import { bytesToHex } from '../../helpers/bytes';
 
 let webpSupported = false;
 export const isWebpSupported = () => {
@@ -45,23 +46,6 @@ const taskListeners = {
     }
   },
 
-  requestFilePart: async(task: RequestFilePartTask) => {
-    const responseTask: RequestFilePartTaskResponse = {
-      type: task.type,
-      id: task.id
-    };
-
-    try {
-      const res = await apiFileManager.requestFilePart(...task.payload);
-      responseTask.payload = res;
-    } catch(err) {
-      responseTask.originalPayload = task.payload;
-      responseTask.error = err;
-    }
-
-    notifyAll(responseTask);
-  },
-
   webpSupport: (task: any) => {
     webpSupported = task.payload;
   },
@@ -101,6 +85,18 @@ const taskListeners = {
     const enabled = task.payload;
     // AppStorage.toggleStorage(enabled);
     CacheStorageController.toggleStorage(enabled);
+  },
+
+  refreshReference: (task: RefreshReferenceTaskResponse) => {
+    const hex = bytesToHex(task.originalPayload);
+    const deferred = apiFileManager.refreshReferencePromises[hex];
+    if(deferred) {
+      if(task.error) {
+        deferred.reject(task.error);
+      } else {
+        deferred.resolve(task.payload);
+      }
+    }
   }
 };
 
@@ -128,6 +124,7 @@ const onMessage = async(e: any) => {
           notifyAll({taskId, result});
         });
   
+      case 'requestFilePart':
       case 'setQueueId':
       case 'cancelDownload':
       case 'uploadFile':

src/lib/mtproto/mtprotoworker.ts

@@ -6,8 +6,8 @@
 
 import type { LocalStorageProxyTask, LocalStorageProxyTaskResponse } from '../localStorage';
 //import type { LocalStorageProxyDeleteTask, LocalStorageProxySetTask } from '../storage';
-import type { InvokeApiOptions, WorkerTaskVoidTemplate } from '../../types';
-import type { MethodDeclMap } from '../../layer';
+import type { Awaited, InvokeApiOptions, WorkerTaskVoidTemplate } from '../../types';
+import type { InputFile, MethodDeclMap } from '../../layer';
 import MTProtoWorker from 'worker-loader!./mtproto.worker';
 //import './mtproto.worker';
 import { isObject } from '../../helpers/object';
@@ -15,8 +15,8 @@ import CryptoWorkerMethods from '../crypto/crypto_methods';
 import { logger } from '../logger';
 import rootScope from '../rootScope';
 import webpWorkerController from '../webp/webpWorkerController';
-import type { DownloadOptions } from './apiFileManager';
-import type { ServiceWorkerTask } from '../serviceWorker/index.service';
+import { ApiFileManager, DownloadOptions } from './apiFileManager';
+import type { RequestFilePartTask, RequestFilePartTaskResponse, ServiceWorkerTask } from '../serviceWorker/index.service';
 import { UserAuth } from './mtproto_config';
 import type { MTMessage } from './networker';
 import DEBUG, { MOUNT_CLASS_TO } from '../../config/debug';
@@ -268,9 +268,23 @@ export class ApiManagerProxy extends CryptoWorkerMethods {
       }
     });
 
-    this.addServiceWorkerTaskListener('requestFilePart', (task) => {
-      this.postMessage(task);
+    this.addServiceWorkerTaskListener('requestFilePart', (task: RequestFilePartTask) => {
+      const responseTask: RequestFilePartTaskResponse = {
+        type: task.type,
+        id: task.id
+      };
+      
+      this.performTaskWorker<Awaited<ReturnType<ApiFileManager['requestFilePart']>>>('requestFilePart', ...task.payload)
+      .then((uploadFile) => {
+        responseTask.payload = uploadFile;
+        this.postSWMessage(responseTask);
+      }, (err) => {
+        responseTask.originalPayload = task.payload;
+        responseTask.error = err;
+        this.postSWMessage(responseTask);
+      });
     });
+
     /// #endif
 
     worker.addEventListener('messageerror', (e) => {
@@ -542,11 +556,11 @@ export class ApiManagerProxy extends CryptoWorkerMethods {
   }
 
   public downloadFile(options: DownloadOptions) {
-    return this.performTaskWorker('downloadFile', options);
+    return this.performTaskWorker<Blob>('downloadFile', options);
   }
 
   public uploadFile(options: {file: Blob | File, fileName: string}) {
-    return this.performTaskWorker('uploadFile', options);
+    return this.performTaskWorker<InputFile>('uploadFile', options);
   }
 
   public toggleStorage(enabled: boolean) {

src/lib/mtproto/referenceDatabase.ts

@@ -5,13 +5,15 @@
  */
 
 import type { RequestFilePartTask, RequestFilePartTaskResponse } from "../serviceWorker/index.service";
+import { RefreshReferenceTask, RefreshReferenceTaskResponse } from "./apiFileManager";
 import type { ApiError } from "./apiManager";
 import appMessagesManager from "../appManagers/appMessagesManager";
-import { Photo } from "../../layer";
+import { InputFileLocation, Photo } from "../../layer";
 import { bytesToHex } from "../../helpers/bytes";
 import { deepEqual } from "../../helpers/object";
 import { MOUNT_CLASS_TO } from "../../config/debug";
 import apiManager from "./mtprotoworker";
+import assumeType from "../../helpers/assumeType";
 
 export type ReferenceContext = ReferenceContext.referenceContextProfilePhoto | ReferenceContext.referenceContextMessage;
 export namespace ReferenceContext {
@@ -38,32 +40,19 @@ class ReferenceDatabase {
   private links: {[hex: string]: ReferenceBytes} = {};
 
   constructor() {
-    apiManager.addTaskListener('requestFilePart', (task: RequestFilePartTaskResponse) => {
-      if(task.error) {
-        const onError = (error: ApiError) => {
-          if(error?.type === 'FILE_REFERENCE_EXPIRED') {
-            // @ts-ignore
-            const bytes = task.originalPayload[1].file_reference;
-            referenceDatabase.refreshReference(bytes).then(() => {
-              // @ts-ignore
-              task.originalPayload[1].file_reference = referenceDatabase.getReferenceByLink(bytes);
-              const newTask: RequestFilePartTask = {
-                type: task.type,
-                id: task.id,
-                payload: task.originalPayload
-              };
-
-              apiManager.postMessage(newTask);
-            }).catch(onError);
-          } else {
-            navigator.serviceWorker.controller.postMessage(task);
-          }
-        };
-
-        onError(task.error);
-      } else {
-        navigator.serviceWorker.controller.postMessage(task);
-      }
+    apiManager.addTaskListener('refreshReference', (task: RefreshReferenceTask) => {
+      const bytes = task.payload;
+
+      assumeType<RefreshReferenceTaskResponse>(task);
+      task.originalPayload = bytes;
+
+      this.refreshReference(bytes).then(() => {
+        task.payload = this.getReferenceByLink(bytes);
+        apiManager.postMessage(task);
+      }, (err) => {
+        task.error = err;
+        apiManager.postMessage(task);
+      });
     });
   }
 

src/lib/storage.ts

@@ -161,7 +161,7 @@ export default class AppStorage<Storage extends Record<string, any>, T extends D
           const deferred = this.getPromises.get(key);
           if(deferred) {
             //deferred.reject(error);
-            deferred.resolve();
+            deferred.resolve(undefined);
             this.getPromises.delete(key);
           }
         }
@@ -196,8 +196,8 @@ export default class AppStorage<Storage extends Record<string, any>, T extends D
       const r = this.getPromises.get(key);
       if(r) return r as any;
 
-      const p = deferredPromise<Storage[typeof key]>();
-      this.getPromises.set(key, p);
+      const p = deferredPromise<Storage[T]>();
+      this.getPromises.set(key, p as any);
 
       this.getThrottled();
 
@@ -282,7 +282,7 @@ export default class AppStorage<Storage extends Record<string, any>, T extends D
       if(!enabled) {
         storage.keysToSet.clear();
         storage.keysToDelete.clear();
-        storage.getPromises.forEach((deferred) => deferred.resolve());
+        storage.getPromises.forEach((deferred) => deferred.resolve(undefined));
         storage.getPromises.clear();
         return storage.clear(true);
       } else {

src/scripts/in/schema_additional_params.json

@@ -258,4 +258,9 @@
     {"name": "hidden", "type": "true"},
     {"name": "fromId", "type": "number"}
   ]
+}, {
+  "predicate": "inputDocumentFileLocation",
+  "params": [
+    {"name": "checkedReference", "type": "boolean"}
+  ]
 }]