private_image/backend/src/modules/upload/services/UploadService.ts

import { BadRequestException, Injectable } from '@nestjs/common';
import { InjectDataSource } from '@nestjs/typeorm';
import { fileTypeFromBuffer } from 'file-type';
import { createHash, randomBytes, randomUUID } from 'node:crypto';
import { DataSource } from 'typeorm';
import { fromBase64 } from '../../../utils/fromBase64';
import { StorageService } from '../../storage/services/StorageService';
import { CreateHotlinkUploadDto } from '../dto/CreateHotlinkUploadDto';
import { CreateZkUploadDto } from '../dto/CreateZkUploadDto';
import { Upload } from '../entities/Upload';
import { CreatedUploadResult } from '../types/CreatedUploadResult';
import { UploadMode } from '../types/UploadMode';

@Injectable()
export class UploadService {
    constructor(
        @InjectDataSource()
        private readonly dataSource: DataSource,
        private readonly storageService: StorageService
    ) {}

    async createZkUpload({
        iv,
        ciphertext,
        expiresAt,
        maxViews,
        passwordProtected,
        wrappedDek,
        wrapIv,
        salt
    }: CreateZkUploadDto): Promise<CreatedUploadResult> {
        const ciphertextBuffer = fromBase64(ciphertext);

        await this.throwIfUnencryptedImage(ciphertextBuffer);

        const deleteToken = randomBytes(32).toString('base64url');
        const id = randomUUID();
        const storagePath = id;

        await this.dataSource.transaction(async manager => {
            const uploadRepository = manager.getRepository(Upload);

            const upload = uploadRepository.create({
                id,
                mode: UploadMode.Zk,
                storagePath,
                iv: fromBase64(iv),
                maxViews,
                deleteTokenHash: this.hashDeleteToken(deleteToken),
                expiresAt,
                zk: {
                    passwordProtected,
                    wrappedDek: passwordProtected && wrappedDek ? fromBase64(wrappedDek) : null,
                    wrapIv: passwordProtected && wrapIv ? fromBase64(wrapIv) : null,
                    salt: passwordProtected && salt ? fromBase64(salt) : null
                }
            });

            await uploadRepository.save(upload);

            await this.storageService.write(storagePath, ciphertextBuffer);
        });

        return {
            id,
            mode: UploadMode.Zk,
            deleteToken
        };
    }

    async createHotlinkUpload({
        iv,
        ciphertext,
        expiresAt,
        maxViews,
        mimeType
    }: CreateHotlinkUploadDto): Promise<CreatedUploadResult> {
        const ciphertextBuffer = fromBase64(ciphertext);

        await this.throwIfUnencryptedImage(ciphertextBuffer);

        const deleteToken = randomBytes(32).toString('base64url');
        const id = randomUUID();
        const storagePath = id;

        await this.dataSource.transaction(async manager => {
            const uploadRepository = manager.getRepository(Upload);

            const upload = uploadRepository.create({
                id,
                mode: UploadMode.Hotlink,
                storagePath,
                iv: fromBase64(iv),
                maxViews,
                deleteTokenHash: this.hashDeleteToken(deleteToken),
                expiresAt,
                hotlink: { mimeType }
            });

            await uploadRepository.save(upload);

            await this.storageService.write(storagePath, ciphertextBuffer);
        });

        return {
            id,
            mode: UploadMode.Hotlink,
            deleteToken
        };
    }

    private hashDeleteToken(token: string): string {
        return createHash('sha256').update(token, 'utf8').digest('hex');
    }

    private async throwIfUnencryptedImage(ciphertext: Buffer): Promise<void> {
        const detected = await fileTypeFromBuffer(ciphertext);

        if (detected?.mime.startsWith('image/')) {
            throw new BadRequestException('Ciphertext must not be a plaintext image');
        }
    }
}