edr-platform/apps/edr-freight-api/src/common/freight-jwt.guard.ts

import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { Reflector } from '@nestjs/core';
import { InjectDataSource } from '@nestjs/typeorm';
import { JwtGuard as IamJwtGuard } from '@tria-plc/api-common/modules/auth/services/jwt.guard';
import type { TCurrentUser } from '@tria-plc/api-common/modules/auth/types/current-user.type';
import { CURRENT_POSITION_ID } from '@tria-plc/api-common/utils/constants/tenant.constant';
import { DataSource } from 'typeorm';

/** One position as the login snapshot stores it (`iam.sessions.userInfo`). */
export type SnapshotPosition = {
  id?: string;
  employeePositionId?: string;
  isDelegate?: boolean;
  [key: string]: unknown;
};

/** One employee row as the snapshot stores it. A user may hold several. */
export type SnapshotEmployee = {
  id?: string;
  positions?: SnapshotPosition[];
  [key: string]: unknown;
};

type SessionUserInfo = { employee?: SnapshotEmployee[] };

/**
 * `x-current-position-id` is sent with two different meanings by two different
 * frontends, and the IAM guard reads it both ways in the same function: it
 * picks the EMPLOYEE row by `position.id` but the POSITION by
 * `employeePositionId`. Freight sends `employeePositionId`, Smart Office sends
 * `position.id` — so whichever value arrives, one of the two lookups silently
 * matches nothing and falls back to the first entry.
 *
 * Matching both fields is what makes the header mean one thing again.
 */
const identifies = (position: SnapshotPosition, id: string): boolean =>
  position?.id === id || position?.employeePositionId === id;

/**
 * Every post the user holds, across every employee row, first occurrence kept.
 *
 * IAM keeps one employee row per ORGANIZATION, and "EDR" and "EDR Freight" are
 * separate organizations — so a user given a freight post and a Smart Office
 * post owns two rows, one post on each. Only one row can be the active one, and
 * a permission check that reads only that row cannot see the other post at all.
 */
export const collectAllPositions = (
  employees: SnapshotEmployee[],
): SnapshotPosition[] => {
  const seen = new Set<string>();
  const all: SnapshotPosition[] = [];

  for (const employee of employees) {
    for (const position of employee.positions ?? []) {
      const key = position.employeePositionId ?? position.id;
      if (key) {
        if (seen.has(key)) continue;
        seen.add(key);
      }
      all.push(position);
    }
  }

  return all;
};

/**
 * Which employee row the caller is acting as, and which of its positions the
 * request selected. Pure so it can be tested without a session or a token.
 *
 * `owner` is the row holding the requested position; failing that the row the
 * parent guard already picked; failing that the first. `active` is undefined
 * when no header was sent or it names nothing — the caller then leaves the
 * parent's choice of `employee.position` alone.
 */
export const resolveActiveEmployee = (
  employees: SnapshotEmployee[],
  requestedId: string | undefined,
  parentEmployeeId: string | undefined,
): { owner: SnapshotEmployee | undefined; active: SnapshotPosition | undefined } => {
  const owner =
    (requestedId &&
      employees.find((candidate) =>
        (candidate.positions ?? []).some((position) =>
          identifies(position, requestedId),
        ),
      )) ||
    employees.find(
      (candidate) => candidate.id && candidate.id === parentEmployeeId,
    ) ||
    employees[0];

  const active = requestedId
    ? (owner?.positions ?? []).find((position) =>
        identifies(position, requestedId),
      )
    : undefined;

  return { owner, active };
};

/**
 * Like the IAM JwtGuard, but resolves the caller's position honestly.
 *
 * IAM models an employee as holding many positions — and a user as possibly
 * holding several employee rows — and the login snapshot in
 * `iam.sessions.userInfo` carries all of them. `JwtGuard.parseToken` collapses
 * that to a single `employee.position` and drops the rest, so staff holding two
 * posts resolve to one post's permissions and every check on the other one
 * rejects them.
 *
 * This guard re-reads the snapshot and fixes three things the parent gets wrong:
 *
 *  1. re-attaches the full position list as `employee.positions`, which is what
 *     the permission utils union over;
 *  2. selects the employee row that actually owns the requested position, so a
 *     post held on a second employee row is reachable at all;
 *  3. sets `employee.position` to the requested position when the parent's
 *     one-sided id match missed it, keeping `auditUser` in step.
 *
 * Every correction is skipped unless the snapshot positively resolves it, so an
 * unreadable session degrades to the parent's single-position behaviour rather
 * than to no position at all.
 */
