keanu-weblite/src/models/attachmentManager.ts

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

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