keanu-weblite/src/models/attachmentManager.ts

import { MatrixClient, MatrixEvent, Room } from "matrix-js-sdk";
import { EventAttachment, KeanuEventExtension } from "./eventAttachment";
import axios, { AxiosRequestConfig, AxiosResponse } from "axios";
import { Counter, ModeOfOperation } from "aes-js";
import { Attachment, AttachmentBatch, AttachmentSendInfo } from "./attachment";
import proofmode from "../plugins/proofmode";
import imageSize from "image-size";
import imageResize from "image-resize";
import { computed, isRef, Reactive, reactive, ref, Ref } from "vue";
import utils from "@/plugins/utils";

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

  cache: Map<string | undefined, Reactive<EventAttachment>>;

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

    this.cache = new Map();

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

  public createUpload(room: Room) {
    return createUploadBatch(this.matrixClient, room);
  }

  public createAttachment(file: File): Attachment {
    let a: Attachment = {
      status: "loading",
      file: file,
      useScaled: false,
    };
    const ra = reactive(a);
    this.prepareUpload(ra);
    return ra;
  }

  private async prepareUpload(attachment: Attachment): Promise<Attachment> {
    const file = attachment.file;
    if (file.type.startsWith("image/")) {
      attachment.proof = await proofmode.proofCheckFile(file);

      var reader = new FileReader();
      await new Promise((resolve) => {
        reader.onload = (evt) => {
          attachment.src = (evt.target?.result as string) ?? undefined;
          if (attachment.src) {
            try {
              const buffer = Uint8Array.from(window.atob(attachment.src.replace(/^data[^,]+,/, "")), (v) =>
                v.charCodeAt(0)
              );
              attachment.dimensions = imageSize(buffer);

              // Need to resize?
              const w = attachment.dimensions.width;
              const h = attachment.dimensions.height;
              if (w > 640 || h > 640) {
                var aspect = w / h;
                var newWidth = parseInt((w > h ? 640 : 640 * aspect).toFixed());
                var newHeight = parseInt((w > h ? 640 / aspect : 640).toFixed());
                imageResize(attachment.src, {
                  format: "webp",
                  width: newWidth,
                  height: newHeight,
                  outputType: "blob",
                })
                  .then((img) => {
                    attachment.scaledFile = new File([img as BlobPart], file.name, {
                      type: "image/webp",
                      lastModified: Date.now(),
                    });
                    attachment.scaledDimensions = {
                      width: newWidth,
                      height: newHeight,
                    };

                    // Use scaled version if the image does not contain C2PA
                    attachment.useScaled =
                      attachment.scaledFile !== undefined &&
                      (attachment.proof === undefined ||
                        attachment.proof.integrity === undefined ||
                        attachment.proof.integrity.c2pa === undefined);
                  })
                  .catch((err) => {
                    console.error("Resize failed:", err);
                  });
              }
            } catch (error) {
              console.error("Failed to get image dimensions: " + error);
            }
          }
          resolve(true);
        };
        reader.readAsDataURL(file);
      });
    }

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

  public getEventAttachment(event: MatrixEvent & KeanuEventExtension): Reactive<EventAttachment> {
    let entry = this.cache.get(event.getId());
    if (entry !== undefined) {
      return entry;
    }
    const attachment: Reactive<EventAttachment> = reactive({
      event: event,
      srcProgress: -1,
      thumbnailProgress: -1,
      loadSrc: () => Promise.reject("Not implemented"),
      loadThumbnail: () => Promise.reject("Not implemented"),
      release: () => Promise.reject("Not implemented"),
    });
    attachment.loadSrc = () => {
      if (attachment.src) {
        return Promise.resolve(attachment.src);
      } else if (attachment.srcPromise) {
        return attachment.srcPromise;
      }
      attachment.srcPromise = this._loadEventAttachmentOrThumbnail(event, false, (percent) => {
        attachment.srcProgress = percent;
      }).then((src) => {
        attachment.src = src;
        return src;
      });
      return attachment.srcPromise;
    };
    attachment.loadThumbnail = () => {
      if (attachment.thumbnail) {
        return Promise.resolve(attachment.thumbnail);
      } else if (attachment.thumbnailPromise) {
        return attachment.thumbnailPromise;
      }
      attachment.thumbnailPromise = this._loadEventAttachmentOrThumbnail(event, true, (percent) => {
        attachment.thumbnailProgress = percent;
      }).then((thummbnail) => {
        attachment.thumbnail = thummbnail;
        return thummbnail;
      });
      return attachment.thumbnailPromise;
    };
    attachment.release = (src: boolean, thumbnail: boolean) => {
      // TODO - figure out logic
      if (entry) {
        // TODO - abortable promises
        this.cache.delete(event.getId());
        if (attachment.src) {
          URL.revokeObjectURL(attachment.src);
        }
        if (attachment.thumbnail) {
          URL.revokeObjectURL(attachment.thumbnail);
        }
      }
    };
    this.cache.set(event.getId(), attachment!);
    return attachment;
  }

  private async _loadEventAttachmentOrThumbnail(
    event: MatrixEvent & KeanuEventExtension,
    thumbnail: boolean,
    progress?: (percent: number) => void
  ): Promise<string> {
    await this.matrixClient.decryptEventIfNeeded(event);

    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 &&
      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;
    }

    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 };
    return URL.createObjectURL(new Blob([bytes.buffer], { type: mime }));
  }

  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");
        });
    });
  }
}

