keanu-weblite/src/models/attachmentManager.ts

import { MatrixClient, MatrixEvent, Room, RoomEvent } from "matrix-js-sdk";
import {
  EventAttachment,
  EventAttachmentUrlData,
  EventAttachmentUrlType,
  KeanuEvent,
  KeanuEventExtension,
  KeanuRoom,
} from "./eventAttachment";
import axios, { AxiosRequestConfig, AxiosResponse } from "axios";
import { Counter, ModeOfOperation } from "aes-js";
import {
  Attachment,
  AttachmentBatch,
  AttachmentSendInfo,
  AttachmentSendStatus,
  AttachmentThumbnail,
} from "./attachment";
import proofmode from "../plugins/proofmode";
import imageResize from "image-resize";
import { computed, ref, Ref, shallowReactive, unref } from "vue";
import utils, { THUMBNAIL_MAX_WIDTH, THUMBNAIL_MAX_HEIGHT, CLIENT_EVENT_MEDIA_INTERVENTION_FLAGS } from "@/plugins/utils";
import { extractMediaMetadata } from "./proof";

export class AttachmentManager {
  matrixClient: MatrixClient;
  useAuthedMedia: boolean;
  maxSizeUploads: number;
  maxSizeAutoDownloads: number;

  cache: Map<string, EventAttachment>;
  cacheUploads: Map<string, AttachmentBatch>;

  constructor(matrixClient: MatrixClient, useAuthedMedia: boolean, maxSizeAutoDownloads: number) {
    this.matrixClient = matrixClient;
    this.useAuthedMedia = useAuthedMedia;
    this.maxSizeUploads = 0;
    this.maxSizeAutoDownloads = maxSizeAutoDownloads;

    this.cache = new Map();
    this.cacheUploads = new Map();

    // Get max upload size
    this.matrixClient
      .getMediaConfig(useAuthedMedia)
      .then((config) => {
        this.maxSizeUploads = config["m.upload.size"] ?? 0;
      })
      .catch(() => {});
  }

  public createUpload(room: KeanuRoom) {
    return createUploadBatch(this, room);
  }

  public createAttachment(file: File, room: KeanuRoom): Attachment {
    let a: Attachment = {
      status: "loading",
      file: file,
      useCompressed: false,
      detailsViewed: false,
      sendInfo: {
        status: "initial",
        statusDate: 0,
        mediaEventId: undefined,
        progress: 0,
        promise: undefined,
        randomRotation: 0,
        randomTranslationX: 0,
        randomTranslationY: 0,
      },
    };
    const ra = shallowReactive(a);
    this.prepareUpload(ra, room);
    return ra;
  }

