Auth Service

import { Inject, Injectable, Logger } from "@nestjs/common";
import { ConfigService } from "@nestjs/config";
import type { FastifySessionObject } from "@fastify/session";

import { randomUUID } from "node:crypto";
import { addMinutes, compareDesc, getUnixTime } from "date-fns";

import type { RedisClientType } from "redis";

import type { ClientAuth, ClientMetadata, ServerMetadata, ServerMetadataHelpers, TokenEndpointResponse } from "openid-client";
import * as oidc from "openid-client";
import type { JSONWebKeySet, JWK, JWSHeaderParameters, JWTPayload, JWTVerifyResult } from "jose";
import * as jose from "jose";

import { DefaultAzureCredential } from "@azure/identity";
import type { GetKeyOptions, JsonWebKey, KeyProperties, KeyVaultKey, SignResult } from "@azure/keyvault-keys";
import { CryptographyClient, KeyClient, KnownSignatureAlgorithms } from "@azure/keyvault-keys";

import { REDIS_CLIENT } from "../app.symbols.js";
import { assertExists } from "../app.utils.ts";
import type { Provider } from "../app.config.ts";

interface IDTokenPayload extends JWTPayload {
    sub: string;
    email?: string;
}

interface LogoutTokenPayload extends JWTPayload {
    sid?: string;
    events?: {
        "http://schemas.openid.net/event/backchannel-logout"?: Record<string, unknown>;
    };
}

@Injectable()
export class AuthService {
    private readonly logger: Logger = new Logger(AuthService.name);
    private readonly configurations: Map<string, oidc.Configuration> = new Map<string, oidc.Configuration>();

    constructor(
        private readonly configService: ConfigService,
        @Inject(REDIS_CLIENT) private readonly redisClient: RedisClientType,
    ) {}

    async onModuleInit(): Promise<void> {
        const providers: Provider[] = assertExists(this.configService.get<Provider[]>(`providers`));
        await Promise.all(
            providers.map(async (provider: Provider): Promise<void> => {
                const issuer: URL = new URL(provider.issuer);

                const clientMetadata: ClientMetadata = provider.metadata as ClientMetadata;
                const clientId: string = clientMetadata.client_id;

                const configuration: oidc.Configuration = await oidc.discovery(issuer, clientId, clientMetadata);
                this.configurations.set(provider.name, configuration);
            }),
        );
    }

    async loginRedirect(provider: Provider, session: FastifySessionObject): Promise<URL> {
        const configuration: oidc.Configuration = assertExists(this.configurations.get(provider.name));

        const codeVerifier: string = oidc.randomPKCECodeVerifier();
        const codeChallenge: string = await oidc.calculatePKCECodeChallenge(codeVerifier);
        const state: string = oidc.randomState();

        session.codeVerifier = codeVerifier;
        session.state = state;

        const parameters: Record<string, string> = {
            redirect_uri: provider.loginCallbackUrl,
            code_challenge_method: "S256",
            code_challenge: codeChallenge,
            state: state,
            scope: provider.scope,
            resource: provider.resource,
        };

        return oidc.buildAuthorizationUrl(configuration, parameters);
    }

