From cede383aac29189dc4509154f28391c3de1e0c10 Mon Sep 17 00:00:00 2001 From: jzunigax2 <125698953+jzunigax2@users.noreply.github.com> Date: Thu, 2 Jul 2026 17:35:06 -0600 Subject: [PATCH] feat(mta-hooks): implement MTA hooks for quota enforcement - Added MTA hooks module with authentication guard, controller, and service. - Implemented `handleRcpt` method in `MtaHooksService` to manage recipient quota checks. - Introduced `getUserUsage` method in `BridgeClient` for fetching user storage usage. - Updated `.env.template` to include MTA hooks credentials. - Added unit tests for MTA hooks functionality and authentication guard to ensure reliability. --- .env.template | 4 + src/app.module.ts | 2 + src/config/configuration.ts | 5 + .../bridge/bridge.service.spec.ts | 43 +++++ .../infrastructure/bridge/bridge.service.ts | 31 ++- .../mta-hooks/mta-hooks-auth.guard.spec.ts | 72 +++++++ src/modules/mta-hooks/mta-hooks-auth.guard.ts | 57 ++++++ src/modules/mta-hooks/mta-hooks.controller.ts | 23 +++ src/modules/mta-hooks/mta-hooks.module.ts | 13 ++ .../mta-hooks/mta-hooks.service.spec.ts | 178 ++++++++++++++++++ src/modules/mta-hooks/mta-hooks.service.ts | 79 ++++++++ src/modules/mta-hooks/mta-hooks.types.ts | 49 +++++ 12 files changed, 555 insertions(+), 1 deletion(-) create mode 100644 src/modules/mta-hooks/mta-hooks-auth.guard.spec.ts create mode 100644 src/modules/mta-hooks/mta-hooks-auth.guard.ts create mode 100644 src/modules/mta-hooks/mta-hooks.controller.ts create mode 100644 src/modules/mta-hooks/mta-hooks.module.ts create mode 100644 src/modules/mta-hooks/mta-hooks.service.spec.ts create mode 100644 src/modules/mta-hooks/mta-hooks.service.ts create mode 100644 src/modules/mta-hooks/mta-hooks.types.ts diff --git a/.env.template b/.env.template index 99c1b5f..426c904 100644 --- a/.env.template +++ b/.env.template @@ -30,3 +30,7 @@ SERVER_PRIVATE_KEY= # Stalwart webhook (ingest events) STALWART_WEBHOOK_USERNAME=stalwart STALWART_WEBHOOK_SECRET= + +# MTA hooks (RCPT) +MTA_HOOKS_USERNAME=stalwart +MTA_HOOKS_SECRET= diff --git a/src/app.module.ts b/src/app.module.ts index d49e129..d425e65 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -15,6 +15,7 @@ import { GatewayModule } from './modules/gateway/gateway.module'; import { HttpGlobalExceptionFilter } from './common/filters/http-global-exception.filter'; import { AddressesModule } from './modules/addresses/addresses.module'; import { StalwartEventsModule } from './modules/stalwart-events/stalwart-events.module'; +import { MtaHooksModule } from './modules/mta-hooks/mta-hooks.module'; @Module({ imports: [ @@ -89,6 +90,7 @@ import { StalwartEventsModule } from './modules/stalwart-events/stalwart-events. AddressesModule, GatewayModule, StalwartEventsModule, + MtaHooksModule, ], controllers: [], providers: [ diff --git a/src/config/configuration.ts b/src/config/configuration.ts index ebb486d..a897e3f 100644 --- a/src/config/configuration.ts +++ b/src/config/configuration.ts @@ -53,4 +53,9 @@ export default () => ({ username: process.env.STALWART_WEBHOOK_USERNAME ?? 'stalwart', secret: process.env.STALWART_WEBHOOK_SECRET ?? '', }, + + mtaHooks: { + username: process.env.MTA_HOOKS_USERNAME ?? 'stalwart', + secret: process.env.MTA_HOOKS_SECRET ?? '', + }, }); diff --git a/src/modules/infrastructure/bridge/bridge.service.spec.ts b/src/modules/infrastructure/bridge/bridge.service.spec.ts index 9910091..aa14691 100644 --- a/src/modules/infrastructure/bridge/bridge.service.spec.ts +++ b/src/modules/infrastructure/bridge/bridge.service.spec.ts @@ -235,4 +235,47 @@ describe('BridgeClient', () => { expect(error.details).toBe('entry not found'); }); }); + + describe('getUserUsage', () => { + it('when Bridge returns 200, then signs a gateway token, GETs the usage, and returns the snapshot', async () => { + const snapshot = { maxSpaceBytes: 5000, totalUsedSpaceBytes: 1200 }; + jwtService.sign.mockReturnValue('signed-jwt'); + httpRequest.mockResolvedValue({ + statusCode: 200, + body: { text: () => Promise.resolve(JSON.stringify(snapshot)) }, + }); + + const result = await service.getUserUsage('user-1'); + + expect(result).toStrictEqual(snapshot); + expect(httpRequest).toHaveBeenCalledWith( + expect.objectContaining({ + method: 'GET', + path: '/v2/gateway/users/user-1/usage', + headers: expect.objectContaining({ + authorization: 'Bearer signed-jwt', + }) as unknown, + }), + ); + }); + + it('when Bridge returns a non-200 status, then throws BridgeApiError with statusCode and details', async () => { + jwtService.sign.mockReturnValue('signed-jwt'); + httpRequest.mockResolvedValue({ + statusCode: 500, + body: { text: () => Promise.resolve('internal error') }, + }); + + const error: unknown = await service + .getUserUsage('user-1') + .catch((e: unknown) => e); + + expect(error).toBeInstanceOf(BridgeApiError); + if (!(error instanceof BridgeApiError)) { + throw new Error('expected BridgeApiError'); + } + expect(error.statusCode).toBe(500); + expect(error.details).toBe('internal error'); + }); + }); }); diff --git a/src/modules/infrastructure/bridge/bridge.service.ts b/src/modules/infrastructure/bridge/bridge.service.ts index 1068606..6c10fac 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 { @@ -157,6 +161,31 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy { } } + async getUserUsage(userUuid: string): Promise { + const token = this.signGatewayToken(userUuid); + + const { statusCode, body } = await this.httpClient.request({ + method: 'GET', + path: `${this.basePath}/v2/gateway/users/${encodeURIComponent(userUuid)}/usage`, + headers: { + accept: 'application/json', + authorization: `Bearer ${token}`, + }, + }); + + const text = await body.text(); + + if (statusCode !== 200) { + throw new BridgeApiError( + `Failed to fetch usage for user '${userUuid}': HTTP ${statusCode}`, + statusCode, + text, + ); + } + + return JSON.parse(text) as UserSpaceSnapshot; + } + private signGatewayToken(userUuid: string): string { return this.jwtService.sign( { payload: { uuid: userUuid } }, diff --git a/src/modules/mta-hooks/mta-hooks-auth.guard.spec.ts b/src/modules/mta-hooks/mta-hooks-auth.guard.spec.ts new file mode 100644 index 0000000..0568821 --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks-auth.guard.spec.ts @@ -0,0 +1,72 @@ +import { describe, it, expect, beforeEach } from 'vitest'; +import { createMock } from '@golevelup/ts-vitest'; +import { type ConfigService } from '@nestjs/config'; +import { UnauthorizedException, type ExecutionContext } from '@nestjs/common'; +import { MtaHooksAuthGuard } from './mta-hooks-auth.guard.js'; + +const username = 'stalwart'; +const secret = 'secret'; + +describe('MtaHooksAuthGuard', () => { + let guard: MtaHooksAuthGuard; + + const contextWithAuth = (header?: string): ExecutionContext => + createMock({ + switchToHttp: () => ({ + getRequest: () => ({ + headers: header ? { authorization: header } : {}, + }), + }), + }); + + const basic = (username: string, secret: string): string => + `Basic ${Buffer.from(`${username}:${secret}`).toString('base64')}`; + + beforeEach(() => { + const configService = createMock(); + configService.getOrThrow.mockImplementation((key: string) => { + if (key === 'mtaHooks.username') return username; + if (key === 'mtaHooks.secret') return secret; + throw new Error(`unknown key: ${key}`); + }); + guard = new MtaHooksAuthGuard(configService); + }); + + it('when credentials match, then allows the request', () => { + const context = contextWithAuth(basic(username, secret)); + + expect(guard.canActivate(context)).toBe(true); + }); + + it('when the secret is wrong, then throws Unauthorized', () => { + const context = contextWithAuth(basic(username, 'wrong')); + + expect(() => guard.canActivate(context)).toThrow(UnauthorizedException); + }); + + it('when the username is wrong, then throws Unauthorized', () => { + const context = contextWithAuth(basic('wrong', secret)); + + expect(() => guard.canActivate(context)).toThrow(UnauthorizedException); + }); + + it('when the Authorization header is missing, then throws Unauthorized', () => { + const context = contextWithAuth(undefined); + + expect(() => guard.canActivate(context)).toThrow(UnauthorizedException); + }); + + it('when the scheme is not Basic, then throws Unauthorized', () => { + const context = contextWithAuth('Bearer some-token'); + + expect(() => guard.canActivate(context)).toThrow(UnauthorizedException); + }); + + it('when the decoded credentials lack a colon separator, then throws Unauthorized', () => { + const context = contextWithAuth( + `Basic ${Buffer.from('nocolon').toString('base64')}`, + ); + + expect(() => guard.canActivate(context)).toThrow(UnauthorizedException); + }); +}); diff --git a/src/modules/mta-hooks/mta-hooks-auth.guard.ts b/src/modules/mta-hooks/mta-hooks-auth.guard.ts new file mode 100644 index 0000000..2cd27b8 --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks-auth.guard.ts @@ -0,0 +1,57 @@ +import { timingSafeEqual } from 'node:crypto'; +import { + CanActivate, + type ExecutionContext, + Injectable, + UnauthorizedException, +} from '@nestjs/common'; +import { ConfigService } from '@nestjs/config'; +import type { Request } from 'express'; + +@Injectable() +export class MtaHooksAuthGuard implements CanActivate { + private readonly expectedUsername: string; + private readonly expectedSecret: string; + + constructor(configService: ConfigService) { + this.expectedUsername = + configService.getOrThrow('mtaHooks.username'); + this.expectedSecret = configService.getOrThrow('mtaHooks.secret'); + } + + canActivate(context: ExecutionContext): boolean { + const request = context.switchToHttp().getRequest(); + const header = request.headers.authorization ?? ''; + + const [scheme, encoded] = header.split(' '); + if (scheme !== 'Basic' || !encoded) { + throw new UnauthorizedException('Missing or malformed Basic credentials'); + } + + const decoded = Buffer.from(encoded, 'base64').toString('utf8'); + const separatorIndex = decoded.indexOf(':'); + if (separatorIndex === -1) { + throw new UnauthorizedException('Malformed Basic credentials'); + } + + const username = decoded.slice(0, separatorIndex); + const secret = decoded.slice(separatorIndex + 1); + + const usernameMatches = this.safeEqual(username, this.expectedUsername); + const secretMatches = this.safeEqual(secret, this.expectedSecret); + if (!usernameMatches || !secretMatches) { + throw new UnauthorizedException('Invalid MTA hook credentials'); + } + + return true; + } + + private safeEqual(a: string, b: string): boolean { + const bufferA = Buffer.from(a, 'utf8'); + const bufferB = Buffer.from(b, 'utf8'); + if (bufferA.length !== bufferB.length) { + return false; + } + return timingSafeEqual(bufferA, bufferB); + } +} diff --git a/src/modules/mta-hooks/mta-hooks.controller.ts b/src/modules/mta-hooks/mta-hooks.controller.ts new file mode 100644 index 0000000..b85e2cf --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks.controller.ts @@ -0,0 +1,23 @@ +import { Body, Controller, Post, UseGuards } from '@nestjs/common'; +import { ApiBasicAuth, ApiOperation, ApiTags } from '@nestjs/swagger'; +import { Public } from '../auth/decorators/public.decorator.js'; +import { MtaHooksAuthGuard } from './mta-hooks-auth.guard.js'; +import { MtaHooksService } from './mta-hooks.service.js'; +import type { MtaHookRequest, MtaHookResponse } from './mta-hooks.types.js'; + +@ApiTags('MTA Hooks') +@ApiBasicAuth('mta-hooks') +@Public() +@UseGuards(MtaHooksAuthGuard) +@Controller('mta-hooks') +export class MtaHooksController { + constructor(private readonly mtaHooksService: MtaHooksService) {} + + @Post('rcpt') + @ApiOperation({ + summary: 'RCPT-stage hook', + }) + rcpt(@Body() request: MtaHookRequest): Promise { + return this.mtaHooksService.handleRcpt(request); + } +} diff --git a/src/modules/mta-hooks/mta-hooks.module.ts b/src/modules/mta-hooks/mta-hooks.module.ts new file mode 100644 index 0000000..850b794 --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks.module.ts @@ -0,0 +1,13 @@ +import { Module } from '@nestjs/common'; +import { AccountModule } from '../account/account.module.js'; +import { BridgeModule } from '../infrastructure/bridge/bridge.module.js'; +import { MtaHooksAuthGuard } from './mta-hooks-auth.guard.js'; +import { MtaHooksController } from './mta-hooks.controller.js'; +import { MtaHooksService } from './mta-hooks.service.js'; + +@Module({ + imports: [AccountModule, BridgeModule], + controllers: [MtaHooksController], + providers: [MtaHooksService, MtaHooksAuthGuard], +}) +export class MtaHooksModule {} diff --git a/src/modules/mta-hooks/mta-hooks.service.spec.ts b/src/modules/mta-hooks/mta-hooks.service.spec.ts new file mode 100644 index 0000000..3236b23 --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks.service.spec.ts @@ -0,0 +1,178 @@ +import { describe, it, expect, beforeEach } from 'vitest'; +import { Test, type TestingModule } from '@nestjs/testing'; +import { createMock, type DeepMocked } from '@golevelup/ts-vitest'; +import { AccountService } from '../account/account.service.js'; +import { BridgeClient } from '../infrastructure/bridge/bridge.service.js'; +import { MtaHooksService } from './mta-hooks.service.js'; +import type { MtaHookRequest } from './mta-hooks.types.js'; + +describe('MtaHooksService', () => { + let service: MtaHooksService; + let accountService: DeepMocked; + let bridgeClient: DeepMocked; + + const buildRequest = ( + to: string[], + opts: { size?: number } = {}, + ): MtaHookRequest => ({ + context: { stage: 'rcpt' }, + envelope: { + from: { + address: 'sender@external.com', + parameters: + opts.size !== undefined ? { size: String(opts.size) } : undefined, + }, + to: to.map((address) => ({ address })), + }, + }); + + beforeEach(async () => { + const module: TestingModule = await Test.createTestingModule({ + providers: [MtaHooksService], + }) + .useMocker(() => createMock()) + .compile(); + + service = module.get(MtaHooksService); + accountService = module.get(AccountService); + bridgeClient = module.get(BridgeClient); + }); + + describe('handleRcpt', () => { + it('when the recipient stays within quota, then accepts', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockResolvedValue({ + maxSpaceBytes: 5000, + totalUsedSpaceBytes: 2000, + }); + + const result = await service.handleRcpt( + buildRequest(['jane@inxt.com'], { size: 500 }), + ); + + expect(result).toStrictEqual({ action: 'accept' }); + expect(accountService.findUserIdByAddress).toHaveBeenCalledWith( + 'jane@inxt.com', + ); + expect(bridgeClient.getUserUsage).toHaveBeenCalledWith('user-1'); + }); + + it('when the declared SIZE pushes the recipient over quota, then rejects with 452 4.2.2', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockResolvedValue({ + maxSpaceBytes: 5000, + totalUsedSpaceBytes: 4800, + }); + + const result = await service.handleRcpt( + buildRequest(['jane@inxt.com'], { size: 500 }), + ); + + expect(result).toStrictEqual({ + action: 'reject', + response: { + status: 452, + enhancedStatus: '4.2.2', + message: 'Recipient mailbox is over quota', + }, + }); + }); + + it('when projected usage exactly equals the quota, then accepts', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockResolvedValue({ + maxSpaceBytes: 5000, + totalUsedSpaceBytes: 4500, + }); + + const result = await service.handleRcpt( + buildRequest(['jane@inxt.com'], { size: 500 }), + ); + + expect(result).toStrictEqual({ action: 'accept' }); + }); + + it('when no SIZE is declared but the mailbox is already over quota, then rejects', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockResolvedValue({ + maxSpaceBytes: 5000, + totalUsedSpaceBytes: 5500, + }); + + const result = await service.handleRcpt(buildRequest(['jane@inxt.com'])); + + expect(result.action).toBe('reject'); + }); + + it('when the address resolves to no internxt user, then skips it and accepts', async () => { + accountService.findUserIdByAddress.mockResolvedValue(null); + + const result = await service.handleRcpt( + buildRequest(['external@gmail.com'], { size: 500 }), + ); + + expect(result).toStrictEqual({ action: 'accept' }); + expect(bridgeClient.getUserUsage).not.toHaveBeenCalled(); + }); + + it('when the recipient address is upper-cased, then it is lowercased before resolution', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockResolvedValue({ + maxSpaceBytes: 5000, + totalUsedSpaceBytes: 0, + }); + + await service.handleRcpt(buildRequest(['Jane@INXT.com'], { size: 10 })); + + expect(accountService.findUserIdByAddress).toHaveBeenCalledWith( + 'jane@inxt.com', + ); + }); + + it('when several recipients are present, then only the current (last) one is evaluated', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockResolvedValue({ + maxSpaceBytes: 5000, + totalUsedSpaceBytes: 0, + }); + + await service.handleRcpt( + buildRequest(['first@inxt.com', 'current@inxt.com']), + ); + + expect(accountService.findUserIdByAddress).toHaveBeenCalledTimes(1); + expect(accountService.findUserIdByAddress).toHaveBeenCalledWith( + 'current@inxt.com', + ); + }); + + it('when recipient resolution throws, then fails open and accepts', async () => { + accountService.findUserIdByAddress.mockRejectedValue( + new Error('database down'), + ); + + const result = await service.handleRcpt(buildRequest(['jane@inxt.com'])); + + expect(result).toStrictEqual({ action: 'accept' }); + expect(bridgeClient.getUserUsage).not.toHaveBeenCalled(); + }); + + it('when the Bridge usage lookup throws, then fails open and accepts', async () => { + accountService.findUserIdByAddress.mockResolvedValue('user-1'); + bridgeClient.getUserUsage.mockRejectedValue(new Error('bridge down')); + + const result = await service.handleRcpt(buildRequest(['jane@inxt.com'])); + + expect(result).toStrictEqual({ action: 'accept' }); + }); + + it('when the request carries no recipients, then accepts without lookups', async () => { + const result = await service.handleRcpt({ + context: { stage: 'rcpt' }, + }); + + expect(result).toStrictEqual({ action: 'accept' }); + expect(accountService.findUserIdByAddress).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/src/modules/mta-hooks/mta-hooks.service.ts b/src/modules/mta-hooks/mta-hooks.service.ts new file mode 100644 index 0000000..b934f5e --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks.service.ts @@ -0,0 +1,79 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { AccountService } from '../account/account.service.js'; +import { BridgeClient } from '../infrastructure/bridge/bridge.service.js'; +import type { + MtaHookEnvelope, + MtaHookRequest, + MtaHookResponse, +} from './mta-hooks.types.js'; + +const ACCEPT: MtaHookResponse = { action: 'accept' }; + +const REJECT_OVER_QUOTA: MtaHookResponse = { + action: 'reject', + response: { + status: 452, + enhancedStatus: '4.2.2', + message: 'Recipient mailbox is over quota', + }, +}; + +@Injectable() +export class MtaHooksService { + private readonly logger = new Logger(MtaHooksService.name); + + constructor( + private readonly accountService: AccountService, + private readonly bridgeClient: BridgeClient, + ) {} + + async handleRcpt(request: MtaHookRequest): Promise { + const recipients = request.envelope?.to ?? []; + const recipient = recipients.at(-1); + if (!recipient) { + return ACCEPT; + } + + const declaredSize = this.parseDeclaredSize(request.envelope); + + try { + const overQuota = await this.isRecipientOverQuota( + recipient.address, + declaredSize, + ); + return overQuota ? REJECT_OVER_QUOTA : ACCEPT; + } catch (error) { + this.logger.error( + `MTA hook quota check failed, accepting recipient (fail-open): ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return ACCEPT; + } + } + + private async isRecipientOverQuota( + rawAddress: string, + incomingSize: number, + ): Promise { + const address = rawAddress.toLowerCase(); + const userUuid = await this.accountService.findUserIdByAddress(address); + if (!userUuid) { + return false; + } + + const { maxSpaceBytes, totalUsedSpaceBytes } = + await this.bridgeClient.getUserUsage(userUuid); + + return totalUsedSpaceBytes + incomingSize > maxSpaceBytes; + } + + private parseDeclaredSize(envelope?: MtaHookEnvelope): number { + const raw = envelope?.from.parameters?.size; + if (!raw) { + return 0; + } + const parsed = Number.parseInt(raw, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : 0; + } +} diff --git a/src/modules/mta-hooks/mta-hooks.types.ts b/src/modules/mta-hooks/mta-hooks.types.ts new file mode 100644 index 0000000..6050781 --- /dev/null +++ b/src/modules/mta-hooks/mta-hooks.types.ts @@ -0,0 +1,49 @@ +/** + * Types for Stalwart's MTA Hooks protocol. + * + * Stalwart POSTs a JSON {@link MtaHookRequest} to the configured endpoint at a + * given SMTP stage and expects a JSON {@link MtaHookResponse} telling it whether + * to accept or reject the message. + * + * We only model the fields this service reads + * + * @see https://stalw.art/docs/api/mta-hooks/overview + */ + +export type MtaHookStage = + | 'connect' + | 'ehlo' + | 'auth' + | 'mail' + | 'rcpt' + | 'data'; + +export interface MtaHookAddress { + address: string; + parameters?: Record | null; +} + +export interface MtaHookEnvelope { + from: MtaHookAddress; + to: MtaHookAddress[]; +} + +export interface MtaHookRequest { + context: { + stage: MtaHookStage; + }; + envelope?: MtaHookEnvelope; +} + +export type MtaHookAction = 'accept' | 'reject'; + +export interface MtaHookSmtpResponse { + status?: number; + enhancedStatus?: string; + message?: string; +} + +export interface MtaHookResponse { + action: MtaHookAction; + response?: MtaHookSmtpResponse; +}