  private async prepareUpload(attachment: Attachment, room: KeanuRoom): Promise<Attachment> {
    const file = attachment.file;
    if (utils.isSupportedImageType(file.type)) {
      let url = URL.createObjectURL(file);
      attachment.src = url;
      if (attachment.src) {
        try {
          let img = new Image();
          img.src = url;
          attachment.dimensions = await new Promise((response) => {
            img.onload = (event) => {
              response({ width: img.width, height: img.height });
            };
            img.onerror = (event) => {
              response(undefined);
            };
          });

          // Need to resize?
          const w = attachment.dimensions?.width ?? 0;
          const h = attachment.dimensions?.height ?? 0;
          const sizeDown = (w > 640 || h > 640);
          var aspect = w / h;
          var newWidth = sizeDown ? parseInt((w > h ? 640 : 640 * aspect).toFixed()) : w;
          var newHeight = sizeDown ? parseInt((w > h ? 640 / aspect : 640).toFixed()) : h;
          const compressedImg = await imageResize(file, {
            format: "webp",
            width: newWidth,
            height: newHeight,
            outputType: "blob",
          });
          attachment.compressedFile = new File([compressedImg as BlobPart], file.name, {
            type: "image/webp",
            lastModified: Date.now(),
          });
          attachment.compressedDimensions = {
            width: newWidth,
            height: newHeight,
          };
        } catch (error) {
          console.error("Failed to get image dimensions: " + error);
        }
      }
      try {
        attachment.proof = await proofmode.proofCheckFile(file);
        attachment.mediaMetadata = extractMediaMetadata(attachment.proof);

        // Default to scaled version if the image does not contain Content Credentials
        //
        const isDirectRoom = (room: Room) => {
            // TODO - Use the is_direct accountData flag (m.direct). WE (as the client)
            // apprently need to set this...
            if (room && room.getJoinRule() == "invite" && room.getInvitedAndJoinedMemberCount() == 2) {
                return true;
            }
            return false;
        };
        const isChannel = room.displayType == "im.keanu.room_type_channel";
        const isFileDrop = room.displayType == "im.keanu.room_type_file";

        let useOriginal = false;
        if (isChannel) {
          useOriginal = false;
        } else if (isFileDrop) {
          useOriginal = true;
        } else {
          if (isDirectRoom(room) && attachment.proof?.integrity?.c2pa !== undefined) {
            useOriginal = true;
          }
        }

        attachment.useCompressed = attachment.compressedFile !== undefined && !useOriginal;
      } catch (error) {
        console.error("Failed to get content credentials: " + error);
      }
    } else if (file.type.startsWith("video/")) {
      let url = URL.createObjectURL(file);
      const thumb: AttachmentThumbnail | undefined = await new Promise((resolve) => {
        if (url) {
          try {
            let canvas: HTMLCanvasElement = document.createElement("canvas");

            canvas.width = 640;
            canvas.height = 480;
            let video: HTMLVideoElement = document.createElement("video");
            video.preload = "auto";
            video.addEventListener("loadedmetadata", function () {
              if (video.videoWidth > THUMBNAIL_MAX_WIDTH || video.videoHeight > THUMBNAIL_MAX_HEIGHT) {
                var aspect = video.videoWidth / video.videoHeight;
                canvas.width = parseInt(
                  (video.videoWidth > video.videoHeight ? THUMBNAIL_MAX_WIDTH : THUMBNAIL_MAX_HEIGHT * aspect).toFixed()
                );
                canvas.height = parseInt(
                  (video.videoWidth > video.videoHeight ? THUMBNAIL_MAX_WIDTH / aspect : THUMBNAIL_MAX_HEIGHT).toFixed()
                );
              } else {
                canvas.width = video.videoWidth;
                canvas.height = video.videoHeight;
              }
            });
            video.addEventListener("loadeddata", (e) => {
              const ctx = canvas.getContext("2d");
              if (ctx) {
                ctx.drawImage(video, 0, 0, video.videoWidth, video.videoHeight, 0, 0, canvas.width, canvas.height);
              }
              canvas.toBlob((b) => {
                b?.arrayBuffer().then((data) => {
                  resolve({
                    data: new Uint8Array(data),
                    mimetype: b.type,
                    size: b.size,
                    w: canvas.width,
                    h: canvas.height,
                  });
                });
              }, "image/png");
            });
            video.src = url;
            video.currentTime = 0.1;
            video.load();
          } catch (error) {
            console.error("Failed to get video thumbnail: " + error);
            resolve(undefined);
          }
        } else {
          resolve(undefined);
        }
      });
      attachment.thumbnail = thumb;
    }

    attachment.status = "loaded";
    return attachment;
  }

  private getFileName(event: KeanuEvent) {
    const content = event.getContent();
    return (content.body || content.filename || "").toLowerCase();
  }

  private getSrcFileSize(event: KeanuEvent) {
    const content = event.getContent();
    if (content.info) {
      return content.info.size;
    }
    return 0;
  }

