kv-music/vite-plugin-auth-gate.js

import { loadEnv } from 'vite';
import cookieSession from 'cookie-session';
import { createRemoteJWKSet, jwtVerify } from 'jose';
import { readFileSync, existsSync } from 'fs';
import { join, extname } from 'path';

function parseBody(req) {
    return new Promise((resolve, reject) => {
        let body = '';
        req.on('data', (chunk) => (body += chunk));
        req.on('end', () => {
            try {
                resolve(JSON.parse(body));
            } catch {
                reject(new Error('Invalid JSON'));
            }
        });
        req.on('error', reject);
    });
}

function buildInjectionScript(env) {
    const AUTH_ENABLED = (env.AUTH_ENABLED ?? 'false') !== 'false';
    const APPWRITE_ENDPOINT = env.APPWRITE_ENDPOINT;
    const APPWRITE_PROJECT_ID = env.APPWRITE_PROJECT_ID;
    const POCKETBASE_URL = env.POCKETBASE_URL;
    const AUTH_GOOGLE_ENABLED = env.AUTH_GOOGLE_ENABLED;
    const AUTH_EMAIL_ENABLED = env.AUTH_EMAIL_ENABLED;

    const flags = [];
    if (AUTH_ENABLED) flags.push('window.__AUTH_GATE__=true');
    const authProviderOverrides = {};
    if (AUTH_GOOGLE_ENABLED !== undefined) {
        authProviderOverrides.google = AUTH_GOOGLE_ENABLED !== 'false';
    }
    if (AUTH_EMAIL_ENABLED !== undefined) {
        authProviderOverrides.password = AUTH_EMAIL_ENABLED !== 'false';
    }
    if (Object.keys(authProviderOverrides).length > 0) {
        flags.push(`window.__AUTH_PROVIDERS__=${JSON.stringify(authProviderOverrides)}`);
    }
    if (APPWRITE_ENDPOINT) flags.push(`window.__APPWRITE_ENDPOINT__=${JSON.stringify(APPWRITE_ENDPOINT)}`);
    if (APPWRITE_PROJECT_ID) flags.push(`window.__APPWRITE_PROJECT_ID__=${JSON.stringify(APPWRITE_PROJECT_ID)}`);
    if (POCKETBASE_URL) flags.push(`window.__POCKETBASE_URL__=${JSON.stringify(POCKETBASE_URL)}`);

    return flags.length > 0 ? `<script>${flags.join(';')};</script>` : null;
}

export default function authGatePlugin() {
    let env = {};

    return {
        name: 'auth-gate',

        config(_, { mode }) {
            env = loadEnv(mode, process.cwd(), '');
        },

        transformIndexHtml(html) {
            const scriptTag = buildInjectionScript(env);
            return scriptTag ? html.replace('</head>', `${scriptTag}\n</head>`) : html;
        },

        configurePreviewServer(server) {
            const configScript = buildInjectionScript(env);

            // --- Pre-build injected HTML pages ---

            const distDir = join(process.cwd(), 'dist');

            let indexHtml = null;
            const indexPath = join(distDir, 'index.html');
            if (configScript && existsSync(indexPath)) {
                indexHtml = readFileSync(indexPath, 'utf-8');
                indexHtml = indexHtml.replace('</head>', `${configScript}\n</head>`);
            }

            let loginHtml = null;
            const AUTH_ENABLED = (env.AUTH_ENABLED ?? 'false') !== 'false';
            if (AUTH_ENABLED) {
                const loginPath = join(distDir, 'login.html');
                if (existsSync(loginPath)) {
                    loginHtml = readFileSync(loginPath, 'utf-8');
                    if (configScript) loginHtml = loginHtml.replace('</head>', `${configScript}\n</head>`);
                }
            }

            // --- /health (always available) ---

            server.middlewares.use((req, res, next) => {
                if (req.url.split('?')[0] === '/health') {
                    res.end('OK');
                    return;
                }
                next();
            });

            // --- Auth gate (only when enabled) ---

            if (AUTH_ENABLED) {
                const AUTH_SECRET = env.AUTH_SECRET;
                const SESSION_MAX_AGE = Number(env.SESSION_MAX_AGE) || 7 * 24 * 60 * 60 * 1000;

                if (!AUTH_SECRET) {
                    console.error('AUTH_SECRET is required when AUTH_ENABLED=true');
                    process.exit(1);
                }

                console.log(`Auth gate enabled (Project: ${env.APPWRITE_PROJECT_ID})`);

                server.middlewares.use(
                    cookieSession({
                        name: 'mono_session',
                        keys: [AUTH_SECRET],
                        maxAge: SESSION_MAX_AGE,
                        httpOnly: true,
                        sameSite: 'lax',
                    })
                );

                server.middlewares.use(async (req, res, next) => {
                    const url = req.url.split('?')[0];

                    if (url === '/login' || url === '/login.html') {
                        if (req.session && req.session.uid) {
                            res.writeHead(302, { Location: '/' });
                            res.end();
                            return;
                        }
                        if (loginHtml) {
                            res.setHeader('Content-Type', 'text/html');
                            res.setHeader('Cache-Control', 'no-store');
                            res.end(loginHtml);
                        } else {
                            res.statusCode = 404;
                            res.end('Login page not found');
                        }
                        return;
                    }

                    if (url === '/api/auth/login' && req.method === 'POST') {
                        try {
                            const body = await parseBody(req);
                            if (!body.userId) {
                                res.statusCode = 400;
                                res.setHeader('Content-Type', 'application/json');
                                res.end(JSON.stringify({ error: 'Missing userId' }));
                                return;
                            }
                            req.session.uid = body.userId;
                            req.session.email = body.email;
                            req.session.iat = Date.now();
                            res.setHeader('Content-Type', 'application/json');
                            res.end(JSON.stringify({ ok: true }));
                        } catch (err) {
                            console.error('Login session creation failed:', err.message);
                            res.statusCode = 401;
                            res.setHeader('Content-Type', 'application/json');
                            res.end(JSON.stringify({ error: 'Login failed' }));
                        }
                        return;
                    }

                    if (url === '/api/auth/logout' && req.method === 'POST') {
                        req.session = null;
                        res.setHeader('Clear-Site-Data', '"cache", "storage"');
                        res.setHeader('Content-Type', 'application/json');
                        res.end(JSON.stringify({ ok: true }));
                        return;
                    }

                    if (!req.session || !req.session.uid) {
                        const ext = extname(url);
                        if (ext && ext !== '.html') {
                            res.statusCode = 401;
                            res.end('Unauthorized');
                        } else {
                            res.writeHead(302, { Location: '/login' });
                            res.end();
                        }
                        return;
                    }

                    // Authenticated: serve injected index.html for HTML routes
                    const ext = extname(url);
                    if ((!ext || ext === '.html') && indexHtml) {
                        res.setHeader('Content-Type', 'text/html');
                        res.setHeader('Cache-Control', 'no-store');
                        res.end(indexHtml);
                        return;
                    }

                    next();
                });
            } else if (indexHtml) {
                // No auth gate, but env config needs injection into HTML
                server.middlewares.use((req, res, next) => {
                    const url = req.url.split('?')[0];
                    const ext = extname(url);
                    if (!ext || ext === '.html') {
                        res.setHeader('Content-Type', 'text/html');
                        res.end(indexHtml);
                        return;
                    }
                    next();
                });
            }
        },
    };
}