    async loginCallback(provider: Provider, url: URL, session: FastifySessionObject): Promise<void> {
        const configuration: oidc.Configuration = assertExists(this.configurations.get(provider.name));
        const clientMetadata: Readonly<ClientMetadata> = configuration.clientMetadata();
        const serverMetadata: Readonly<ServerMetadata> & ServerMetadataHelpers = configuration.serverMetadata();
        const clientId: string = assertExists(clientMetadata.client_id);
        const jwksUri: string = assertExists(serverMetadata.jwks_uri);
        const jwksUrl: URL = new URL(jwksUri);

        const clientAssertion: string = await this.privateKeyJwt(provider);
        const clientAuth: ClientAuth = (server: ServerMetadata, client: ClientMetadata, body: URLSearchParams): void => {
            body.set("client_assertion_type", "urn:ietf:params:oauth:client-assertion-type:jwt-bearer");
            body.set("client_assertion", clientAssertion);
        };
        const requestConfiguration: oidc.Configuration = new oidc.Configuration(serverMetadata, clientId, clientMetadata, clientAuth);

        const codeVerifier: string = assertExists(session.codeVerifier, "Session does not have 'codeVerifier' property.");
        const state: string = assertExists(session.state, "Session does not have 'state' property.");
        const parameters: Record<string, string> = {
            pkceCodeVerifier: codeVerifier,
            expectedState: state,
        };

        const tokens: TokenEndpointResponse = await oidc.authorizationCodeGrant(requestConfiguration, url, parameters);
        const idToken: string = assertExists(tokens.id_token);

        const jwtVerifyResult: JWTVerifyResult<IDTokenPayload> = await jose.jwtVerify<IDTokenPayload>(idToken, jose.createRemoteJWKSet(jwksUrl));
        const email: string = assertExists(jwtVerifyResult.payload.email);

        await this.redisClient.multi().hSet(session.sessionId, "access_token", assertExists(tokens.access_token)).hSet(session.sessionId, "refresh_token", assertExists(tokens.refresh_token)).hSet(session.sessionId, "id_token", assertExists(tokens.id_token)).exec();

        delete session.codeVerifier;
        delete session.state;

        session.email = email;
    }

    async logoutRedirect(provider: Provider, session: FastifySessionObject): Promise<URL> {
        const configuration: oidc.Configuration = assertExists(this.configurations.get(provider.name));

        const idToken: string = assertExists(await this.redisClient.hGet(session.sessionId, "id_token"));

        const parameters: Record<string, string> = {
            next: provider.logoutCallbackUrl,
            post_logout_redirect_uri: provider.logoutCallbackUrl,
            id_token_hint: idToken,
        };

        return oidc.buildEndSessionUrl(configuration, parameters);
    }

    async logoutCallback(provider: Provider, session: FastifySessionObject): Promise<void> {
        await this.redisClient.del(session.sessionId);
    }

    async backchannelLogout(provider: Provider, logoutToken: string): Promise<void> {
        const configuration: oidc.Configuration = assertExists(this.configurations.get(provider.name));
        const serverMetadata: Readonly<ServerMetadata> & ServerMetadataHelpers = configuration.serverMetadata();
        const jwksUri: string = assertExists(serverMetadata.jwks_uri);
        const jwksUrl: URL = new URL(jwksUri);

        const jwtVerifyResult: JWTVerifyResult<LogoutTokenPayload> = await jose.jwtVerify<LogoutTokenPayload>(logoutToken, jose.createRemoteJWKSet(jwksUrl));

        this.logger.log(`sub: ${jwtVerifyResult.payload.sub}`);
        this.logger.log(`sid: ${jwtVerifyResult.payload.sid}`);

        assertExists(jwtVerifyResult.payload.sub);
    }

    async session(session: FastifySessionObject): Promise<boolean> {
        return Boolean(await this.redisClient.exists((session.sessionId)));
    }

