Update server.js

server.js

@@ -7,7 +7,6 @@ require('dotenv').config();
 
 const app = express();
 
-// --- SECURITY HEADERS ---
 app.use(helmet({
     contentSecurityPolicy: false,
     crossOriginEmbedderPolicy: false
@@ -16,14 +15,48 @@ app.use(helmet({
 app.set('trust proxy', 1);
 app.use(express.json({ limit: '1mb' }));
 
-// --- CORS PROTECTION ---
+// --- API KEY AUTHENTICATION ---
+const API_KEYS = process.env.API_KEYS 
+    ? process.env.API_KEYS.split(',').map(k => k.trim())
+    : [];
+
+function authenticateRequest(req) {
+    const origin = req.headers.origin;
+    const apiKey = req.headers['x-api-key'];
+    
+    // CASE 1: Browser request (has Origin)
+    if (origin) {
+        if (allowedOrigins.length === 0) return { valid: true, source: 'open-mode' };
+        return { 
+            valid: allowedOrigins.includes(origin), 
+            source: origin 
+        };
+    }
+    
+    // CASE 2: Backend request (no Origin) - MUST have API key
+    if (apiKey) {
+        if (API_KEYS.length === 0) return { valid: true, source: 'no-keys-configured' };
+        return { 
+            valid: API_KEYS.includes(apiKey), 
+            source: 'api-key' 
+        };
+    }
+    
+    // CASE 3: No Origin, No API Key - BLOCKED
+    return { valid: false, source: 'unauthorized' };
+}
+
+// --- CORS (For browsers) ---
 const allowedOrigins = process.env.ALLOWED_ORIGINS 
     ? process.env.ALLOWED_ORIGINS.split(',').map(o => o.trim())
     : [];
 
 app.use(cors({
     origin: function (origin, callback) {
-        if (!origin || allowedOrigins.length === 0) return callback(null, true);
+        if (!origin) {
+            return callback(null, true);
+        }
+        if (allowedOrigins.length === 0) return callback(null, true);
         if (allowedOrigins.includes(origin)) {
             callback(null, true);
         } else {
@@ -41,7 +74,27 @@ app.use((err, req, res, next) => {
     next(err);
 });
 
-// --- RATE LIMITING (15 minutes) ---
+// --- API KEY AUTHENTICATION MIDDLEWARE ---
+app.use((req, res, next) => {
+    if (req.path === '/' || req.path === '/health') {
+        return next();
+    }
+    
+    const auth = authenticateRequest(req);
+    
+    if (!auth.valid) {
+        console.log(`[Security] Blocked unauthorized request from ${req.ip} (${auth.source})`);
+        return res.status(401).json({ 
+            error: 'Unauthorized',
+            message: 'Valid origin or API key required' 
+        });
+    }
+    
+    console.log(`[Auth] Request authorized from: ${auth.source}`);
+    next();
+});
+
+// --- RATE LIMITING ---
 const limiter = rateLimit({
     windowMs: 15 * 60 * 1000,
     max: 100,
@@ -56,8 +109,9 @@ app.use(limiter);
 // --- REQUEST LOGGING ---
 app.use((req, res, next) => {
     const ip = (req.ip || 'unknown').replace(/:\d+[^:]*$/, '');
-    const origin = (req.headers.origin || 'unknown').substring(0, 100);
-    console.log(`[${new Date().toISOString()}] ${ip} -> ${req.method} ${req.path} from ${origin}`);
+    const origin = req.headers.origin || 'no-origin';
+    const apiKey = req.headers['x-api-key'] ? '***' : 'none';
+    console.log(`[${new Date().toISOString()}] ${ip} -> ${req.method} ${req.path} | Origin: ${origin} | Key: ${apiKey}`);
     next();
 });
 
@@ -92,7 +146,6 @@ app.use((req, res, next) => {
         
         dailyUsage.set(ip, count + 1);
         
-        // FIXED: Prevent unbounded growth
         if (dailyUsage.size > 10000) {
             console.warn('[System] Daily usage map too large, clearing oldest entries');
             const entries = Array.from(dailyUsage.entries()).slice(0, 1000);
@@ -104,7 +157,6 @@ app.use((req, res, next) => {
 
 // --- BOT DETECTION ---
 app.use((req, res, next) => {
-    // FIXED: Check all POST requests, not just /prediction/
     if (req.method !== 'POST') {
         return next();
     }
@@ -112,6 +164,11 @@ app.use((req, res, next) => {
     const userAgent = (req.headers['user-agent'] || '').toLowerCase();
     const suspiciousBots = ['python-requests', 'curl/', 'wget/', 'scrapy', 'crawler'];
     
+    const hasValidApiKey = req.headers['x-api-key'] && API_KEYS.includes(req.headers['x-api-key']);
+    if (hasValidApiKey) {
+        return next();
+    }
+    
     const isBot = suspiciousBots.some(bot => userAgent.includes(bot));
     
     if (isBot) {
@@ -142,7 +199,6 @@ const flowCache = new Map();
 setInterval(() => {
     const now = Date.now();
     for (const [key, value] of flowCache.entries()) {
-        // FIXED: Safety check for timestamp
         if (value.timestamp && now - value.timestamp > 10 * 60 * 1000) {
             flowCache.delete(key);
         }
@@ -232,7 +288,6 @@ async function handleStreamingResponse(flowiseResponse, clientRes) {
     flowiseResponse.body.on('error', (err) => {
         console.error('[Streaming Error]', err.message);
         
-        // FIXED: Send error event if stream already started
         if (streamStarted) {
             clientRes.write(`\n\nevent: error\ndata: {"error": "Stream interrupted"}\n\n`);
         }
@@ -240,7 +295,7 @@ async function handleStreamingResponse(flowiseResponse, clientRes) {
     });
 }
 
-// --- ROUTE 1: PREDICTION (WITH STREAMING SUPPORT) ---
+// --- ROUTES ---
 app.post('/api/v1/prediction/:instanceNum/:botName', async (req, res) => {
     try {
         const instanceNum = parseInt(req.params.instanceNum);
@@ -288,13 +343,11 @@ app.post('/api/v1/prediction/:instanceNum/:botName', async (req, res) => {
         
         const contentType = response.headers.get('content-type') || '';
         
-        // STREAMING RESPONSE
         if (contentType.includes('text/event-stream')) {
             console.log('[Streaming] Detected SSE response');
             return handleStreamingResponse(response, res);
         }
         
-        // NON-STREAMING RESPONSE
         console.log('[Non-streaming] Parsing JSON response');
         const text = await response.text();
         
@@ -315,7 +368,6 @@ app.post('/api/v1/prediction/:instanceNum/:botName', async (req, res) => {
     }
 });
 
-// --- ROUTE 2: CHATBOT CONFIG ---
 app.get('/api/v1/public-chatbotConfig/:instanceNum/:botName', async (req, res) => {
     try {
         const instanceNum = parseInt(req.params.instanceNum);
@@ -347,7 +399,6 @@ app.get('/api/v1/public-chatbotConfig/:instanceNum/:botName', async (req, res) =
     }
 });
 
-// --- ROUTE 3: STREAMING CHECK ---
 app.get('/api/v1/chatflows-streaming/:instanceNum/:botName', async (req, res) => {
     try {
         const instanceNum = parseInt(req.params.instanceNum);
@@ -379,7 +430,6 @@ app.get('/api/v1/chatflows-streaming/:instanceNum/:botName', async (req, res) =>
     }
 });
 
-// --- HEALTH CHECK ---
 app.get('/', (req, res) => res.send('Federated Proxy Active'));
 
 app.get('/health', (req, res) => {
@@ -392,18 +442,15 @@ app.get('/health', (req, res) => {
     });
 });
 
-// --- 404 HANDLER ---
 app.use((req, res) => {
     res.status(404).json({ error: 'Route not found' });
 });
 
-// --- GLOBAL ERROR HANDLER ---
 app.use((err, req, res, next) => {
     console.error('[Error] Unhandled error:', err);
     res.status(500).json({ error: 'Internal server error' });
 });
 
-// --- GRACEFUL SHUTDOWN ---
 const server = app.listen(7860, '0.0.0.0', () => {
     console.log('Federated Proxy running on port 7860');
 });