  public getEventAttachment(event: KeanuEvent): EventAttachment {
    let entry = this.cache.get(event.getId() ?? "invalid");
    if (entry !== undefined) {
      return entry;
    }

    const fileSize = this.getSrcFileSize(event);

    let mediaInterventionFlags = event.getContent()[CLIENT_EVENT_MEDIA_INTERVENTION_FLAGS];

    const attachment: EventAttachment = {
      event: event,
      name: this.getFileName(event),
      srcSize: fileSize,
      srcProgress: -1,
      thumbnailProgress: -1,
      autoDownloadable: fileSize <= this.maxSizeAutoDownloads,
      mediaInterventionFlags: mediaInterventionFlags ? JSON.parse(mediaInterventionFlags) : undefined,
      mediaMetadata: undefined,
      proof: undefined,
      loadSrc: () => Promise.reject("Not implemented"),
      loadThumbnail: () => Promise.reject("Not implemented"),
      loadBlob: () => Promise.reject("Not implemented"),
      release: () => Promise.reject("Not implemented"),
    };
    attachment.loadSrc = () => {
      if (attachment.src) {
        return Promise.resolve({ data: attachment.src, type: "src" });
      } else if (attachment.srcPromise) {
        return attachment.srcPromise;
      }
      attachment.srcPromise = this._loadEventAttachmentOrThumbnail(event, false, false, (percent) => {
        attachment.srcProgress = percent;
      }).then((res) => {
        attachment.src = (res as EventAttachmentUrlData).data;
        return res;
      }) as Promise<EventAttachmentUrlData>;
      return attachment.srcPromise as Promise<{ data: string; type: EventAttachmentUrlType }>;
    };
    attachment.loadThumbnail = () => {
      if (attachment.thumbnail) {
        return Promise.resolve({ data: attachment.thumbnail, type: "thumbnail" });
      } else if (attachment.thumbnailPromise) {
        return attachment.thumbnailPromise;
      }
      attachment.thumbnailPromise = this._loadEventAttachmentOrThumbnail(event, true, false, (percent) => {
        attachment.thumbnailProgress = percent;
      }).then((res) => {
        attachment.thumbnail = res.data as string;
        if (res.type == "src") {
          // Downloaded the src as thumb, so set "src" as well!
          attachment.src = res.data as string;
        }
        return res;
      }) as Promise<EventAttachmentUrlData>;
      return attachment.thumbnailPromise;
    };
    attachment.loadBlob = () => {
      const promise = this._loadEventAttachmentOrThumbnail(event, false, true, (percent) => {
        attachment.srcProgress = percent;
      }).then((res) => {
        return { data: res.data as Blob };
      });
      return promise;
    };
    attachment.release = (src: boolean, thumbnail: boolean) => {
      // TODO - figure out logic
      if (entry) {
        // TODO - abortable promises
        this.cache.delete(event.getId() ?? "invalid");
        if (attachment.src) {
          URL.revokeObjectURL(attachment.src);
        }
        if (attachment.thumbnail) {
          URL.revokeObjectURL(attachment.thumbnail);
        }
      }
    };
    if (event.getId()) {
      this.cache.set(event.getId()!, attachment!);
    }
    return attachment;
  }

