dd0c/products/05-aws-cost-anomaly/src/auth/middleware.ts

import { z } from 'zod';
import jwt from 'jsonwebtoken';
import crypto from 'node:crypto';
import type { FastifyInstance, FastifyRequest, FastifyReply } from 'fastify';
import type { Pool } from 'pg';
import pino from 'pino';

const logger = pino({ name: 'auth' });

export interface AuthPayload {
  tenantId: string;
  userId: string;
  email: string;
  role: 'owner' | 'admin' | 'member' | 'viewer';
}

/**
 * Returns an onRequest hook that validates JWT or API key auth.
 * No URL matching — only register this hook inside a protected plugin scope.
 */
export function authHook(jwtSecret: string, pool: Pool) {
  return async (req: FastifyRequest, reply: FastifyReply) => {
    const apiKey = req.headers['x-api-key'] as string | undefined;
    const authHeader = req.headers['authorization'];

    if (apiKey) {
      if (!apiKey.startsWith('dd0c_') || apiKey.length !== 37) {
        return reply.status(401).send({ error: 'Invalid API key format' });
      }

      const prefix = apiKey.slice(0, 13);
      const keyHash = crypto.createHash('sha256').update(apiKey).digest('hex');

      const result = await pool.query(
        `SELECT ak.tenant_id, ak.user_id, u.role
         FROM api_keys ak JOIN users u ON ak.user_id = u.id
         WHERE ak.key_prefix = $1 AND ak.key_hash = $2 AND ak.revoked = false`,
        [prefix, keyHash],
      );

      if (!result.rows[0]) {
        return reply.status(401).send({ error: 'Invalid API key' });
      }

      (req as any).tenantId = result.rows[0].tenant_id;
      (req as any).userId = result.rows[0].user_id;
      (req as any).userRole = result.rows[0].role;
      return;
    }

    if (authHeader?.startsWith('Bearer ')) {
      const token = authHeader.slice(7);
      try {
        const payload = jwt.verify(token, jwtSecret, { algorithms: ['HS256'] }) as AuthPayload;
        (req as any).tenantId = payload.tenantId;
        (req as any).userId = payload.userId;
        (req as any).userRole = payload.role;
        return;
      } catch {
        return reply.status(401).send({ error: 'Invalid or expired token' });
      }
    }

    return reply.status(401).send({ error: 'Missing authentication' });
  };
}

/**
 * Decorate the Fastify request with auth properties.
 * Call this once on the root app instance before registering any routes.
 */
export function decorateAuth(app: FastifyInstance) {
  app.decorateRequest('tenantId', '');
  app.decorateRequest('userId', '');
  app.decorateRequest('userRole', 'viewer');
}

export function requireRole(req: FastifyRequest, reply: FastifyReply, minRole: AuthPayload['role']): boolean {
  const hierarchy: Record<string, number> = { owner: 4, admin: 3, member: 2, viewer: 1 };
  const userLevel = hierarchy[(req as any).userRole] ?? 0;
  const requiredLevel = hierarchy[minRole] ?? 0;

  if (userLevel < requiredLevel) {
    reply.status(403).send({ error: `Requires ${minRole} role or higher` });
    return false;
  }
  return true;
}

export function signToken(payload: AuthPayload, secret: string, expiresIn = '24h'): string {
  return jwt.sign(payload, secret, { expiresIn } as jwt.SignOptions);
}

// --- Password hashing (scrypt — no native bcrypt dep needed) ---

async function hashPassword(password: string): Promise<string> {
  const salt = crypto.randomBytes(16).toString('hex');
  return new Promise((resolve, reject) => {
    crypto.scrypt(password, salt, 64, { N: 65536, r: 8, p: 1, maxmem: 128 * 1024 * 1024 }, (err, derived) => {
      if (err) reject(err);
      resolve(`${salt}:${derived.toString('hex')}`);
    });
  });
}