const createUploadBatch = (matrixClient: MatrixClient, room: Room): AttachmentBatch => {
  const sendingStatus: Ref<"initial" | "sending" | "sent" | "canceled" | "failed"> = ref("initial");
  const sendingRootEventId: Ref<string | undefined> = ref(undefined);
  const sendingPromise: Ref<Promise<any> | undefined> = ref(undefined);
  const attachments: Ref<Attachment[]> = ref([]);

  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 addAttachment = (attachment: Attachment) => {
    if (sendingStatus.value == "initial") {
      attachments.value.push(attachment);
    }
  };

  const removeAttachment = (attachment: Attachment) => {
    if (sendingStatus.value == "initial") {
      attachments.value = attachments.value.filter((a) => a !== attachment);
    }
  };

  const cancel = () => {
    if (sendingStatus.value !== "initial") {
      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 = (message: string): Promise<any> => {
    sendingStatus.value = "sending";
    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 = reactive(sendInfo);
    });

    const sendingPromise = utils
      .sendTextMessage(matrixClient, room.roomId, message)
      .then((eventId: string) => {
        sendingRootEventId.value = eventId;

        // Use the eventId as a thread root for all the media
        let promiseChain = Promise.resolve();
        const getItemPromise = (index: number) => {
          if (index < attachments.value.length) {
            const attachment = attachments.value[index];
            const item = attachment.sendInfo!;
            if (item.status !== "initial") {
              return getItemPromise(++index);
            }
            item.status = "sending";

            let file = (() => {
              if (attachment.scaledFile && attachment.useScaled) {
                // Send scaled version of image instead!
                return attachment.scaledFile;
              } else {
                // Send actual file image when not scaled!
                return attachment.file;
              }
            })();

            const itemPromise = utils
              .sendFile(
                matrixClient,
                room.roomId,
                file,
                ({ loaded, total }: { loaded: number; total: number }) => {
                  if (loaded == total) {
                    item.progress = 100;
                  } else if (total > 0) {
                    item.progress = (100 * loaded) / total;
                  }
                },
                eventId,
                attachment.dimensions
              )
              .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.randomRotation = signR * (2 + Math.random() * 10);
                item.randomTranslationX = signX * Math.random() * 20;
                item.randomTranslationY = signY * Math.random() * 20;
                item.mediaEventId = mediaEventId;
                item.status = "sent";
                item.statusDate = Date.now();
              })
              .catch((ignorederr: any) => {
                if (item.promise?.aborted) {
                  item.status = "canceled";
                } else {
                  console.error("ERROR", ignorederr);
                  item.status = "failed";
                }
                return Promise.resolve();
              });
            item.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("ERROR", err);
      });
    return sendingPromise;
  };

  return {
    sendingStatus,
    sendingRootEventId,
    sendingPromise,
    attachments,
    attachmentsSentCount,
    attachmentsSending,
    attachmentsSent,
    addAttachment,
    removeAttachment,
    send,
    cancel,
    cancelSendAttachment,
  };
};