  private async _loadEventAttachmentOrThumbnail(
    event: MatrixEvent & KeanuEventExtension,
    thumbnail: boolean,
    asBlob: boolean,
    progress?: (percent: number) => void
  ): Promise<EventAttachmentUrlData | { data: Blob; type: EventAttachmentUrlType }> {
    await this.matrixClient.decryptEventIfNeeded(event);

    let urltype: EventAttachmentUrlType = thumbnail ? "thumbnail" : "src";

    const content = event.getContent();
    var url = null;
    var mime = "image/png";
    var file = null;
    let decrypt = true;
    if (thumbnail && !!content.info && !!content.info.thumbnail_url) {
      url = this.matrixClient.mxcUrlToHttp(
        content.info.thumbnail_url,
        undefined,
        undefined,
        undefined,
        undefined,
        undefined,
        this.useAuthedMedia
      );
      decrypt = false;
      if (content.info.thumbnail_info) {
        mime = content.info.thumbnail_info.mimetype;
      }
    } else if (content.url != null) {
      url = this.matrixClient.mxcUrlToHttp(
        content.url,
        undefined,
        undefined,
        undefined,
        undefined,
        undefined,
        this.useAuthedMedia
      );
      decrypt = false;
      if (content.info) {
        mime = content.info.mimetype;
      }
    } else if (thumbnail && content && content.info && content.info.thumbnail_file && content.info.thumbnail_file.url) {
      file = content.info.thumbnail_file;
      url = this.matrixClient.mxcUrlToHttp(
        file.url,
        undefined,
        undefined,
        undefined,
        undefined,
        undefined,
        this.useAuthedMedia
      );
      mime = file.mimetype;
    } else if (
      content.file &&
      content.file.url &&
      event.getContent()?.info?.size > 0 &&
      (!thumbnail || event.getContent()?.info?.size < this.maxSizeAutoDownloads)
    ) {
      // No thumb, use real url
      file = content.file;
      url = this.matrixClient.mxcUrlToHttp(
        file.url,
        undefined,
        undefined,
        undefined,
        undefined,
        undefined,
        this.useAuthedMedia
      );
      mime = file.mimetype;
      urltype = "src";
    }

    if (url == null) {
      throw new Error("No url found!");
    }

    let options: AxiosRequestConfig = {
      responseType: "arraybuffer",
      onDownloadProgress: (progressEvent) => {
        if (progress) {
          progress(Math.floor((progressEvent.progress ?? 0) * 100));
        }
      },
    };
    if (this.useAuthedMedia) {
      options.headers = {
        Authorization: `Bearer ${this.matrixClient.getAccessToken()}`,
      };
    }

    const response = await axios.get(url, options);
    const bytes = decrypt ? await this.decryptData(file, response) : { buffer: response.data };
    const blob = new Blob([bytes.buffer], { type: mime });
    return asBlob ? { data: blob, type: urltype } : { data: URL.createObjectURL(blob), type: urltype };
  }

  private b64toBuffer(val: any) {
    const baseValue = val.replaceAll("-", "+").replaceAll("_", "/");
    return Buffer.from(baseValue, "base64");
  }

  private decryptData(file: any, response: AxiosResponse): Promise<Uint8Array> {
    return new Promise((resolve, reject) => {
      const key = this.b64toBuffer(file.key.k);
      const iv = this.b64toBuffer(file.iv);
      const originalHash = this.b64toBuffer(file.hashes.sha256);

      var aesCtr = new ModeOfOperation.ctr(key, new Counter(iv));
      const data = new Uint8Array(response.data);

      crypto.subtle
        .digest("SHA-256", data)
        .then((hash) => {
          // Calculate sha256 and compare hashes
          if (Buffer.compare(Buffer.from(hash), originalHash) != 0) {
            reject("Hashes don't match!");
            return;
          }
          var decryptedBytes = aesCtr.decrypt(data);
          resolve(decryptedBytes);
        })
        .catch((err) => {
          reject("Failed to calculate hash value");
        });
    });
  }
}