async function verifyPassword(password: string, hash: string): Promise<boolean> {
  const [salt, key] = hash.split(':');
  return new Promise((resolve, reject) => {
    crypto.scrypt(password, salt, 64, { N: 65536, r: 8, p: 1, maxmem: 128 * 1024 * 1024 }, (err, derived) => {
      if (err) reject(err);
      resolve(crypto.timingSafeEqual(Buffer.from(key, 'hex'), derived));
    });
  });
}

// --- Auth Routes ---

const loginSchema = z.object({
  email: z.string().email(),
  password: z.string().min(8),
});

const signupSchema = z.object({
  email: z.string().email(),
  password: z.string().min(8),
  tenant_name: z.string().min(1).max(100).optional(),
  invite_token: z.string().optional(),
}).refine(
  (data) => data.invite_token || data.tenant_name,
  { message: 'Either tenant_name or invite_token is required', path: ['tenant_name'] },
);

const inviteSchema = z.object({
  email: z.string().email(),
  role: z.enum(['admin', 'member', 'viewer']).default('member'),
});

/** Public auth routes — login/signup. No auth required. */
export function registerAuthRoutes(app: FastifyInstance, jwtSecret: string, pool: Pool) {
  app.post('/api/v1/auth/login', { config: { rateLimit: { max: 5, timeWindow: '1 minute' } } }, async (req, reply) => {
    const body = loginSchema.parse(req.body);

    const result = await pool.query(
      `SELECT u.id, u.tenant_id, u.email, u.password_hash, u.role
       FROM users u WHERE u.email = $1`,
      [body.email],
    );

    const user = result.rows[0];
    if (!user) return reply.status(401).send({ error: 'Invalid credentials' });

    const valid = await verifyPassword(body.password, user.password_hash);
    if (!valid) return reply.status(401).send({ error: 'Invalid credentials' });

    const token = signToken({
      tenantId: user.tenant_id,
      userId: user.id,
      email: user.email,
      role: user.role,
    }, jwtSecret);

    return { token, expires_in: '24h' };
  });

  app.post('/api/v1/auth/signup', { config: { rateLimit: { max: 3, timeWindow: '1 minute' } } }, async (req, reply) => {
    const body = signupSchema.parse(req.body);

    const existing = await pool.query('SELECT id FROM users WHERE email = $1', [body.email]);
    if (existing.rows[0]) return reply.status(409).send({ error: 'Email already registered' });

    const passwordHash = await hashPassword(body.password);

    const client = await pool.connect();
    try {
      await client.query('BEGIN');

      let tenantId: string;
      let role: string;

      if (body.invite_token) {
        const tokenHash = crypto.createHash('sha256').update(body.invite_token).digest('hex');
        const invite = await client.query(
          `SELECT id, tenant_id, email, role, expires_at, accepted_at FROM tenant_invites WHERE token = $1 FOR UPDATE`,
          [tokenHash],
        );
        if (!invite.rows[0]) {
          await client.query('ROLLBACK');
          return reply.status(400).send({ error: 'Invalid invite token' });
        }
        const inv = invite.rows[0];
        if (inv.accepted_at) {
          await client.query('ROLLBACK');
          return reply.status(400).send({ error: 'Invite already accepted' });
        }
        if (new Date(inv.expires_at) < new Date()) {
          await client.query('ROLLBACK');
          return reply.status(400).send({ error: 'Invite expired' });
        }
        if (inv.email && inv.email.toLowerCase() !== body.email.toLowerCase()) {
          await client.query('ROLLBACK');
          return reply.status(400).send({ error: 'Email does not match invite' });
        }

        tenantId = inv.tenant_id;
        role = inv.role;

        await client.query(
          `UPDATE tenant_invites SET accepted_at = NOW() WHERE id = $1`,
          [inv.id],
        );
      } else {
        if (!body.tenant_name) {
          await client.query('ROLLBACK');
          return reply.status(400).send({ error: 'tenant_name is required for new signups' });
        }
        const slug = body.tenant_name.toLowerCase().replace(/[^a-z0-9]+/g, '-').slice(0, 42) + '-' + crypto.randomBytes(3).toString('hex');
        const tenant = await client.query(
          `INSERT INTO tenants (name, slug) VALUES ($1, $2) RETURNING id`,
          [body.tenant_name, slug],
        );
        tenantId = tenant.rows[0].id;
        role = 'owner';
      }

      // Set tenant context for RLS before inserting user
      await client.query("SELECT set_config('app.tenant_id', $1, true)", [tenantId]);

      const user = await client.query(
        `INSERT INTO users (tenant_id, email, password_hash, role) VALUES ($1, $2, $3, $4) RETURNING id`,
        [tenantId, body.email, passwordHash, role],
      );

      await client.query('COMMIT');

      const token = signToken({
        tenantId,
        userId: user.rows[0].id,
        email: body.email,
        role: role as AuthPayload['role'],
      }, jwtSecret);

      return reply.status(201).send({ token, tenant_id: tenantId, expires_in: '24h' });
    } catch (err) {
      await client.query('ROLLBACK');
      throw err;
    } finally {
      client.release();
    }
  });
}

