1  import mongoose from "mongoose";
  2  import crypto from "crypto";
  3  import config from "@/config.js";
  4  import { logger } from "@/middleware/logging.js";
  5  
  6  let isConnected = false;
  7  let reconnectAttempts = 0;
  8  let reconnectTimer = null;
  9  let connectionPool = null;
 10  const queryMetrics = new Map();
 11  
 12  const sanitizeQuery = (query, depth = 0) => {
 13    if (!query || typeof query !== "object") {
 14      return {};
 15    } else if (depth > 3) {
 16      return {};
 17    }
 18  
 19    const sanitized = {};
 20    const dangerousOperators = [
 21      "$where", "$function", "$accumulator", "$reduce", "$expr",
 22      "$jsonSchema", "$eval", "$comment",
 23    ];
 24  
 25    const safeOperators = [
 26      "$gt", "$gte", "$lt", "$lte", "$ne", "$in", "$nin", "$exists",
 27      "$not", "$nor", "$and", "$or", "$regex", "$text", "$elemMatch",
 28      "$size", "$mod", "$all", "$max", "$min", "$mul", "$type",
 29      "$setOnInsert", "$each", "$slice", "$sort", "$addToSet", "$push",
 30    ];
 31  
 32    for (const [key, value] of Object.entries(query)) {
 33      if (dangerousOperators.includes(key)) {
 34        logger.warn(`Blocked dangerous operator in query: ${key}`);
 35        continue;
 36      }
 37  
 38      if (key === "__proto__" || key === "constructor" || key === "prototype") {
 39        continue;
 40      }
 41  
 42      if (typeof value === "string") {
 43        if (value.length > 1000) {
 44          continue;
 45        }
 46        if (value.includes("$where") || value.includes("function(")) {
 47          logger.warn(`Blocked potentially malicious string in query: ${key}`);
 48          continue;
 49        }
 50        sanitized[key] = value;
 51      } else if (Array.isArray(value)) {
 52        if (value.length > 100) {
 53          continue;
 54        }
 55        sanitized[key] = value.slice(0, 100).map(item => {
 56          if (typeof item === "object" && item !== null) {
 57            return depth < 2 ? sanitizeQuery(item, depth + 1) : {};
 58          }
 59          return item;
 60        });
 61      } else if (typeof value === "object" && value !== null) {
 62        if (safeOperators.includes(key)) {
 63          sanitized[key] = depth < 2 ? sanitizeQuery(value, depth + 1) : {};
 64        } else {
 65          sanitized[key] = depth < 2 ? sanitizeQuery(value, depth + 1) : {};
 66        }
 67      } else {
 68        sanitized[key] = value;
 69      }
 70    }
 71    return sanitized;
 72  };
 73  
 74  const validateConnectionString = (uri) => {
 75    if (!uri || typeof uri !== "string") {
 76      return false;
 77    } else if (uri.length > 2048) {
 78      return false;
 79    } else {
 80      const validProtocols = ["mongodb:"];
 81      const hasValidProtocol = validProtocols.some(protocol => uri.startsWith(protocol));
 82      if (!hasValidProtocol) {
 83        return false;
 84      }
 85      if (uri.includes("@") && !uri.includes("mongodb+srv:") && !uri.includes("ssl=true")) {
 86        return false;
 87      }
 88      return true;
 89    }
 90  };
 91  
 92  export const connectDatabase = async () => {
 93    if (isConnected && mongoose.connection.readyState === 1) {
 94      logger.info("Database already connected");
 95      return mongoose.connection;
 96    }
 97  
 98    if (!validateConnectionString(config.database.uri)) {
 99      logger.error("Invalid database connection string");
100      throw new Error("Invalid database configuration");
101    }
102  
103    try {
104      mongoose.set("strictQuery", true);
105      mongoose.set("sanitizeFilter", true);
106      mongoose.set("sanitizeProjection", true);
107  
108      const secureOptions = {
109        ...config.database.options,
110        autoIndex: false,
111        serverSelectionTimeoutMS: config.timeouts.short || 5000,
112        socketTimeoutMS: config.timeouts.medium || 10000,
113        maxPoolSize: Math.min(config.database.options.maxPoolSize || 50, 100),
114        minPoolSize: Math.max(config.database.options.minPoolSize || 5, 2),
115        maxIdleTimeMS: config.timeouts.long || 30000,
116        waitQueueTimeoutMS: config.timeouts.medium || 10000,
117        family: 4,
118        retryReads: true,
119        retryWrites: true,
120        readPreference: "primary",
121        writeConcern: { w: "majority", j: true },
122      };
123  
124      await mongoose.connect(config.database.uri, secureOptions);
125      connectionPool = mongoose.connection;
126      isConnected = true;
127      reconnectAttempts = 0;
128  
129      logger.important(`Connected to MongoDB: ${config.database.dbName}`);
130      return mongoose.connection;
131    } catch (error) {
132      logger.error("Database connection error:", { message: error.message });
133      await handleReconnect();
134      throw error;
135    }
136  };
137  
138  const handleReconnect = async () => {
139    if (reconnectTimer) {
140      return;
141    }
142  
143    if (reconnectAttempts >= config.retries.some) {
144      logger.error("Max reconnection attempts reached");
145      process.exit(1);
146    }
147  
148    reconnectAttempts++;
149    const jitter = crypto.randomInt(0, 1000);
150    const baseDelay = (config.timeouts.short || 5000) * Math.pow(2, reconnectAttempts - 1);
151    const delay = Math.min(baseDelay + jitter, 60000);
152    logger.info(`Reconnection attempt ${reconnectAttempts}/${config.retries.some} in ${delay}ms`);
153  
154    reconnectTimer = setTimeout(async () => {
155      reconnectTimer = null;
156      try {
157        await connectDatabase();
158      } catch (err) {
159        logger.error("Reconnection failed:", { message: err.message });
160      }
161    }, delay);
162  };
163  
164  export { sanitizeQuery };
165  
166  const trackQueryMetrics = (operation, collection, duration) => {
167    const key = `${collection}:${operation}`;
168    const existing = queryMetrics.get(key) || { count: 0, totalDuration: 0, maxDuration: 0 };
169    existing.count++;
170    existing.totalDuration += duration;
171    existing.maxDuration = Math.max(existing.maxDuration, duration);
172    existing.lastSeen = Date.now();
173    queryMetrics.set(key, existing);
174  
175    const maxSize = Math.min(config.database.performance.maxQueryMetricsSize || 10000, 50000);
176    if (queryMetrics.size > maxSize) {
177      const now = Date.now();
178      const staleThreshold = 3600000;
179      const entries = Array.from(queryMetrics.entries());
180  
181      const staleEntries = entries.filter(([_, metrics]) =>
182        now - (metrics.lastSeen || 0) > staleThreshold,
183      );
184  
185      if (staleEntries.length > 0) {
186        staleEntries.forEach(([key]) => queryMetrics.delete(key));
187      } else {
188        const oldestEntries = entries
189          .sort((a, b) => (a[1].lastSeen || 0) - (b[1].lastSeen || 0))
190          .slice(0, Math.floor(maxSize * 0.2));
191        oldestEntries.forEach(([key]) => queryMetrics.delete(key));
192      }
193    }
194  
195    if (duration > config.database.performance.slowQueryThreshold) {
196      logger.warn("Slow query detected", { collection, operation, duration });
197    }
198  };
199  
200  mongoose.plugin((schema) => {
201    schema.pre("find", function () {
202      this._startTime = Date.now();
203      if (this._conditions) {
204        this._conditions = sanitizeQuery(this._conditions);
205      }
206    });
207  
208    schema.post("find", function () {
209      if (this._startTime) {
210        trackQueryMetrics("find", this.model.collection.name, Date.now() - this._startTime);
211      }
212    });
213  
214    schema.pre("findOne", function () {
215      this._startTime = Date.now();
216      if (this._conditions) {
217        this._conditions = sanitizeQuery(this._conditions);
218      }
219    });
220  
221    schema.post("findOne", function () {
222      if (this._startTime) {
223        trackQueryMetrics("findOne", this.model.collection.name, Date.now() - this._startTime);
224      }
225    });
226  
227    schema.pre("save", function () {
228      this._startTime = Date.now();
229    });
230  
231    schema.post("save", function () {
232      if (this._startTime) {
233        trackQueryMetrics("save", this.model?.collection?.name || "unknown", Date.now() - this._startTime);
234      }
235    });
236  });
237  
238  mongoose.connection.on("connected", () => {
239    isConnected = true;
240    reconnectAttempts = 0;
241    if (reconnectTimer) {
242      clearTimeout(reconnectTimer);
243      reconnectTimer = null;
244    }
245    logger.info("Mongoose connected to database");
246  });
247  
248  mongoose.connection.on("error", (err) => {
249    isConnected = false;
250    logger.error("Mongoose connection error:", { message: err.message, code: err.code });
251  });
252  
253  mongoose.connection.on("disconnected", () => {
254    isConnected = false;
255    logger.warn("Mongoose disconnected");
256    if (!reconnectTimer) {
257      handleReconnect();
258    }
259  });
260  
261  mongoose.connection.on("reconnected", () => {
262    isConnected = true;
263    reconnectAttempts = 0;
264    logger.info("Mongoose reconnected to database");
265  });
266  
267  export const close = async () => {
268    if (reconnectTimer) {
269      clearTimeout(reconnectTimer);
270      reconnectTimer = null;
271    }
272    try {
273      queryMetrics.clear();
274      await mongoose.connection.close(false);
275      isConnected = false;
276      logger.info("Database connection closed gracefully");
277    } catch (err) {
278      logger.error("Error closing database connection:", { message: err.message });
279    }
280  };
281  
282  process.on("SIGINT", async () => {
283    await close();
284    process.exit(0);
285  });
286  
287  process.on("SIGTERM", async () => {
288    await close();
289    process.exit(0);
290  });
291  
292  export const getConnectionStatus = () => ({
293    connected: isConnected,
294    readyState: mongoose.connection.readyState,
295    poolSize: connectionPool?.client?.topology?.s?.pool?.size || 0,
296  });
297  
298  export const getQueryMetrics = () => {
299    const metrics = {};
300    for (const [key, value] of queryMetrics.entries()) {
301      metrics[key] = {
302        ...value,
303        avgDuration: value.count > 0 ? Math.round(value.totalDuration / value.count) : 0,
304      };
305    }
306    return metrics;
307  };
308  
309  export const healthCheck = async () => {
310    try {
311      if (!isConnected || mongoose.connection.readyState !== 1) {
312        return { healthy: false, reason: "Not connected" };
313      }
314      const start = Date.now();
315      await mongoose.connection.db.admin().ping();
316      const latency = Date.now() - start;
317      return { healthy: true, latency, readyState: mongoose.connection.readyState };
318    } catch (err) {
319      return { healthy: false, reason: err.message };
320    }
321  };