TREK/server/src/nest/oauth/oauth.module.ts

import { OauthClients } from '../../db/entities/OauthClients.entity';
import { OauthConsents } from '../../db/entities/OauthConsents.entity';
import { OauthTokens } from '../../db/entities/OauthTokens.entity';
import { Users } from '../../db/entities/Users.entity';
import { AddonsModule } from '../addons/addons.module';
import { AddonsService } from '../addons/addons.service';
import { createMcpAddonGate } from '../addons/mcp-addon-gate';
import { AuditModule } from '../audit/audit.module';
import { RateLimitModule } from '../common/rate-limit.module';
import { SchedulingModule } from '../scheduling/scheduling.module';
import { OauthApiController } from './oauth-api.controller';
import { CimdClientResolver } from './oauth-cimd';
import { OauthPublicController } from './oauth-public.controller';
import { TrekClientsStore, TrekOAuthProvider } from './oauth-sdk.provider';
import { OauthTokenRetentionJob } from './oauth-token-retention.job';
import { PendingCodeStore, pendingCodesSlot, processPendingCodes } from './oauth.pending-codes';
import { OauthService } from './oauth.service';
import { MikroOrmModule } from '@mikro-orm/nestjs';
import { authorizationHandler } from '@modelcontextprotocol/sdk/server/auth/handlers/authorize';
import { clientRegistrationHandler } from '@modelcontextprotocol/sdk/server/auth/handlers/register';
import { Module, type OnModuleDestroy } from '@nestjs/common';
import type { MiddlewareConsumer, NestModule } from '@nestjs/common';

/**
 * OAuth 2.1 server (MCP). Public token/userinfo/revoke endpoints + the SPA's
 * authenticated consent/client/session management, plus the SDK's own
 * /oauth/authorize and /oauth/register Express routers, mounted here through
 * MiddlewareConsumer over the injected TrekOAuthProvider/TrekClientsStore.
 * forRoutes('<path>') is an Express prefix mount — the same app.use('/path')
 * shape (and the same req.url stripping) the pre-init platform mount used, and
 * exactly what the SDK routers are built for. The handlers are constructed once
 * per app instance in configure(), preserving the per-app lifetime of the DCR
 * handler's built-in rate limiter.
 *
 * The handlers bring their own body parsers; the global 100kb parsers run
 * first and theirs no-op on the already-consumed stream. The one observable
 * difference vs the pre-init mount: bracketed keys in a POST /oauth/authorize
 * body now parse with qs (extended:true) rather than querystring — the SDK
 * zod-rejects such bodies either way.
 *
 * Pending authorization codes live behind the PendingCodeStore port
 * (oauth.pending-codes.ts), provided here as the process-wide in-memory
 * instance: the consent controller (container singleton) writes them, the
 * SDK exchange path reads them back through the same injected singleton. The
 * constructor installs whichever store the container resolved in
 * pendingCodesSlot, so the sweep and a hand-built OauthService follow a
 * swapped provider too.
 *
 * Client ID Metadata Documents (oauth-cimd.ts) resolve through the injected
 * CimdClientResolver.
 *
 * Exports OauthService for AdminController (admin OAuth-session panel) and the
 * MCP transport's token verification.
 *
 * `MikroOrmModule.forFeature([OauthClients, OauthTokens, OauthConsents,
 * Users])` (Plan 3b Task 4 deliverable): `Users` is registered for cross-
 * module availability (inventory §6 — `OauthTokensRepository` needs to be
 * importable from `AuthModule` too, Task 5's `revokeAllForUser`, the same
 * shape `McpTokensRepository`/`InviteTokensRepository` already use across
 * `nest/tokens`/`nest/auth`/`nest/oidc`) even though `OauthService` itself
 * never injects `UsersRepository` directly — OA16/OA31 read `users` columns
 * THROUGH `OauthTokensRepository`'s own join methods, not a separate
 * `Users` call (Task 1's ruling: "do not add users methods" for this join).
 */
@Module({
  imports: [
    RateLimitModule,
    AuditModule,
    AddonsModule,
    SchedulingModule,
    MikroOrmModule.forFeature([OauthClients, OauthTokens, OauthConsents, Users]),
  ],
  controllers: [OauthPublicController, OauthApiController],
  providers: [
    OauthService,
    CimdClientResolver,
    TrekClientsStore,
    TrekOAuthProvider,
    OauthTokenRetentionJob,
    { provide: PendingCodeStore, useValue: processPendingCodes },
  ],
  exports: [OauthService],
})
export class OauthModule implements NestModule, OnModuleDestroy {
  constructor(
    private readonly addons: AddonsService,
    private readonly provider: TrekOAuthProvider,
    private readonly clients: TrekClientsStore,
    private readonly pendingCodes: PendingCodeStore,
  ) {
    pendingCodesSlot.install(pendingCodes);
  }

  onModuleDestroy(): void {
    pendingCodesSlot.release(this.pendingCodes);
  }

  configure(consumer: MiddlewareConsumer): void {
    const mcpAddonGate = createMcpAddonGate(this.addons);

    // SDK authorize handler: validates OAuth params, calls provider.authorize()
    // which redirects to the SPA consent page at /oauth/consent
    consumer.apply(mcpAddonGate, authorizationHandler({ provider: this.provider })).forRoutes('oauth/authorize');

    // SDK DCR handler: accepts registrations without scope (fixes issue #959 bug 2)
    consumer.apply(mcpAddonGate, clientRegistrationHandler({ clientsStore: this.clients })).forRoutes('oauth/register');
  }
}