Media: Fix repairing file references (#4804)

src/api/gramjs/apiBuilders/messageContent.ts

@@ -22,13 +22,16 @@ import type {
   ApiWebPage,
   ApiWebPageStickerData,
   ApiWebPageStoryData,
+  BoughtPaidMedia,
   MediaContent,
 } from '../../types';
 import type { UniversalMessage } from './messages';
 
 import { SUPPORTED_IMAGE_CONTENT_TYPES, SUPPORTED_VIDEO_CONTENT_TYPES, VIDEO_WEBM_TYPE } from '../../../config';
 import { pick } from '../../../util/iteratees';
-import { addMediaToLocalDb, addStoryToLocalDb, serializeBytes } from '../helpers';
+import {
+  addMediaToLocalDb, addStoryToLocalDb, type MediaRepairContext, serializeBytes,
+} from '../helpers';
 import {
   buildApiFormattedText,
   buildApiMessageEntity,
@@ -46,8 +49,9 @@ export function buildMessageContent(
   let content: MediaContent = {};
 
   if (mtpMessage.media) {
+    const repairContext = 'peerId' in mtpMessage ? mtpMessage : undefined;
     content = {
-      ...buildMessageMediaContent(mtpMessage.media),
+      ...buildMessageMediaContent(mtpMessage.media, repairContext),
     };
   }
 
@@ -74,8 +78,10 @@ export function buildMessageTextContent(
   };
 }
 
-export function buildMessageMediaContent(media: GramJs.TypeMessageMedia): MediaContent | undefined {
+export function buildMessageMediaContent(
-  addMediaToLocalDb(media);
+  media: GramJs.TypeMessageMedia, context?: MediaRepairContext,
+): MediaContent | undefined {
+  addMediaToLocalDb(media, context);
 
   const ttlSeconds = 'ttlSeconds' in media ? media.ttlSeconds : undefined;
 
@@ -102,7 +108,7 @@ export function buildMessageMediaContent(media: GramJs.TypeMessageMedia): MediaC
   }
 
   if (media instanceof GramJs.MessageMediaInvoice && media.extendedMedia instanceof GramJs.MessageExtendedMedia) {
-    return buildMessageMediaContent(media.extendedMedia.media);
+    return buildMessageMediaContent(media.extendedMedia.media, context);
   }
 
   const sticker = buildSticker(media);
@@ -785,12 +791,7 @@ function buildPaidMedia(media: GramJs.TypeMessageMedia): ApiPaidMedia | undefine
       mediaType: 'paidMedia',
       starsAmount: starsAmount.toJSNumber(),
       isBought,
-      extendedMedia: extendedMedia
+      extendedMedia: buildBoughtMediaContent(extendedMedia)!,
-        .filter((paidMedia): paidMedia is GramJs.MessageExtendedMedia => (
-          paidMedia instanceof GramJs.MessageExtendedMedia
-        ))
-        .map((paidMedia) => buildMessageMediaContent(paidMedia.media))
-        .filter(Boolean),
     };
   }
 
@@ -857,7 +858,9 @@ export function buildApiWebDocument(document?: GramJs.TypeWebDocument): ApiWebDo
   };
 }
 