export const createUploadBatch = (manager: AttachmentManager | null, room: KeanuRoom | null): AttachmentBatch => {
  const matrixClient = manager?.matrixClient;
  const maxSizeUploads = manager?.maxSizeUploads ?? 0;

  const txnId = utils.randomPass();

  const sendingStatus: Ref<AttachmentSendStatus> = ref("initial");
  const sendingRootEventId: Ref<string | undefined> = ref(undefined);
  const sendingRootMessage: Ref<string | undefined> = ref(undefined);
  const sendingPromise: Ref<Promise<any> | undefined> = ref(undefined);
  const attachments: Ref<Attachment[]> = ref([]);

  const totalOriginalFileSize: Ref<number> = ref(0);
  const progressPercent: Ref<number> = ref(0);

  const updateProgress = () => {
    // Use relative sizes of the original files to determine how many percent
    // the individual files contribute to the total progress.
    const progress = attachments.value.reduce((cb, item) => {
      const info = item.sendInfo;
      const thisFileCurrent = item.file.size;
      const max = totalOriginalFileSize.value > 0 ? thisFileCurrent / totalOriginalFileSize.value : 0;
      const q = (info.progress * max) / 100;
      return cb + q;
    }, 0);
    const percent = parseInt(Math.floor(progress * 100).toFixed());
    progressPercent.value = percent;
  };

  const attachmentsSentCount = computed(() => {
    return attachments.value.reduce((a, elem) => (elem.sendInfo?.status == "sent" ? a + 1 : a), 0);
  });

  const attachmentsSending = computed(() => {
    return attachments.value.filter((elem) => elem.sendInfo?.status == "initial" || elem.sendInfo?.status == "sending");
  });

  const attachmentsSent = computed(() => {
    sortSendingAttachments();
    return attachments.value.filter((elem) => elem.sendInfo?.status == "sent");
  });

  const sortSendingAttachments = () => {
    attachments.value.sort((a, b) => (b.sendInfo?.statusDate ?? 0) - (a.sendInfo?.statusDate ?? 0));
  };

  const addFiles = (filesToAdd: File[]) => {
    if (sendingStatus.value == "initial" && manager !== null && room !== null) {
      attachments.value.push(...filesToAdd.map((f) => manager!.createAttachment(f, room)));
    }
  };

  const removeAttachment = (attachment: Attachment) => {
    if (sendingStatus.value == "initial") {
      const index = attachments.value.indexOf(attachment);
      if (index > -1) {
        attachments.value.splice(index, 1);
      }
    }
  };

  const isTooLarge = (attachment: Attachment) => {
    const file = attachment.compressedFile && attachment.useCompressed ? attachment.compressedFile : attachment.file;
    return file.size > maxSizeUploads;
  };

  const canSend = computed(() => {
    return (
      attachments.value.length > 0 && !attachments.value.some((a: Attachment) => isTooLarge(a) || a.status !== "loaded")
    );
  });

  const cancel = () => {
    if (sendingStatus.value !== "initial" && matrixClient && room) {
      attachments.value.toReversed().forEach((attachment) => {
        cancelSendAttachment(attachment);
      });
      sendingStatus.value = "canceled";
      if (sendingRootEventId.value) {
        // Redact all media we already sent, plus the root event
        let promises = attachments.value.reduce((val: Promise<any>[], attachment: Attachment) => {
          if (attachment.sendInfo?.mediaEventId) {
            val.push(
              matrixClient.redactEvent(room.roomId, attachment.sendInfo!.mediaEventId, undefined, {
                reason: "cancel",
              })
            );
          }
          return val;
        }, [] as Promise<any>[]);
        if (sendingRootEventId.value) {
          promises.push(
            matrixClient.redactEvent(room.roomId, sendingRootEventId.value, undefined, {
              reason: "cancel",
            })
          );
        }
        Promise.allSettled(promises)
          .then(() => {
            console.log("Message redacted");
          })
          .catch((err) => {
            console.log("Redaction failed: ", err);
          });
      }
    }
  };

  const cancelSendAttachment = (attachment: Attachment) => {
    if (attachment.sendInfo) {
      if (attachment.sendInfo.promise && attachment.sendInfo.status != "initial") {
        attachment.sendInfo.promise.abort();
      }
      attachment.sendInfo.status = "canceled";
    }
  };

  const send = async (message: string): Promise<any> => {
    if (!matrixClient || !room) return Promise.reject("Not configured");
    sendingStatus.value = "sending";
    progressPercent.value = 0;
    sendingRootMessage.value = message;

    attachments.value.forEach((attachment) => {
      let sendInfo: AttachmentSendInfo = {
        status: "initial",
        statusDate: Date.now(),
        mediaEventId: undefined,
        progress: 0,
        randomRotation: 0,
        randomTranslationX: 0,
        randomTranslationY: 0,
        promise: undefined,
      };
      attachment.sendInfo = shallowReactive(sendInfo);
      attachment.proof = undefined;
    });

    totalOriginalFileSize.value = attachments.value.reduce((cb, item) => {
      return cb + item.file.size;
    }, 0);

    const onLocalEchoUpdated = (event: KeanuEvent) => {
      if (!event.uploadBatch && event.getTxnId() === txnId) {
        event.uploadBatch = {
          sendingStatus,
          sendingRootEventId,
          sendingRootMessage,
          progressPercent,
          attachments: unref(attachments),
          attachmentsSentCount,
          attachmentsSending,
          attachmentsSent,
          addFiles,
          removeAttachment,
          isTooLarge,
          canSend,
          send,
          cancel,
          cancelSendAttachment,
        };
      }
    };
    room.on(RoomEvent.LocalEchoUpdated, onLocalEchoUpdated);

    const promise = (sendingRootEventId.value ? Promise.resolve(sendingRootEventId.value) : utils.sendTextMessage(matrixClient, room.roomId, message, undefined, undefined, txnId))
      .then((eventId: string) => {
        sendingRootEventId.value = eventId;

        let promiseChain = Promise.resolve();
        const getItemPromise = (index: number) => {
          if (index < attachments.value.length) {
            const attachment = attachments.value[index];
            const item = attachment;
            if (item.sendInfo.status !== "initial" && item.sendInfo.status !== "failed") {
              return getItemPromise(++index);
            }
            item.sendInfo.status = "sending";

            let file = (() => {
              if (attachment.compressedFile && attachment.useCompressed) {
                // Send compressed!
                return attachment.compressedFile;
              } else {
                // Send original
                return attachment.file;
              }
            })();

            const itemPromise = utils
              .sendFile(
                matrixClient,
                room.roomId,
                file,
                ({ loaded, total }: { loaded: number; total: number }) => {
                  if (loaded == total) {
                    item.sendInfo.progress = 100;
                  } else if (total > 0) {
                    item.sendInfo.progress = (100 * loaded) / total;
                  }
                  updateProgress();
                },
                eventId,
                attachment.dimensions,
                attachment.thumbnail,
                attachment.mediaMetadata
              )
              .then((mediaEventId: string) => {
                // Look at last item rotation, flipping the sign on this, so looks more like a true stack
                let signR = 1;
                let signX = 1;
                let signY = 1;
                if (attachmentsSent.value.length > 0) {
                  if (attachmentsSent.value[0].sendInfo!.randomRotation >= 0) {
                    signR = -1;
                  }
                  if (attachmentsSent.value[0].sendInfo!.randomTranslationX >= 0) {
                    signX = -1;
                  }
                  if (attachmentsSent.value[0].sendInfo!.randomTranslationY >= 0) {
                    signY = -1;
                  }
                }
                item.sendInfo.randomRotation = signR * (2 + Math.random() * 10);
                item.sendInfo.randomTranslationX = signX * Math.random() * 20;
                item.sendInfo.randomTranslationY = signY * Math.random() * 20;
                item.sendInfo.mediaEventId = mediaEventId;
                item.sendInfo.status = "sent";
                item.sendInfo.statusDate = Date.now();
              })
              .catch((error: any) => {
                if (item.sendInfo.promise?.aborted) {
                  item.sendInfo.status = "canceled";
                } else {
                  console.error("ERROR", error);
                  item.sendInfo.status = "failed";
                  return Promise.reject(error)
                }
                return Promise.resolve();
              });
            item.sendInfo.promise = itemPromise;
            return itemPromise.then(() => getItemPromise(++index));
          } else return Promise.resolve();
        };

        return promiseChain.then(() => getItemPromise(0));
      })
      .then(() => {
        sendingStatus.value = "sent";
        sendingRootEventId.value = undefined;
      })
      .catch((err: any) => {
        console.error("Upload error", err);
        sendingStatus.value = "failed";
      })
      .finally(() => {
        room.off(RoomEvent.LocalEchoUpdated, onLocalEchoUpdated);
      });
    sendingPromise.value = promise;
    return promise;
  };

  return {
    sendingStatus,
    sendingRootEventId,
    sendingRootMessage,
    progressPercent,
    attachments: unref(attachments),
    attachmentsSentCount,
    attachmentsSending,
    attachmentsSent,
    addFiles,
    removeAttachment,
    isTooLarge,
    canSend,
    send,
    cancel,
    cancelSendAttachment,
  };
};