DiscordClone/Backend/routes/auth.js

const express = require('express');
const router = express.Router();
const db = require('../db');
const crypto = require('crypto');

// Helper to generate fake salt for user privacy
function generateFakeSalt(username) {
    return crypto.createHmac('sha256', 'SERVER_SECRET_KEY') // In prod, use env var
        .update(username)
        .digest('hex');
}

router.post('/register', async (req, res) => {
    const { username, salt, encryptedMK, hak, publicKey, signingKey, encryptedPrivateKeys, inviteCode } = req.body;
    
    try {
        // Step 1: Enforce Invite (unless first user)
        const userCountRes = await db.query('SELECT count(*) FROM users');
        const userCount = parseInt(userCountRes.rows[0].count);

        if (userCount > 0) {
            if (!inviteCode) {
                return res.status(403).json({ error: 'Invite code required' });
            }

            // Check Invite validity
            const inviteRes = await db.query('SELECT * FROM invites WHERE code = $1', [inviteCode]);
            if (inviteRes.rows.length === 0) {
                return res.status(403).json({ error: 'Invalid invite code' });
            }
            
            var invite = inviteRes.rows[0];

            // Check Expiration
            if (invite.expires_at && new Date() > new Date(invite.expires_at)) {
                return res.status(410).json({ error: 'Invite expired' });
            }

            // Check Usage Limits
            if (invite.max_uses !== null && invite.uses >= invite.max_uses) {
                return res.status(410).json({ error: 'Invite max uses reached' });
            }
        }

        // START TRANSACTION - To ensure invite usage and user creation are atomic
        await db.query('BEGIN');

        try {
            // Update Invite Usage (only if enforced)
            if (userCount > 0) {
                await db.query('UPDATE invites SET uses = uses + 1 WHERE code = $1', [inviteCode]);
            }

            // Create User
            // Create User
            const result = await db.query(
                `INSERT INTO users (username, client_salt, encrypted_master_key, hashed_auth_key, public_identity_key, public_signing_key, encrypted_private_keys)
                 VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING id`,
                [username, salt, encryptedMK, hak, publicKey, signingKey, encryptedPrivateKeys]
            );
            const newUserId = result.rows[0].id;

            // Assign Roles
            // 1. @everyone (Always)
            await db.query(`
                INSERT INTO user_roles (user_id, role_id)
                SELECT $1, id FROM roles WHERE name = '@everyone'
            `, [newUserId]);

            // 2. Owner (If first user or if admin logic allows)
            if (userCount === 0) {
                 await db.query(`
                    INSERT INTO user_roles (user_id, role_id)
                    SELECT $1, id FROM roles WHERE name = 'Owner'
                `, [newUserId]);
                
                // Also set is_admin = true for legacy support
                await db.query('UPDATE users SET is_admin = TRUE WHERE id = $1', [newUserId]);
            }

            await db.query('COMMIT');
            res.json({ success: true, userId: result.rows[0].id });

        } catch (txErr) {
            await db.query('ROLLBACK');
            throw txErr;
        }

    } catch (err) {
        console.error(err);
        if (err.code === '23505') { // Unique violation
            res.status(400).json({ error: 'Username taken' });
        } else {
            res.status(500).json({ error: 'Server error' });
        }
    }
});

router.post('/login/salt', async (req, res) => {
    const { username } = req.body;
    try {
        const result = await db.query('SELECT client_salt FROM users WHERE username = $1', [username]);
        if (result.rows.length > 0) {
            res.json({ salt: result.rows[0].client_salt });
        } else {
            // Return fake salt to prevent enumeration
            res.json({ salt: generateFakeSalt(username) });
        }
    } catch (err) {
        console.error(err);
        res.status(500).json({ error: 'Server error' });
    }
});

router.post('/login/verify', async (req, res) => {
    const { username, dak } = req.body;

    try {
        const result = await db.query(
            'SELECT id, hashed_auth_key, encrypted_master_key, encrypted_private_keys, public_identity_key FROM users WHERE username = $1',
            [username]
        );

        if (result.rows.length === 0) {
            return res.status(401).json({ error: 'Invalid credentials' });
        }

        const user = result.rows[0];
        const hashedDAK = crypto.createHash('sha256').update(dak).digest('hex');

        if (hashedDAK === user.hashed_auth_key) {
            res.json({
                success: true,
                userId: user.id,
                encryptedMK: user.encrypted_master_key,
                encryptedPrivateKeys: user.encrypted_private_keys,
                publicKey: user.public_identity_key // Return Public Key
            });
        } else {
            res.status(401).json({ error: 'Invalid credentials' });
        }
    } catch (err) {
        console.error(err);
        res.status(500).json({ error: 'Server error' });
    }
});

router.get('/users/public-keys', async (req, res) => {
    try {
        const result = await db.query('SELECT id, username, public_identity_key FROM users');
        res.json(result.rows);
    } catch (err) {
        console.error(err);
        res.status(500).json({ error: 'Server error' });
    }
});

module.exports = router;