Refresh Tokens and Secure Token Rotation

// Helper functions for token generation
const generateAccessToken = (userId: number) => {
  return jwt.sign({ userId }, JWT_SECRET_KEY, { expiresIn: '15m' });  // Short-lived token (15 minutes)
};

const generateRefreshToken = (userId: number) => {
  // Generate a unique jti using Math.random
  return jwt.sign(
    { userId, jti: Math.random().toString(36).substring(2) }, // Unique JWT ID
    JWT_SECRET_KEY,
    { expiresIn: '7d' }
  );
};

// Generate both access and refresh tokens during login
router.post('/login', async (req, res) => {
    const { username, password } = req.body;
    try {
        const user = await User.findOne({ where: { username } });
        if (user && await bcrypt.compare(password, user.password)) {
            // Generate tokens using helper functions
            const accessToken = generateAccessToken(user.id);
            const refreshToken = generateRefreshToken(user.id);
            
            // Store tokens in HTTP-only cookies
            res.cookie('token', accessToken, {
                httpOnly: true,
                secure: process.env.NODE_ENV === 'production',
                sameSite: 'strict',
                maxAge: 15 * 60 * 1000 // 15 minutes
            });
            
            res.cookie('refreshToken', refreshToken, {
                httpOnly: true,
                secure: process.env.NODE_ENV === 'production',
                sameSite: 'strict',
                maxAge: 7 * 24 * 60 * 60 * 1000 // 7 days
            });
            
            return res.json({ message: 'Login successful' });
        }
        res.status(401).json({ error: 'Invalid credentials' });
    } catch (error) {
        console.error('Error during login:', error);
        res.status(500).json({ error: 'Internal server error' });
    }
});

// Helper functions for token generation
const generateAccessToken = (userId: number) => {
  return jwt.sign({ userId }, JWT_SECRET_KEY, { expiresIn: '15m' });  // Short-lived token (15 minutes)
};

// Refresh token endpoint
router.post('/refresh', async (req, res) => {
    const refreshToken = req.cookies.refreshToken;
    
    if (!refreshToken) {
        return res.status(401).json({ error: 'Refresh token not provided' });
    }
    
    try {
        // Verify the refresh token
        const decoded = jwt.verify(refreshToken, JWT_SECRET_KEY);
        
        // Find the user
        const user = await User.findByPk(decoded.userId);
        if (!user) {
            return res.status(401).json({ error: 'User not found' });
        }
        
        // Generate a new access token using the helper function
        const newAccessToken = generateAccessToken(user.id);
        
        // Set the new access token as a cookie
        res.cookie('token', newAccessToken, {
            httpOnly: true,
            secure: process.env.NODE_ENV === 'production',
            sameSite: 'strict',
            maxAge: 15 * 60 * 1000 // 15 minutes
        });
        
        return res.json({ message: 'Token refreshed successfully' });
    } catch (error) {
        return res.status(401).json({ error: 'Invalid refresh token' });
    }
});

// In-memory store for refresh tokens: userId -> Set of valid jtis
const refreshTokenStore = new Map();

// Helper functions for token generation
const generateAccessToken = (userId: number) => {
  return jwt.sign({ userId }, JWT_SECRET_KEY, { expiresIn: '15m' });
};

const generateRefreshToken = (userId: number) => {
  // Generate a unique jti using Math.random
  return jwt.sign(
    { userId, jti: Math.random().toString(36).substring(2) },
    JWT_SECRET_KEY,
    { expiresIn: '7d' }
  );
};

// Modified login endpoint to store refresh token jti
router.post('/login', async (req, res) => {
    const { username, password } = req.body;
    try {
        const user = await User.findOne({ where: { username } });
        if (user && await bcrypt.compare(password, user.password)) {
            const accessToken = generateAccessToken(user.id);
            const refreshToken = generateRefreshToken(user.id);

            // Decode the refresh token to extract jti
            const { jti } = jwt.decode(refreshToken);

            // Store the refresh token's jti
            let userJtis = refreshTokenStore.get(user.id) || new Set();
            userJtis.add(jti);
            refreshTokenStore.set(user.id, userJtis);

            // Set cookies
            res.cookie('token', accessToken, {
                httpOnly: true,
                secure: process.env.NODE_ENV === 'production',
                sameSite: 'strict',
                maxAge: 15 * 60 * 1000
            });
            res.cookie('refreshToken', refreshToken, {
                httpOnly: true,
                secure: process.env.NODE_ENV === 'production',
                sameSite: 'strict',
                maxAge: 7 * 24 * 60 * 60 * 1000
            });

            return res.json({ message: 'Login successful' });
        }
        res.status(401).json({ error: 'Invalid credentials' });
    } catch (error) {
        console.error('Error during login:', error);
        res.status(500).json({ error: 'Internal server error' });
    }
});

// Modified refresh endpoint with token rotation using jti
router.post('/refresh', async (req, res) => {
    const refreshToken = req.cookies.refreshToken;

    if (!refreshToken) {
        return res.status(401).json({ error: 'Refresh token not provided' });
    }

    try {
        // Verify the refresh token
        const decoded = jwt.verify(refreshToken, JWT_SECRET_KEY);
        const userId = decoded.userId;
        const jti = decoded.jti;

        // Check if the jti exists in our store
        let userJtis = refreshTokenStore.get(userId) || new Set();
        if (!userJtis.has(jti)) {
            return res.status(401).json({ error: 'Invalid refresh token' });
        }
        const user = await User.findByPk(userId);
        if (!user) {
            return res.status(401).json({ error: 'User not found' });
        }
        // Remove the used jti (token rotation)
        userJtis.delete(jti);

        // Generate new tokens
        const newAccessToken = generateAccessToken(user.id);
        const newRefreshToken = generateRefreshToken(user.id);

        // Store the new refresh token's jti
        const { jti: newJti } = jwt.decode(newRefreshToken);
        userJtis.add(newJti);
        refreshTokenStore.set(userId, userJtis);

        // Set the new tokens as cookies
        res.cookie('token', newAccessToken, {
            httpOnly: true,
            secure: process.env.NODE_ENV === 'production',
            sameSite: 'strict',
            maxAge: 15 * 60 * 1000
        });
        res.cookie('refreshToken', newRefreshToken, {
            httpOnly: true,
            secure: process.env.NODE_ENV === 'production',
            sameSite: 'strict',
            maxAge: 7 * 24 * 60 * 60 * 1000
        });

        return res.json({ message: 'Tokens refreshed successfully' });
    } catch (error) {
        return res.status(401).json({ error: 'Invalid refresh token' });
    }
});