feat(api-quota): enforce user storage quota checks

This commit is contained in:
2026-04-05 00:08:27 +08:00
parent bd3241504f
commit 32022c1437
+89 -14
View File
@@ -1,5 +1,5 @@
import { randomUUID } from "node:crypto"; import { randomUUID } from "node:crypto";
import { Injectable, NotFoundException } from "@nestjs/common"; import { Injectable, NotFoundException, PayloadTooLargeException } from "@nestjs/common";
import { ConfigService } from "@nestjs/config"; import { ConfigService } from "@nestjs/config";
import { PutObjectCommand, S3Client } from "@aws-sdk/client-s3"; import { PutObjectCommand, S3Client } from "@aws-sdk/client-s3";
import { getSignedUrl } from "@aws-sdk/s3-request-presigner"; import { getSignedUrl } from "@aws-sdk/s3-request-presigner";
@@ -8,6 +8,11 @@ import { PrismaService } from "../prisma/prisma.service";
import { CompleteAttachmentDto } from "./dto/complete-attachment.dto"; import { CompleteAttachmentDto } from "./dto/complete-attachment.dto";
import { PresignAttachmentDto } from "./dto/presign-attachment.dto"; import { PresignAttachmentDto } from "./dto/presign-attachment.dto";
type QuotaInfo = {
totalBytes: bigint;
usedBytes: bigint;
};
export type PresignAttachmentResponse = { export type PresignAttachmentResponse = {
method: "PUT"; method: "PUT";
uploadUrl: string; uploadUrl: string;
@@ -15,6 +20,11 @@ export type PresignAttachmentResponse = {
objectKey: string; objectKey: string;
objectUrl: string; objectUrl: string;
expiresInSeconds: number; expiresInSeconds: number;
quota: {
totalBytes: string;
usedBytes: string;
remainingBytes: string;
};
headers: { headers: {
"Content-Type": string; "Content-Type": string;
}; };
@@ -49,6 +59,9 @@ export class AttachmentService {
userId: string, userId: string,
body: PresignAttachmentDto body: PresignAttachmentDto
): Promise<PresignAttachmentResponse> { ): Promise<PresignAttachmentResponse> {
const quotaInfo = await this.getQuotaSnapshot(userId);
this.assertQuotaAvailable(quotaInfo.totalBytes, quotaInfo.usedBytes, body.fileSize);
if (body.taskId) { if (body.taskId) {
await this.ensureTaskOwnership(userId, body.taskId); await this.ensureTaskOwnership(userId, body.taskId);
} }
@@ -76,6 +89,11 @@ export class AttachmentService {
objectKey, objectKey,
objectUrl, objectUrl,
expiresInSeconds, expiresInSeconds,
quota: {
totalBytes: quotaInfo.totalBytes.toString(),
usedBytes: quotaInfo.usedBytes.toString(),
remainingBytes: (quotaInfo.totalBytes - quotaInfo.usedBytes).toString()
},
headers: { headers: {
"Content-Type": body.mimeType "Content-Type": body.mimeType
} }
@@ -92,20 +110,45 @@ export class AttachmentService {
const bucket = body.bucket ?? this.getDefaultBucket(); const bucket = body.bucket ?? this.getDefaultBucket();
const objectUrl = this.resolveObjectUrl(bucket, body.objectKey); const objectUrl = this.resolveObjectUrl(bucket, body.objectKey);
const attachment = await this.prismaService.attachment.create({
data: { const attachment = await this.prismaService.$transaction(async (tx) => {
userId, const quotaInfo = await this.getQuotaSnapshot(userId, tx);
taskId: body.taskId ?? null, this.assertQuotaAvailable(quotaInfo.totalBytes, quotaInfo.usedBytes, body.fileSize);
type: body.type ?? this.resolveAttachmentType(body.mimeType),
url: objectUrl, const uploadBytes = BigInt(body.fileSize);
mimeType: body.mimeType, const maxUsedBeforeUpload = quotaInfo.totalBytes - uploadBytes;
fileName: body.fileName, const updatedUser = await tx.user.updateMany({
fileSize: body.fileSize, where: {
width: body.width ?? null, id: userId,
height: body.height ?? null, usedStorageBytes: {
durationMs: body.durationMs ?? null, lte: maxUsedBeforeUpload
checksum: body.checksum ?? null }
},
data: {
usedStorageBytes: {
increment: uploadBytes
}
}
});
if (updatedUser.count === 0) {
throw new PayloadTooLargeException("存储配额不足");
} }
return tx.attachment.create({
data: {
userId,
taskId: body.taskId ?? null,
type: body.type ?? this.resolveAttachmentType(body.mimeType),
url: objectUrl,
mimeType: body.mimeType,
fileName: body.fileName,
fileSize: body.fileSize,
width: body.width ?? null,
height: body.height ?? null,
durationMs: body.durationMs ?? null,
checksum: body.checksum ?? null
}
});
}); });
return { return {
@@ -204,4 +247,36 @@ export class AttachmentService {
throw new NotFoundException("任务不存在"); throw new NotFoundException("任务不存在");
} }
} }
private async getQuotaSnapshot(
userId: string,
tx: Pick<PrismaService, "user"> = this.prismaService
): Promise<QuotaInfo> {
const user = await tx.user.findUnique({
where: {
id: userId
},
select: {
id: true,
defaultStorageQuotaMb: true,
usedStorageBytes: true
}
});
if (!user) {
throw new NotFoundException("用户不存在");
}
return {
totalBytes: BigInt(user.defaultStorageQuotaMb) * 1024n * 1024n,
usedBytes: user.usedStorageBytes
};
}
private assertQuotaAvailable(totalBytes: bigint, usedBytes: bigint, fileSize: number): void {
const uploadBytes = BigInt(fileSize);
if (uploadBytes > totalBytes || usedBytes + uploadBytes > totalBytes) {
throw new PayloadTooLargeException("存储配额不足");
}
}
} }