-export function buildBoughtMediaContent(media: GramJs.TypeMessageExtendedMedia[]): MediaContent[] | undefined {
+export function buildBoughtMediaContent(
+  media: GramJs.TypeMessageExtendedMedia[],
+): BoughtPaidMedia[] | undefined {
   const boughtMedia = media
     .filter((m): m is GramJs.MessageExtendedMedia => m instanceof GramJs.MessageExtendedMedia)
     .map((m) => buildMessageMediaContent(m.media))

src/api/gramjs/apiBuilders/messages.ts

@@ -49,6 +49,7 @@ import { interpolateArray } from '../../../util/waveform';
 import { buildPeer } from '../gramjsBuilders';
 import {
   addPhotoToLocalDb,
+  type MediaRepairContext,
   resolveMessageApiChatId,
   serializeBytes,
 } from '../helpers';
@@ -122,6 +123,7 @@ export function buildApiMessageFromShort(mtpMessage: GramJs.UpdateShortMessage):
   return buildApiMessageWithChatId(chatId, {
     ...mtpMessage,
     fromId: buildPeer(mtpMessage.out ? currentUserId : buildApiPeerId(mtpMessage.userId, 'user')),
+    peerId: buildPeer(mtpMessage.out ? buildApiPeerId(mtpMessage.userId, 'user') : currentUserId),
   });
 }
 
@@ -131,6 +133,7 @@ export function buildApiMessageFromShortChat(mtpMessage: GramJs.UpdateShortChatM
   return buildApiMessageWithChatId(chatId, {
     ...mtpMessage,
     fromId: buildPeer(buildApiPeerId(mtpMessage.fromId, 'user')),
+    peerId: buildPeer(buildApiPeerId(mtpMessage.chatId, 'chat')),
   });
 }
 
@@ -151,7 +154,7 @@ export function buildApiMessageFromNotification(
 }
 
 export type UniversalMessage = (
-  Pick<GramJs.Message & GramJs.MessageService, ('id' | 'date')>
+  Pick<GramJs.Message & GramJs.MessageService, ('id' | 'date' | 'peerId')>
   & Partial<GramJs.Message & GramJs.MessageService>
 );
 
@@ -211,7 +214,7 @@ export function buildApiMessageWithChatId(
     isPinned: mtpMessage.pinned,
     reactions: mtpMessage.reactions && buildMessageReactions(mtpMessage.reactions),
     emojiOnlyCount,
-    ...(mtpMessage.replyTo && { replyInfo: buildApiReplyInfo(mtpMessage.replyTo) }),
+    ...(mtpMessage.replyTo && { replyInfo: buildApiReplyInfo(mtpMessage.replyTo, mtpMessage) }),
     forwardInfo,
     isEdited,
     editDate: mtpMessage.editDate,
@@ -286,7 +289,9 @@ function buildApiMessageForwardInfo(fwdFrom: GramJs.MessageFwdHeader, isChatWith
   };
 }
 
-function buildApiReplyInfo(replyHeader: GramJs.TypeMessageReplyHeader): ApiReplyInfo | undefined {
+function buildApiReplyInfo(
+  replyHeader: GramJs.TypeMessageReplyHeader, context?: MediaRepairContext,
+): ApiReplyInfo | undefined {
   if (replyHeader instanceof GramJs.MessageReplyStoryHeader) {
     return {
       type: 'story',
@@ -315,7 +320,7 @@ function buildApiReplyInfo(replyHeader: GramJs.TypeMessageReplyHeader): ApiReply
       isForumTopic: forumTopic,
       replyFrom: replyFrom && buildApiMessageForwardInfo(replyFrom),
       replyToPeerId: replyToPeerId && getApiChatIdFromMtpPeer(replyToPeerId),
-      replyMedia: replyMedia && buildMessageMediaContent(replyMedia),
+      replyMedia: replyMedia && buildMessageMediaContent(replyMedia, context),
       isQuote: quote,
       quoteText: quoteText ? buildMessageTextContent(quoteText, quoteEntities) : undefined,
     };

src/api/gramjs/helpers.ts

@@ -20,6 +20,9 @@ const LOG_SUFFIX = {
   'UNEXPECTED RESPONSE': '#D1191C',
 };
 
+export type MessageRepairContext = Pick<GramJs.TypeMessage, 'peerId' | 'id'>;
+export type MediaRepairContext = MessageRepairContext;
+
 export function resolveMessageApiChatId(mtpMessage: GramJs.TypeMessage) {
   if (!(mtpMessage instanceof GramJs.Message || mtpMessage instanceof GramJs.MessageService)) {
     return undefined;
@@ -53,9 +56,9 @@ export function addMessageToLocalDb(message: GramJs.TypeMessage | GramJs.TypeSpo
   }
 }
 
-export function addMediaToLocalDb(media: GramJs.TypeMessageMedia, message?: GramJs.TypeMessage) {
+export function addMediaToLocalDb(media: GramJs.TypeMessageMedia, context?: MediaRepairContext) {
   if (media instanceof GramJs.MessageMediaDocument && media.document) {
-    const document = addMessageRepairInfo(media.document, message);
+    const document = addMessageRepairInfo(media.document, context);
     addDocumentToLocalDb(document);
   }
 
@@ -63,45 +66,45 @@ export function addMediaToLocalDb(media: GramJs.TypeMessageMedia, message?: Gram
     && media.webpage instanceof GramJs.WebPage
   ) {
     if (media.webpage.document) {
-      const document = addMessageRepairInfo(media.webpage.document, message);
+      const document = addMessageRepairInfo(media.webpage.document, context);
       addDocumentToLocalDb(document);
     }
     if (media.webpage.photo) {
-      const photo = addMessageRepairInfo(media.webpage.photo, message);
+      const photo = addMessageRepairInfo(media.webpage.photo, context);
       addPhotoToLocalDb(photo);
     }
   }
 
   if (media instanceof GramJs.MessageMediaGame) {
     if (media.game.document) {
-      const document = addMessageRepairInfo(media.game.document, message);
+      const document = addMessageRepairInfo(media.game.document, context);
       addDocumentToLocalDb(document);
     }
 
-    const photo = addMessageRepairInfo(media.game.photo, message);
+    const photo = addMessageRepairInfo(media.game.photo, context);
     addPhotoToLocalDb(photo);
   }
 
   if (media instanceof GramJs.MessageMediaPhoto && media.photo) {
-    const photo = addMessageRepairInfo(media.photo, message);
+    const photo = addMessageRepairInfo(media.photo, context);
     addPhotoToLocalDb(photo);
   }
 
   if (media instanceof GramJs.MessageMediaInvoice) {
     if (media.photo) {
-      const photo = addMessageRepairInfo(media.photo, message);
+      const photo = addMessageRepairInfo(media.photo, context);
       addWebDocumentToLocalDb(photo);
     }
 
     if (media.extendedMedia instanceof GramJs.MessageExtendedMedia) {
-      addMediaToLocalDb(media.extendedMedia.media, message);
+      addMediaToLocalDb(media.extendedMedia.media, context);
     }
   }
 
   if (media instanceof GramJs.MessageMediaPaidMedia) {
     media.extendedMedia.forEach((extendedMedia) => {
       if (extendedMedia instanceof GramJs.MessageExtendedMedia) {
-        addMediaToLocalDb(extendedMedia.media, message);
+        addMediaToLocalDb(extendedMedia.media, context);
       }
     });
   }
@@ -156,18 +159,18 @@ export function addStoryRepairInfo<T extends GramJs.TypeDocument | GramJs.TypeWe
 }
 
 export function addMessageRepairInfo<T extends GramJs.TypeDocument | GramJs.TypeWebDocument | GramJs.TypePhoto>(
-  media: T, message?: GramJs.TypeMessage,
+  media: T, context?: MediaRepairContext,
 ) : T & RepairInfo {
-  if (!message?.peerId) return media;
+  if (!context?.peerId) return media;
-  if (!(media instanceof GramJs.Document && media instanceof GramJs.Photo && media instanceof GramJs.WebDocument)) {
+  if (!(media instanceof GramJs.Document || media instanceof GramJs.Photo || media instanceof GramJs.WebDocument)) {
     return media;
   }
 
   const repairableMedia = media as T & RepairInfo;
   repairableMedia.localRepairInfo = {
     type: 'message',
-    peerId: getApiChatIdFromMtpPeer(message.peerId),
+    peerId: getApiChatIdFromMtpPeer(context.peerId),
-    id: message.id,
+    id: context.id,
   };
   return repairableMedia;
 }

src/api/gramjs/localDb.ts

@@ -3,7 +3,7 @@ import { constructors } from '../../lib/gramjs/tl';
 
 import type { Api as GramJs } from '../../lib/gramjs';
 
-import { DATA_BROADCAST_CHANNEL_NAME } from '../../config';
+import { DATA_BROADCAST_CHANNEL_NAME, DEBUG } from '../../config';
 import { throttle } from '../../util/schedulers';
 import { omitVirtualClassFields } from './apiBuilders/helpers';
 
@@ -141,3 +141,7 @@ export function updateFullLocalDb(initial: LocalDb) {
 export function clearLocalDb() {
   Object.assign(localDb, createLocalDbInitial());
 }
+
+if (DEBUG) {
+  (globalThis as any).getLocalDb = () => localDb;
+}

src/api/gramjs/methods/client.ts

@@ -519,10 +519,10 @@ async function repairStoryMedia(peerId: string, storyId: number) {
 
   addEntitiesToLocalDb(result.users);
   result.stories.forEach((story) => {
-    addStoryToLocalDb(story, peerId);
-
     const apiStory = buildApiStory(peerId, story);
     if (!apiStory || 'isDeleted' in apiStory) return;
+
+    addStoryToLocalDb(story, peerId);
     onUpdate({
       '@type': 'updateStory',
       peerId,

src/api/gramjs/methods/messages.ts

@@ -74,6 +74,7 @@ import {
   buildInputTextWithEntities,
   buildMessageFromUpdate,
   buildMtpMessageEntity,
+  buildPeer,
   buildSendMessageAction,
   generateRandomBigInt,
   getEntityTypeById,
@@ -1911,7 +1912,9 @@ function handleLocalMessageUpdate(localMessage: ApiMessage, update: GramJs.TypeU
     if (messageUpdate.media) {
       newContent = {
         ...newContent,
-        ...buildMessageMediaContent(messageUpdate.media),
+        ...buildMessageMediaContent(messageUpdate.media, {
+          peerId: buildPeer(localMessage.chatId), id: messageUpdate.id,
+        }),
       };
     }
 

src/api/gramjs/methods/stories.ts

@@ -70,9 +70,6 @@ export async function fetchAllStories({
 
   addEntitiesToLocalDb(result.users);
   addEntitiesToLocalDb(result.chats);
-  result.peerStories.forEach((peerStories) => (
-    peerStories.stories.forEach((story) => addStoryToLocalDb(story, getApiChatIdFromMtpPeer(peerStories.peer)))
-  ));
 
   const allUserStories = result.peerStories.reduce<Record<string, ApiPeerStories>>((acc, peerStories) => {
     const peerId = getApiChatIdFromMtpPeer(peerStories.peer);
@@ -114,6 +111,11 @@ export async function fetchAllStories({
     return acc;
   }, {});
 
+  // Add after building stories to avoid overwriting repair info
+  result.peerStories.forEach((peerStories) => (
+    peerStories.stories.forEach((story) => addStoryToLocalDb(story, getApiChatIdFromMtpPeer(peerStories.peer)))
+  ));
+
   return {
     users: result.users.map(buildApiUser).filter(Boolean),
     chats: result.chats.map((c) => buildApiChatFromPreview(c)).filter(Boolean),
@@ -138,7 +140,6 @@ export async function fetchPeerStories({
   }
 
   addEntitiesToLocalDb(result.users);
-  result.stories.stories.forEach((story) => addStoryToLocalDb(story, peer.id));
 
   const users = result.users.map(buildApiUser).filter(Boolean);
   const chats = result.chats.map((c) => buildApiChatFromPreview(c)).filter(Boolean);
@@ -146,6 +147,9 @@ export async function fetchPeerStories({
     [story.id, buildApiStory(peer.id, story)]
   ));
 
+  // Add after building stories to avoid overwriting repair info
+  result.stories.stories.forEach((story) => addStoryToLocalDb(story, peer.id));
+
   return {
     chats,
     users,
@@ -199,7 +203,6 @@ export async function fetchPeerStoriesByIds({ peer, ids }: { peer: ApiPeer; ids:
 
   addEntitiesToLocalDb(result.users);
   addEntitiesToLocalDb(result.chats);
-  result.stories.forEach((story) => addStoryToLocalDb(story, peer.id));
 
   const users = result.users.map(buildApiUser).filter(Boolean);
   const chats = result.chats.map((c) => buildApiChatFromPreview(c)).filter(Boolean);
@@ -218,6 +221,9 @@ export async function fetchPeerStoriesByIds({ peer, ids }: { peer: ApiPeer; ids:
     return acc;
   }, {});
 
+  // Add after building stories to avoid overwriting repair info
+  result.stories.forEach((story) => addStoryToLocalDb(story, peer.id));
+
   return {
     chats,
     users,
@@ -426,7 +432,6 @@ async function fetchCommonStoriesRequest({ method, peerId }: {
 
   addEntitiesToLocalDb(result.users);
   addEntitiesToLocalDb(result.chats);
-  result.stories.forEach((story) => addStoryToLocalDb(story, peerId));
 
   const users = result.users.map(buildApiUser).filter(Boolean);
   const chats = result.chats.map((c) => buildApiChatFromPreview(c)).filter(Boolean);
@@ -434,6 +439,9 @@ async function fetchCommonStoriesRequest({ method, peerId }: {
     [story.id, buildApiStory(peerId, story)]
   ));
 
+  // Add after building stories to avoid overwriting repair info
+  result.stories.forEach((story) => addStoryToLocalDb(story, peerId));
+
   return {
     users,
     chats,

src/api/gramjs/updates/updater.ts

@@ -1097,7 +1097,8 @@ export function updater(update: Update) {
 
     const { story } = update;
     const peerId = getApiChatIdFromMtpPeer(update.peer);
-    addStoryToLocalDb(story, peerId);
+    const apiStory = buildApiStory(peerId, story) as ApiStory | ApiStorySkipped;
+    addStoryToLocalDb(story, peerId); // Add after building to prevent repair info overwrite
 
     if (story instanceof GramJs.StoryItemDeleted) {
       onUpdate({
@@ -1109,7 +1110,7 @@ export function updater(update: Update) {
       onUpdate({
         '@type': 'updateStory',
         peerId,
-        story: buildApiStory(peerId, story) as ApiStory | ApiStorySkipped,
+        story: apiStory,
       });
     }
   } else if (update instanceof GramJs.UpdateReadStories) {