Update server.js

server.js

@@ -13,51 +13,89 @@ const openaiKey = process.env.OPENAI_KEY;
 const port = 7860;
 const baseUrl = getExternalUrl(process.env.SPACE_ID);
 
-let requestDetails = []; // Store details of each request
-let ipRequestCounts = {}; // Dictionary to keep track of requests per IP
+let requestDetails = [];
+let apiKeyDetails = {}; // Track API key usage and first use timestamp
 
 io.on('connection', (socket) => {
     console.log('A user connected to the websocket for live logs.');
 });
 
-function logAndEmit(message) {
-    console.log(message);  // Continue to log to the console
-    io.emit('log', message);  // Emit this log to the frontend via WebSocket
+function logAndEmit(message, includeTotal = false) {
+    if (includeTotal) {
+        const totalRequests = getTotalRequests();
+        message += ` Total requests so far: ${totalRequests}`;
+    }
+    console.log(message);
+    io.emit('log', message);
 }
 
-app.use('/api', (req, res, next) => {
-  if (req.method === 'POST' && req.url.startsWith('/v1/chat/completions')) {
-    const ip = req.headers['x-forwarded-for'] || req.connection.remoteAddress;
-    ipRequestCounts[ip] = (ipRequestCounts[ip] || 0) + 1;
+function getTotalRequests() {
+    return Object.values(apiKeyDetails).reduce((acc, details) => acc + details.count, 0);
+}
+
+function authenticateApiKey(req, res, next) {
+    const receivedApiKey = req.headers['authorization'];
+    if (!receivedApiKey) {
+        return res.status(401).send('Unauthorized: No API key provided');
+    }
+
+    const validApiKeys = (process.env.SECRET_API_KEYS || '').split(',');
+    const apiKeyMatch = validApiKeys.find(key => `Bearer ${key.trim()}` === receivedApiKey);
     
-    const logMessage = `IP ${ip} has made ${ipRequestCounts[ip]} requests so far.`;
-    logAndEmit(logMessage);
-  }
+    if (!apiKeyMatch) {
+        return res.status(401).send('Unauthorized: API key is invalid');
+    }
+
+    const displayKey = apiKeyMatch.trim().substring(0, 7); // Only use the first four characters for display
+
+    if (!apiKeyDetails[apiKeyMatch]) {
+        apiKeyDetails[apiKeyMatch] = {
+            count: 0,
+            firstUse: moment()  // Set the first use to now if it's not already set
+        };
+    }
+
+    // Check if more than a week has passed since the first use
+    if (moment().diff(apiKeyDetails[apiKeyMatch].firstUse, 'weeks') >= 1) {
+        return res.status(429).send(`API Key ${displayKey}... is blocked after one week of use.`);
+    }
 
-  const timestamp = moment().tz("Asia/Tbilisi").format("DD-MMM-YYYY HH:mm");
-  const requestData = {
-    requestNumber: ipRequestCounts[req.headers['x-forwarded-for'] || req.connection.remoteAddress],
-    ip: req.headers['x-forwarded-for'] || req.connection.remoteAddress,
-    timestamp: timestamp,
-    text: req.method + ' ' + req.url
-  };
-  requestDetails.push(requestData);
-  
-  const detailedLog = `Request ${requestData.requestNumber} on ${requestData.timestamp} from IP ${requestData.ip} with method "${req.method}" and text "${requestData.text}"`;
-  logAndEmit(detailedLog);
-
-  next();
+    // Increment the request count
+    apiKeyDetails[apiKeyMatch].count++;
+
+    const logMessage = `API Key ${displayKey}... used ${apiKeyDetails[apiKeyMatch].count} times so far.`;
+    logAndEmit(logMessage, true);
+    
+    req.apiKey = apiKeyMatch; // Store the API key for the request
+    next();
+}
+
+app.use('/api', authenticateApiKey, (req, res, next) => {
+    const timestamp = moment().tz("Asia/Tbilisi").format("DD-MMM-YYYY HH:mm");
+    const displayKey = req.apiKey.substring(0, 4);
+    const requestData = {
+        requestNumber: apiKeyDetails[req.apiKey].count,
+        apiKey: displayKey,
+        timestamp: timestamp,
+        text: req.method + ' ' + req.url
+    };
+    requestDetails.push(requestData);
+
+    const detailedLog = `Request ${requestData.requestNumber} on ${requestData.timestamp} from API Key ${displayKey}... with method "${req.method}" and text "${requestData.text}"`;
+    logAndEmit(detailedLog, true);
+
+    next();
 }, proxy(targetUrl, {
-  proxyReqOptDecorator: function(proxyReqOpts, srcReq) {
-    proxyReqOpts.headers['Authorization'] = 'Bearer ' + openaiKey;
-    return proxyReqOpts;
-  }
+    proxyReqOptDecorator: function(proxyReqOpts, srcReq) {
+        proxyReqOpts.headers['Authorization'] = 'Bearer ' + openaiKey;
+        return proxyReqOpts;
+    }
 }));
 
 app.get("/", (req, res) => {
-  let requestsInfo = requestDetails.map(detail =>
-    `Request ${detail.requestNumber} on ${detail.timestamp} from IP ${detail.ip} with text "${detail.text}"`).join('<br>');
-  res.send(`This is your OpenAI Reverse Proxy URL: ${baseUrl}.<br><br>Requests Overview:<br>${requestsInfo}`);
+    let requestsInfo = requestDetails.map(detail =>
+        `Request ${detail.requestNumber} on ${detail.timestamp} from API Key ${detail.apiKey}... with text "${detail.text}"`).join('<br>');
+    res.send(`This is your OpenAI Reverse Proxy URL: ${baseUrl}.<br><br>Requests Overview:<br>${requestsInfo}`);
 });
 
 app.get('/logs', (req, res) => {
@@ -74,5 +112,6 @@ function getExternalUrl(spaceId) {
 }
 
 server.listen(port, () => {
-   logAndEmit(`Reverse proxy server and WebSocket running on port ${port}`);
+    const message = `Reverse proxy server and WebSocket running on port ${port}`;
+    logAndEmit(message, true); // true to include total requests in the log
 });