/** Protected auth routes — me, api-keys, invites. Must be registered inside an auth-protected plugin scope. */
export function registerProtectedAuthRoutes(app: FastifyInstance, jwtSecret: string, pool: Pool) {
  app.get('/api/v1/auth/me', async (req, reply) => {
    return {
      tenant_id: (req as any).tenantId,
      user_id: (req as any).userId,
      role: (req as any).userRole,
    };
  });

  app.post('/api/v1/auth/api-keys', async (req, reply) => {
    const tenantId = (req as any).tenantId;
    const userId = (req as any).userId;

    if (!requireRole(req, reply, 'admin')) return;

    const rawKey = `dd0c_${crypto.randomBytes(16).toString('hex')}`;
    const prefix = rawKey.slice(0, 13);
    const keyHash = crypto.createHash('sha256').update(rawKey).digest('hex');

    await pool.query(
      `INSERT INTO api_keys (tenant_id, user_id, key_prefix, key_hash) VALUES ($1, $2, $3, $4)`,
      [tenantId, userId, prefix, keyHash],
    );

    return reply.status(201).send({ api_key: rawKey, prefix });
  });

  // --- Invite endpoints ---

  app.post('/api/v1/auth/invite', async (req, reply) => {
    if (!requireRole(req, reply, 'admin')) return;
    const tenantId = (req as any).tenantId;
    const userId = (req as any).userId;
    const body = inviteSchema.parse(req.body);

    const token = crypto.randomBytes(32).toString('hex');
    const inviteTokenHash = crypto.createHash('sha256').update(token).digest('hex');
    const result = await pool.query(
      `INSERT INTO tenant_invites (tenant_id, email, role, token, invited_by)
       VALUES ($1, $2, $3, $4, $5)
       RETURNING expires_at`,
      [tenantId, body.email, body.role, inviteTokenHash, userId],
    );

    return reply.status(201).send({ invite_token: token, expires_at: result.rows[0].expires_at });
  });

  app.get('/api/v1/auth/invites', async (req, reply) => {
    if (!requireRole(req, reply, 'admin')) return;
    const tenantId = (req as any).tenantId;

    const result = await pool.query(
      `SELECT id, email, role, expires_at, created_at FROM tenant_invites
       WHERE tenant_id = $1 AND accepted_at IS NULL AND expires_at > NOW()
       ORDER BY created_at DESC`,
      [tenantId],
    );

    return { invites: result.rows };
  });

  app.delete('/api/v1/auth/invites/:id', async (req, reply) => {
    if (!requireRole(req, reply, 'admin')) return;
    const tenantId = (req as any).tenantId;
    const { id } = req.params as { id: string };

    const result = await pool.query(
      `DELETE FROM tenant_invites WHERE id = $1 AND tenant_id = $2 AND accepted_at IS NULL RETURNING id`,
      [id, tenantId],
    );

    if (!result.rows[0]) return reply.status(404).send({ error: 'Invite not found' });
    return { deleted: true };
  });
}