    async jwkSet(): Promise<JSONWebKeySet> {
        const managedIdentity: string = assertExists(this.configService.get<string>(`azure.managedIdentity`), "Configuration 'azure.managedIdentity' property is undefined.");
        const keyVaultName: string = assertExists(this.configService.get<string>(`azure.keyVaultName`), "Configuration 'azure.keyVaultName' property is undefined.");
        const keyVaultUrl = `https://${keyVaultName}.vault.azure.net`;

        const azureCredential: DefaultAzureCredential = new DefaultAzureCredential({
            managedIdentityClientId: managedIdentity,
        });
        const keyClient = new KeyClient(keyVaultUrl, azureCredential);

        const ecKeyName: string = assertExists(this.configService.get<string>(`azure.ecKeyName`), "Configuration 'azure.ecKeyName' property is undefined.");

        const ecKeyVersions: KeyProperties[] = [];
        for await (const ecKeyProperty of keyClient.listPropertiesOfKeyVersions(ecKeyName)) {
            ecKeyVersions.push(ecKeyProperty);
        }
        ecKeyVersions.sort((a: KeyProperties, b: KeyProperties): number => {
            return compareDesc(assertExists(a.createdOn), assertExists(b.createdOn));
        });

        const jwkSet: JSONWebKeySet = {
            keys: [],
        };
        for (const ecKeyVersion of ecKeyVersions.slice(0, 2)) {
            const ecKeyOptions: GetKeyOptions = {
                version: assertExists(ecKeyVersion.version),
            };
            const ecKey: KeyVaultKey = assertExists(await keyClient.getKey(ecKeyName, ecKeyOptions), "Configuration 'azure.ecKeyName' property is undefined.");
            const ecJwk: JsonWebKey = assertExists(ecKey.key);

            const jwk: JWK = {
                kid: assertExists(ecKeyVersion.version),
                kty: assertExists(ecJwk.kty),
                crv: assertExists(ecJwk.crv),
                x: jose.base64url.encode(assertExists(ecJwk.x)),
                y: jose.base64url.encode(assertExists(ecJwk.y)),
                alg: KnownSignatureAlgorithms.ES256,
            };

            jwkSet.keys.push(jwk);
        }

        return jwkSet;
    }

    async privateKeyJwt(provider: Provider): Promise<string> {
        const configuration: oidc.Configuration = assertExists(this.configurations.get(provider.name));
        const clientMetadata: Readonly<ClientMetadata> = configuration.clientMetadata();
        const serverMetadata: Readonly<ServerMetadata> & ServerMetadataHelpers = configuration.serverMetadata();
        const clientId: string = assertExists(clientMetadata.client_id);
        const tokenEndpoint: string = assertExists(serverMetadata.token_endpoint);

        const managedIdentity: string = assertExists(this.configService.get<string>(`azure.managedIdentity`), "Configuration 'azure.managedIdentity' property is undefined.");
        const keyVaultName: string = assertExists(this.configService.get<string>(`azure.keyVaultName`), "Configuration 'azure.keyVaultName' property is undefined.");
        const keyVaultUrl = `https://${keyVaultName}.vault.azure.net`;

        const azureCredential: DefaultAzureCredential = new DefaultAzureCredential({
            managedIdentityClientId: managedIdentity,
        });
        const keyClient = new KeyClient(keyVaultUrl, azureCredential);

        const ecKeyName: string = assertExists(this.configService.get<string>(`azure.ecKeyName`), "Configuration 'azure.ecKeyName' property is undefined.");
        const ecKey: KeyVaultKey = assertExists(await keyClient.getKey(ecKeyName), "Configuration 'azure.ecKeyName' property is undefined.");
        const ecKeyVersion: string = assertExists(ecKey.properties.version);

        const header: JWSHeaderParameters = {
            alg: KnownSignatureAlgorithms.ES256,
            kid: ecKeyVersion,
        };
        const encodedHeader: string = jose.base64url.encode(JSON.stringify(header));

        const now: Date = new Date();
        const payload: JWTPayload = {
            jti: randomUUID(),
            sub: clientId,
            iss: clientId,
            aud: tokenEndpoint,
            iat: getUnixTime(now),
            exp: getUnixTime(addMinutes(now, 5)),
        };
        const encodedPayload: string = jose.base64url.encode(JSON.stringify(payload));

        const cryptoClient = new CryptographyClient(ecKey, azureCredential);
        const signatureResult: SignResult = await cryptoClient.signData(KnownSignatureAlgorithms.ES256, Buffer.from(`${encodedHeader}.${encodedPayload}`));
        const signature: Uint8Array<ArrayBufferLike> = signatureResult.result;
        const encodedSignature: string = Buffer.from(signature).toString("base64url");

        return `${encodedHeader}.${encodedPayload}.${encodedSignature}`;
    }
}