diff --git a/migrations/20260623120000-create-mail-bucket-entries.js b/migrations/20260623120000-create-mail-bucket-entries.js new file mode 100644 index 0000000..f848c68 --- /dev/null +++ b/migrations/20260623120000-create-mail-bucket-entries.js @@ -0,0 +1,53 @@ +'use strict'; + +const TABLE_NAME = 'mail_bucket_entries'; + +/** @type {import('sequelize-cli').Migration} */ +module.exports = { + async up(queryInterface, Sequelize) { + await queryInterface.createTable(TABLE_NAME, { + id: { + type: Sequelize.UUID, + defaultValue: Sequelize.UUIDV4, + primaryKey: true, + allowNull: false, + }, + mail_account_id: { + type: Sequelize.UUID, + allowNull: false, + references: { model: 'mail_accounts', key: 'id' }, + onUpdate: 'CASCADE', + onDelete: 'CASCADE', + }, + entry_key: { + type: Sequelize.STRING(255), + allowNull: false, + unique: true, + }, + bridge_entry_id: { + type: Sequelize.STRING(24), + allowNull: false, + }, + size: { + type: Sequelize.BIGINT, + allowNull: false, + }, + created_at: { + type: Sequelize.DATE, + allowNull: false, + defaultValue: Sequelize.fn('now'), + }, + updated_at: { + type: Sequelize.DATE, + allowNull: false, + defaultValue: Sequelize.fn('now'), + }, + }); + + await queryInterface.addIndex(TABLE_NAME, ['mail_account_id']); + }, + + async down(queryInterface) { + await queryInterface.dropTable(TABLE_NAME); + }, +}; diff --git a/src/modules/account/repositories/address.repository.ts b/src/modules/account/repositories/address.repository.ts index ec11742..e650063 100644 --- a/src/modules/account/repositories/address.repository.ts +++ b/src/modules/account/repositories/address.repository.ts @@ -13,6 +13,7 @@ import { MailProviderAccountModel } from '../models/mail-provider-account.model. const MAX_BATCH_LOOKUP = 50; export interface ProviderAccountBucketContext { + mailAccountId: string; userUuid: string; networkBucketId: string | null; } @@ -110,6 +111,7 @@ export class AddressRepository { if (!model?.account) return null; return { + mailAccountId: model.account.id, userUuid: model.account.userId, networkBucketId: model.networkBucketId, }; @@ -132,6 +134,7 @@ export class AddressRepository { if (!link?.address?.account) return null; return { + mailAccountId: link.address.account.id, userUuid: link.address.account.userId, networkBucketId: link.address.networkBucketId, }; diff --git a/src/modules/email/email.module.ts b/src/modules/email/email.module.ts index 876af55..b4c93b0 100644 --- a/src/modules/email/email.module.ts +++ b/src/modules/email/email.module.ts @@ -1,5 +1,5 @@ import { Module } from '@nestjs/common'; -import { BridgeModule } from '../infrastructure/bridge/bridge.module.js'; +import { MailUsageModule } from '../usage/mail-usage.module.js'; import { JmapModule } from '../infrastructure/jmap/jmap.module.js'; import { SmtpModule } from '../infrastructure/smtp/smtp.module.js'; import { ProvisioningModule } from '../provisioning/provisioning.module.js'; @@ -8,7 +8,7 @@ import { EmailService } from './email.service.js'; import { Reflector } from '@nestjs/core'; @Module({ - imports: [JmapModule, SmtpModule, ProvisioningModule, BridgeModule], + imports: [JmapModule, SmtpModule, ProvisioningModule, MailUsageModule], controllers: [EmailController], providers: [EmailService, Reflector], }) diff --git a/src/modules/email/email.service.ts b/src/modules/email/email.service.ts index a9879d3..aeb7939 100644 --- a/src/modules/email/email.service.ts +++ b/src/modules/email/email.service.ts @@ -6,7 +6,7 @@ import { } from '@nestjs/common'; import { ConfigService } from '@nestjs/config'; import { AccountService } from '../account/account.service.js'; -import { BridgeClient } from '../infrastructure/bridge/bridge.service.js'; +import { MailUsageService } from '../usage/mail-usage.service.js'; import { MailProvider } from './mail-provider.port.js'; import type { DraftEmailDto, @@ -60,7 +60,7 @@ export class EmailService { private readonly accountService: AccountService, private readonly smtp: StalwartSmtpService, private readonly configService: ConfigService, - private readonly bridge: BridgeClient, + private readonly usage: MailUsageService, ) {} getMailboxes(userEmail: string): Promise { @@ -276,11 +276,11 @@ export class EmailService { } try { - await this.bridge.deleteBucketEntry( - context.userUuid, - context.networkBucketId, + await this.usage.releaseStoredMessage({ + userUuid: context.userUuid, + bucketId: context.networkBucketId, entryKey, - ); + }); } catch (error) { this.logger.warn( `Failed to release quota entry '${entryKey}' for '${userEmail}': ${(error as Error).message}`, diff --git a/src/modules/infrastructure/bridge/bridge.service.ts b/src/modules/infrastructure/bridge/bridge.service.ts index 1068606..03335d4 100644 --- a/src/modules/infrastructure/bridge/bridge.service.ts +++ b/src/modules/infrastructure/bridge/bridge.service.ts @@ -7,7 +7,11 @@ import { import { ConfigService } from '@nestjs/config'; import { JwtService } from '@nestjs/jwt'; import { Client } from 'undici'; -import type { BucketEntry, MailBucket } from './bridge.types.js'; +import type { + BucketEntry, + MailBucket, + UserSpaceSnapshot, +} from './bridge.types.js'; @Injectable() export class BridgeClient implements OnModuleInit, OnModuleDestroy { @@ -101,7 +105,6 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy { async createBucketEntry( userUuid: string, bucketId: string, - key: string, size: number, ): Promise { const token = this.signGatewayToken(userUuid); @@ -114,14 +117,14 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy { accept: 'application/json', authorization: `Bearer ${token}`, }, - body: JSON.stringify({ key, size }), + body: JSON.stringify({ size }), }); const text = await body.text(); if (statusCode !== 200) { throw new BridgeApiError( - `Failed to create bucket entry '${key}' on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`, + `Failed to create bucket entry on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`, statusCode, text, ); @@ -133,13 +136,13 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy { async deleteBucketEntry( userUuid: string, bucketId: string, - key: string, - ): Promise { + entryId: string, + ): Promise { const token = this.signGatewayToken(userUuid); const { statusCode, body } = await this.httpClient.request({ method: 'DELETE', - path: `${this.basePath}/v2/gateway/users/${encodeURIComponent(userUuid)}/buckets/${encodeURIComponent(bucketId)}/entries/${encodeURIComponent(key)}`, + path: `${this.basePath}/v2/gateway/users/${encodeURIComponent(userUuid)}/buckets/${encodeURIComponent(bucketId)}/entries/${encodeURIComponent(entryId)}`, headers: { accept: 'application/json', authorization: `Bearer ${token}`, @@ -150,11 +153,13 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy { if (statusCode !== 200) { throw new BridgeApiError( - `Failed to delete bucket entry '${key}' on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`, + `Failed to delete bucket entry '${entryId}' on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`, statusCode, text, ); } + + return JSON.parse(text) as UserSpaceSnapshot; } private signGatewayToken(userUuid: string): string { diff --git a/src/modules/stalwart-events/stalwart-events.module.ts b/src/modules/stalwart-events/stalwart-events.module.ts index 34d3b1d..e8a70a1 100644 --- a/src/modules/stalwart-events/stalwart-events.module.ts +++ b/src/modules/stalwart-events/stalwart-events.module.ts @@ -1,12 +1,12 @@ import { Module } from '@nestjs/common'; import { AccountModule } from '../account/account.module.js'; -import { BridgeModule } from '../infrastructure/bridge/bridge.module.js'; +import { MailUsageModule } from '../usage/mail-usage.module.js'; import { StalwartEventsController } from './stalwart-events.controller.js'; import { StalwartEventsAuthGuard } from './stalwart-events-auth.guard.js'; import { StalwartEventsService } from './stalwart-events.service.js'; @Module({ - imports: [AccountModule, BridgeModule], + imports: [AccountModule, MailUsageModule], controllers: [StalwartEventsController], providers: [StalwartEventsAuthGuard, StalwartEventsService], }) diff --git a/src/modules/stalwart-events/stalwart-events.service.ts b/src/modules/stalwart-events/stalwart-events.service.ts index 615ffa4..352de42 100644 --- a/src/modules/stalwart-events/stalwart-events.service.ts +++ b/src/modules/stalwart-events/stalwart-events.service.ts @@ -1,6 +1,6 @@ import { Injectable, Logger } from '@nestjs/common'; import { AccountService } from '../account/account.service.js'; -import { BridgeClient } from '../infrastructure/bridge/bridge.service.js'; +import { MailUsageService } from '../usage/mail-usage.service.js'; import type { StalwartEvent, StalwartWebhookPayload, @@ -12,7 +12,7 @@ export class StalwartEventsService { constructor( private readonly accounts: AccountService, - private readonly bridge: BridgeClient, + private readonly usage: MailUsageService, ) {} async handleBatch(payload: StalwartWebhookPayload): Promise { @@ -57,16 +57,12 @@ export class StalwartEventsService { return; } - const { totalUsedSpaceBytes } = await this.bridge.createBucketEntry( - context.userUuid, - context.networkBucketId, + await this.usage.trackStoredMessage({ + mailAccountId: context.mailAccountId, + userUuid: context.userUuid, + bucketId: context.networkBucketId, entryKey, size, - ); - - this.logger.log( - { entryKey, size, totalUsedSpaceBytes, type: event.type }, - 'Created bucket entry for ingested message', - ); + }); } } diff --git a/src/modules/usage/domain/mail-bucket-entry.domain.ts b/src/modules/usage/domain/mail-bucket-entry.domain.ts new file mode 100644 index 0000000..346841e --- /dev/null +++ b/src/modules/usage/domain/mail-bucket-entry.domain.ts @@ -0,0 +1,27 @@ +export interface MailBucketEntryAttributes { + id: string; + mailAccountId: string; + entryKey: string; + bridgeEntryId: string; + size: number; + createdAt: Date; + updatedAt: Date; +} + +export class MailBucketEntry { + readonly id!: string; + readonly mailAccountId!: string; + readonly entryKey!: string; + readonly bridgeEntryId!: string; + readonly size!: number; + readonly createdAt!: Date; + readonly updatedAt!: Date; + + private constructor(attributes: MailBucketEntryAttributes) { + Object.assign(this, attributes); + } + + static build(attributes: MailBucketEntryAttributes): MailBucketEntry { + return new MailBucketEntry(attributes); + } +} diff --git a/src/modules/usage/mail-usage.module.ts b/src/modules/usage/mail-usage.module.ts new file mode 100644 index 0000000..79abfc8 --- /dev/null +++ b/src/modules/usage/mail-usage.module.ts @@ -0,0 +1,13 @@ +import { Module } from '@nestjs/common'; +import { SequelizeModule } from '@nestjs/sequelize'; +import { BridgeModule } from '../infrastructure/bridge/bridge.module.js'; +import { MailBucketEntryModel } from './models/mail-bucket-entry.model.js'; +import { MailBucketEntryRepository } from './repositories/mail-bucket-entry.repository.js'; +import { MailUsageService } from './mail-usage.service.js'; + +@Module({ + imports: [SequelizeModule.forFeature([MailBucketEntryModel]), BridgeModule], + providers: [MailBucketEntryRepository, MailUsageService], + exports: [MailUsageService], +}) +export class MailUsageModule {} diff --git a/src/modules/usage/mail-usage.service.ts b/src/modules/usage/mail-usage.service.ts new file mode 100644 index 0000000..06573d2 --- /dev/null +++ b/src/modules/usage/mail-usage.service.ts @@ -0,0 +1,94 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { BridgeClient } from '../infrastructure/bridge/bridge.service.js'; +import { + DuplicateEntryKeyError, + MailBucketEntryRepository, +} from './repositories/mail-bucket-entry.repository.js'; + +export interface TrackStoredMessageParams { + mailAccountId: string; + userUuid: string; + bucketId: string; + entryKey: string; + size: number; +} + +export interface ReleaseStoredMessageParams { + userUuid: string; + bucketId: string; + entryKey: string; +} + +@Injectable() +export class MailUsageService { + private readonly logger = new Logger(MailUsageService.name); + + constructor( + private readonly entries: MailBucketEntryRepository, + private readonly bridge: BridgeClient, + ) {} + + async trackStoredMessage(params: TrackStoredMessageParams): Promise { + const { mailAccountId, userUuid, bucketId, entryKey, size } = params; + + const existing = await this.entries.findByEntryKey(entryKey); + if (existing) { + this.logger.debug({ entryKey }, 'Message already tracked; skipping'); + return; + } + + const { id: bridgeEntryId, totalUsedSpaceBytes } = + await this.bridge.createBucketEntry(userUuid, bucketId, size); + + try { + await this.entries.create({ + mailAccountId, + entryKey, + bridgeEntryId, + size, + }); + } catch (error) { + if (error instanceof DuplicateEntryKeyError) { + this.logger.debug( + { entryKey, bridgeEntryId }, + 'Concurrent tracking detected; rolling back minted bucket entry', + ); + await this.bridge.deleteBucketEntry(userUuid, bucketId, bridgeEntryId); + return; + } + throw error; + } + + this.logger.log( + { entryKey, bridgeEntryId, size, totalUsedSpaceBytes }, + 'Tracked stored message', + ); + } + + async releaseStoredMessage( + params: ReleaseStoredMessageParams, + ): Promise { + const { userUuid, bucketId, entryKey } = params; + + const existing = await this.entries.findByEntryKey(entryKey); + if (!existing) { + this.logger.debug( + { entryKey }, + 'No tracked entry for message; skipping release', + ); + return; + } + + await this.bridge.deleteBucketEntry( + userUuid, + bucketId, + existing.bridgeEntryId, + ); + await this.entries.deleteByEntryKey(entryKey); + + this.logger.log( + { entryKey, bridgeEntryId: existing.bridgeEntryId }, + 'Released stored message', + ); + } +} diff --git a/src/modules/usage/models/mail-bucket-entry.model.ts b/src/modules/usage/models/mail-bucket-entry.model.ts new file mode 100644 index 0000000..3df4943 --- /dev/null +++ b/src/modules/usage/models/mail-bucket-entry.model.ts @@ -0,0 +1,46 @@ +import { + AllowNull, + BelongsTo, + Column, + DataType, + Default, + ForeignKey, + Model, + PrimaryKey, + Table, + Unique, +} from 'sequelize-typescript'; +import { MailAccountModel } from '../../account/models/mail-account.model.js'; + +@Table({ + underscored: true, + timestamps: true, + tableName: 'mail_bucket_entries', +}) +export class MailBucketEntryModel extends Model { + @PrimaryKey + @Default(DataType.UUIDV4) + @Column(DataType.UUID) + declare id: string; + + @AllowNull(false) + @ForeignKey(() => MailAccountModel) + @Column(DataType.UUID) + declare mailAccountId: string; + + @AllowNull(false) + @Unique + @Column(DataType.STRING(255)) + declare entryKey: string; + + @AllowNull(false) + @Column(DataType.STRING(24)) + declare bridgeEntryId: string; + + @AllowNull(false) + @Column(DataType.BIGINT) + declare size: string; + + @BelongsTo(() => MailAccountModel) + declare account: MailAccountModel; +} diff --git a/src/modules/usage/repositories/mail-bucket-entry.repository.ts b/src/modules/usage/repositories/mail-bucket-entry.repository.ts new file mode 100644 index 0000000..9d5ce2a --- /dev/null +++ b/src/modules/usage/repositories/mail-bucket-entry.repository.ts @@ -0,0 +1,64 @@ +import { Injectable } from '@nestjs/common'; +import { InjectModel } from '@nestjs/sequelize'; +import { UniqueConstraintError } from 'sequelize'; +import { + MailBucketEntry, + type MailBucketEntryAttributes, +} from '../domain/mail-bucket-entry.domain.js'; +import { MailBucketEntryModel } from '../models/mail-bucket-entry.model.js'; + +export interface CreateMailBucketEntryParams { + mailAccountId: string; + entryKey: string; + bridgeEntryId: string; + size: number; +} + +export class DuplicateEntryKeyError extends Error { + constructor(entryKey: string) { + super(`Bucket entry already tracked for key '${entryKey}'`); + this.name = 'DuplicateEntryKeyError'; + } +} + +@Injectable() +export class MailBucketEntryRepository { + constructor( + @InjectModel(MailBucketEntryModel) + private readonly entryModel: typeof MailBucketEntryModel, + ) {} + + async create(params: CreateMailBucketEntryParams): Promise { + try { + const model = await this.entryModel.create({ ...params }); + return this.toDomain(model); + } catch (error) { + if (error instanceof UniqueConstraintError) { + throw new DuplicateEntryKeyError(params.entryKey); + } + throw error; + } + } + + async findByEntryKey(entryKey: string): Promise { + const model = await this.entryModel.findOne({ where: { entryKey } }); + return model ? this.toDomain(model) : null; + } + + async deleteByEntryKey(entryKey: string): Promise { + await this.entryModel.destroy({ where: { entryKey } }); + } + + private toDomain(model: MailBucketEntryModel): MailBucketEntry { + const attrs: MailBucketEntryAttributes = { + id: model.id, + mailAccountId: model.mailAccountId, + entryKey: model.entryKey, + bridgeEntryId: model.bridgeEntryId, + size: Number(model.size), + createdAt: model.createdAt as Date, + updatedAt: model.updatedAt as Date, + }; + return MailBucketEntry.build(attrs); + } +}