@Injectable()
export class FreightJwtGuard extends IamJwtGuard implements CanActivate {
  // ponytail: unbounded-until-TTL map, cleared wholesale when it gets big.
  // Sessions are few and the value is small; swap for an LRU if that changes.
  private static readonly CACHE_TTL_MS = 30_000;
  private static readonly CACHE_MAX_ENTRIES = 5_000;
  private readonly cache = new Map<
    string,
    { employees: SnapshotEmployee[]; expiresAt: number }
  >();

  constructor(
    reflector: Reflector,
    @InjectDataSource() private readonly ds: DataSource,
  ) {
    super(reflector, ds);
  }

  async canActivate(context: ExecutionContext): Promise<boolean> {
    if (!(await super.canActivate(context))) return false;

    const request = context.switchToHttp().getRequest();
    const user = request.user as TCurrentUser | undefined;
    const employee = user?.employee as SnapshotEmployee | undefined;
    if (!employee || !user?.sessionId) return true;

    const employees = await this.employeesForSession(user.sessionId);
    if (!employees.length) return true;

    const requestedId = request.headers?.[CURRENT_POSITION_ID] as
      | string
      | undefined;

    const { owner, active } = resolveActiveEmployee(
      employees,
      requestedId,
      employee.id,
    );

    const ownerPositions = owner?.positions ?? [];
    // Never blank out what the parent resolved: a snapshot without positions
    // must degrade to the single-position behaviour, not to no positions.
    if (!ownerPositions.length) return true;

    // Carries the owning row's id / unitId / organizationId too, which unit
    // scoping downstream reads — a swapped row must be swapped whole.
    Object.assign(employee, owner);

    // `collectPermissionKeys` / `collectPositionTypeKeys` union over this, and
    // a user's posts can span several employee rows (one per organization), so
    // it carries every row's — otherwise a freight post is invisible whenever
    // another organization's row wins the active slot.
    employee.positions = collectAllPositions(employees);

    // Delegation stays scoped to the active desk: yard scope widens on
    // `delegatedPositions`, and someone standing in on another organization's
    // row is not this desk's stand-in.
    employee.delegatedPositions = ownerPositions.filter(
      (position) => position.isDelegate,
    );

    // The full set, for `/auth/me` — the position picker has to be able to
    // offer a desk on a row that is not the active one.
    (user as { employeeRows?: SnapshotEmployee[] }).employeeRows = employees;

    if (active) {
      employee.position = active;
      // The parent already built `auditUser` from the position it guessed.
      if (request.auditUser) {
        request.auditUser.employeeId = employee.id;
        request.auditUser.positionId = active.id;
        request.auditUser.employeePositionId = active.employeePositionId;
      }
    }

    return true;
  }

  /** Every employee row the login snapshot holds for this session. */
  private async employeesForSession(
    sessionId: string,
  ): Promise<SnapshotEmployee[]> {
    const now = Date.now();
    const hit = this.cache.get(sessionId);
    if (hit && hit.expiresAt > now) return hit.employees;

    let employees: SnapshotEmployee[] = [];
    try {
      const rows: { userInfo: SessionUserInfo | null }[] = await this.ds.query(
        `SELECT "userInfo" FROM iam.sessions WHERE id = $1`,
        [sessionId],
      );
      employees = rows[0]?.userInfo?.employee ?? [];
    } catch {
      return []; // iam unreachable — caller keeps the parent's single position
    }

    if (this.cache.size >= FreightJwtGuard.CACHE_MAX_ENTRIES)
      this.cache.clear();
    this.cache.set(sessionId, {
      employees,
      expiresAt: now + FreightJwtGuard.CACHE_TTL_MS,
    });
    return employees;
  }
}