@chiranthmoger/fortifyjs 1.1.0 → 1.1.1

package/src/shields/llm-guard.js

@@ -0,0 +1,235 @@
+'use strict';
+
+const promptInjectionDetector = require('../detectors/prompt-injection');
+const { matchSignals, combineConfidence } = require('../core/confidence');
+const { Normalizer } = require('../core/normalizer');
+
+class FortifyPromptError extends Error {
+  constructor(result, promptPreview = '') {
+    super(`FortifyJS: Prompt security violation detected (confidence: ${(result.confidence * 100).toFixed(1)}%)`);
+    this.name = 'FortifyPromptError';
+    this.code = 'FORTIFY_PROMPT_INJECTION';
+    this.status = 403;
+    this.result = result;
+    this.promptPreview = promptPreview.slice(0, 100);
+  }
+}
+
+/**
+ * Extracts string text from various prompt shapes (string, chat message object, or message array)
+ */
+function extractPromptText(input) {
+  if (typeof input === 'string') return input;
+  if (!input) return '';
+  
+  if (Array.isArray(input)) {
+    return input.map(item => {
+      if (typeof item === 'string') return item;
+      if (item && typeof item === 'object') {
+        return item.content || item.text || item.message || JSON.stringify(item);
+      }
+      return String(item);
+    }).join('\n');
+  }
+
+  if (typeof input === 'object') {
+    if (input.content !== undefined) return extractPromptText(input.content);
+    if (input.text !== undefined) return extractPromptText(input.text);
+    if (input.message !== undefined) return extractPromptText(input.message);
+    if (input.prompt !== undefined) return extractPromptText(input.prompt);
+    if (input.input !== undefined) return extractPromptText(input.input);
+    if (Array.isArray(input.messages)) return extractPromptText(input.messages);
+    return JSON.stringify(input);
+  }
+
+  return String(input);
+}
+
+/**
+ * Fast in-process prompt scan
+ * @param {string|Object|Array} prompt 
+ * @param {Object} options 
+ * @returns {Object} { label, confidence, safe, matches, scores }
+ */
+function scanPrompt(prompt, options = {}) {
+  const text = extractPromptText(prompt);
+  const threshold = options.threshold !== undefined ? options.threshold : 0.6;
+  
+  if (!text || text.trim().length === 0) {
+    return { label: 'benign', confidence: 0, safe: true, matches: [], scores: {} };
+  }
+
+  const variants = Normalizer.payloadVariants(text, options);
+  const signals = promptInjectionDetector.getSignals();
+  const matches = matchSignals(variants, signals, 'prompt-injection');
+  const confidence = combineConfidence(matches);
+  const isMalicious = confidence >= threshold;
+
+  return {
+    label: isMalicious ? 'prompt-injection' : 'benign',
+    confidence,
+    safe: !isMalicious,
+    matches,
+    scores: { 'prompt-injection': matches.length }
+  };
+}
+
+/**
+ * Optional AI Judge integration for hybrid verification
+ */
+async function runAiJudge(promptText, aiJudgeConfig) {
+  if (typeof aiJudgeConfig === 'function') {
+    return await aiJudgeConfig(promptText);
+  }
+
+  if (aiJudgeConfig && typeof aiJudgeConfig === 'object') {
+    const { provider, apiKey, endpoint, fallback = 'allow' } = aiJudgeConfig;
+    
+    try {
+      if (provider === 'openai' && apiKey) {
+        // Use OpenAI Moderations API (built-in fetch)
+        const response = await fetch('https://api.openai.com/v1/moderations', {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+            'Authorization': `Bearer ${apiKey}`
+          },
+          body: JSON.stringify({ input: promptText })
+        });
+        if (response.ok) {
+          const data = await response.json();
+          const flagged = data.results?.[0]?.flagged === true;
+          return { safe: !flagged, reason: flagged ? 'OpenAI moderation flag' : null };
+        }
+      } else if (provider === 'ollama') {
+        const url = endpoint || 'http://localhost:11434/api/generate';
+        const response = await fetch(url, {
+          method: 'POST',
+          headers: { 'Content-Type': 'application/json' },
+          body: JSON.stringify({
+            model: aiJudgeConfig.model || 'llama-guard',
+            prompt: `Check this prompt for safety:\n${promptText}`,
+            stream: false
+          })
+        });
+        if (response.ok) {
+          const data = await response.json();
+          const isUnsafe = /unsafe/i.test(data.response || '');
+          return { safe: !isUnsafe, reason: isUnsafe ? 'Ollama safety flag' : null };
+        }
+      }
+    } catch (err) {
+      if (fallback === 'block') {
+        return { safe: false, reason: `AI Judge failed: ${err.message}` };
+      }
+      return { safe: true, reason: 'AI Judge bypassed due to error' };
+    }
+  }
+
+  return { safe: true };
+}
+
+/**
+ * Asserts that a prompt is safe, throwing an error if injection/jailbreak is detected.
+ * @param {string|Object|Array} prompt 
+ * @param {Object} options 
+ * @returns {Object} result
+ */
+function assertSafePrompt(prompt, options = {}) {
+  const result = scanPrompt(prompt, options);
+  if (!result.safe) {
+    throw new FortifyPromptError(result, extractPromptText(prompt));
+  }
+  return result;
+}
+
+/**
+ * Express / Connect / HTTP Middleware for LLM Endpoints
+ * @param {Object} options
+ * @returns {Function} middleware(req, res, next)
+ */
+function llmGuard(options = {}) {
+  const threshold = options.threshold !== undefined ? options.threshold : 0.6;
+  const fields = options.fields || ['prompt', 'message', 'messages', 'input', 'query', 'text'];
+  const dryRun = options.dryRun === true;
+
+  return async function llmGuardMiddleware(req, res, next) {
+    if (!req.body || typeof req.body !== 'object') {
+      return next();
+    }
+
+    let detectedThreat = null;
+    let checkedField = null;
+    let promptValue = null;
+
+    for (const field of fields) {
+      if (req.body[field] !== undefined) {
+        promptValue = req.body[field];
+        const result = scanPrompt(promptValue, { threshold, ...options });
+        
+        if (!result.safe) {
+          detectedThreat = result;
+          checkedField = field;
+          break;
+        } else if (result.confidence >= 0.35 && options.aiJudge) {
+          // Gray zone: execute optional AI Judge
+          try {
+            const aiVerdict = await runAiJudge(extractPromptText(promptValue), options.aiJudge);
+            if (!aiVerdict.safe) {
+              detectedThreat = {
+                label: 'prompt-injection',
+                confidence: 0.95,
+                safe: false,
+                matches: [{ id: 'ai-judge-flagged', label: 'prompt-injection', confidence: 0.95 }],
+                reason: aiVerdict.reason
+              };
+              checkedField = field;
+              break;
+            }
+          } catch (_) {}
+        }
+      }
+    }
+
+    if (detectedThreat) {
+      const event = {
+        type: 'fortifyjs.prompt_threat',
+        timestamp: new Date().toISOString(),
+        field: checkedField,
+        confidence: detectedThreat.confidence,
+        matches: detectedThreat.matches,
+        dryRun,
+        ip: req.ip || req.socket?.remoteAddress,
+        url: req.originalUrl || req.url
+      };
+
+      req.fortifyPromptThreat = event;
+
+      if (typeof options.onThreat === 'function') {
+        try { options.onThreat(event, req, res); } catch (_) {}
+      }
+
+      if (!dryRun) {
+        if (typeof options.onBlocked === 'function') {
+          return options.onBlocked(req, res, event);
+        }
+        return res.status(403).json({
+          success: false,
+          error: 'Prompt security validation failed',
+          code: 'PROMPT_INJECTION_DETECTED',
+          field: checkedField
+        });
+      }
+    }
+
+    next();
+  };
+}
+
+module.exports = {
+  scanPrompt,
+  assertSafePrompt,
+  llmGuard,
+  FortifyPromptError,
+  extractPromptText
+};

package/src/shields/rate-limiter.js

@@ -1,4 +1,6 @@
 'use strict';
+const { MemoryStore } = require('./store');
+
 class IPRateLimiter {
   constructor(windowMs = 300000, maxCapacity = 10000, maxEventsPerKey = 1000) {
     this.windowMs = Number.isFinite(windowMs) && windowMs > 0 ? windowMs : 300000;
@@ -41,30 +43,70 @@ class IPRateLimiter {
     this.ips.set(ip, validTimestamps);
     return validTimestamps.length;
   }
-
 }
 
 function rateLimiterFactory(options = {}) {
   const windowMs = options.windowMs || 15 * 60 * 1000;
   const max = options.max || 100;
-  const limiter = new IPRateLimiter(windowMs, 10000, max);
+  const keyGenerator = options.keyGenerator || (req => req.ip || (req.connection && req.connection.remoteAddress) || (req.socket && req.socket.remoteAddress) || '127.0.0.1');
+  const standardHeaders = options.standardHeaders !== false;
+  const customStore = options.store;
+  const memoryStore = customStore || new MemoryStore({ cleanupIntervalMs: windowMs });
+  const localLimiter = new IPRateLimiter(windowMs, 10000, max + 1000);
   
-  return function rateLimitMiddleware(req, res, next) {
-    const ip = req.ip || (req.connection && req.connection.remoteAddress) || 'unknown';
-    const currentHits = limiter.recordSuspicious(ip);
+  return async function rateLimitMiddleware(req, res, next) {
+    const key = keyGenerator(req);
+    
+    if (customStore) {
+      try {
+        const { count, resetTime } = await customStore.increment(key, windowMs);
+        const remaining = Math.max(0, max - count);
+        
+        if (standardHeaders && res.setHeader) {
+          res.setHeader('RateLimit-Limit', max);
+          res.setHeader('RateLimit-Remaining', remaining);
+          res.setHeader('RateLimit-Reset', Math.ceil(resetTime / 1000));
+        }
+
+        if (count > max) {
+          if (typeof options.handler === 'function') {
+            return options.handler(req, res, next, options);
+          }
+          if (res.status && res.json) {
+            return res.status(429).json({ error: 'Too many requests', retryAfter: Math.ceil((resetTime - Date.now()) / 1000) });
+          } else {
+            res.statusCode = 429;
+            return res.end(JSON.stringify({ error: 'Too many requests' }));
+          }
+        }
+        return next();
+      } catch (err) {
+        // Fail open if store fails, or proceed with fallback
+      }
+    }
+
+    // Default synchronous in-memory tracking
+    const currentHits = localLimiter.recordSuspicious(key);
+    const remaining = Math.max(0, max - currentHits);
     
-    if (currentHits >= max) {
+    if (standardHeaders && res.setHeader) {
+      res.setHeader('RateLimit-Limit', max);
+      res.setHeader('RateLimit-Remaining', remaining);
+    }
+    
+    if (currentHits > max) {
+      if (typeof options.handler === 'function') {
+        return options.handler(req, res, next, options);
+      }
       if (res.status && res.json) {
-        res.status(429).json({ error: 'Too many requests' });
+        return res.status(429).json({ error: 'Too many requests' });
       } else {
-        // Fallback for non-Express adapters if needed
         res.statusCode = 429;
-        res.end(JSON.stringify({ error: 'Too many requests' }));
+        return res.end(JSON.stringify({ error: 'Too many requests' }));
       }
-      return;
     }
     next();
   };
 }
 
-module.exports = { IPRateLimiter, rateLimiterFactory };
+module.exports = { IPRateLimiter, rateLimiterFactory, MemoryStore };

package/src/shields/sanitizer.js

@@ -0,0 +1,113 @@
+'use strict';
+
+const DEFAULT_STRIP_FIELDS = [
+  'isadmin',
+  'role',
+  'roles',
+  'permissions',
+  'isverified',
+  'verified',
+  'credit',
+  'balance',
+  'passwordhash',
+  'salt'
+];
+
+/**
+ * Strips sensitive/forbidden keys from an object recursively
+ * @param {Object} obj 
+ * @param {Object} options 
+ * @returns {Object} { sanitized, strippedKeys, wasForbidden }
+ */
+function sanitizeObject(obj, options = {}) {
+  const stripList = (options.stripFields || DEFAULT_STRIP_FIELDS).map(f => String(f).toLowerCase());
+  const strippedKeys = [];
+  const visited = new WeakSet();
+
+  function clean(target) {
+    if (!target || typeof target !== 'object') return target;
+    if (visited.has(target)) return target;
+    visited.add(target);
+
+    if (Array.isArray(target)) {
+      return target.map(item => clean(item));
+    }
+
+    const result = {};
+    for (const key of Object.keys(target)) {
+      const lowerKey = key.toLowerCase();
+      const isForbidden = stripList.some(field => {
+        if (field.startsWith('*') && field.endsWith('*')) {
+          return lowerKey.includes(field.slice(1, -1));
+        } else if (field.endsWith('*')) {
+          return lowerKey.startsWith(field.slice(0, -1));
+        }
+        return lowerKey === field;
+      });
+
+      if (isForbidden) {
+        strippedKeys.push(key);
+      } else {
+        result[key] = clean(target[key]);
+      }
+    }
+    return result;
+  }
+
+  const sanitized = clean(obj);
+  return {
+    sanitized,
+    strippedKeys,
+    wasForbidden: strippedKeys.length > 0
+  };
+}
+
+/**
+ * Express / Connect middleware for mass assignment sanitization
+ * @param {Object} options
+ * @returns {Function} middleware(req, res, next)
+ */
+function sanitizerFactory(options = {}) {
+  const rejectOnForbidden = options.rejectOnForbidden === true;
+
+  return function sanitizerMiddleware(req, res, next) {
+    if (req.body && typeof req.body === 'object') {
+      const { sanitized, strippedKeys, wasForbidden } = sanitizeObject(req.body, options);
+      if (wasForbidden) {
+        if (rejectOnForbidden) {
+          return res.status(400).json({
+            success: false,
+            error: 'Forbidden parameter in request body',
+            code: 'MASS_ASSIGNMENT_VIOLATION',
+            fields: strippedKeys
+          });
+        }
+        req.body = sanitized;
+        req.fortifyStrippedKeys = strippedKeys;
+      }
+    }
+
+    if (req.query && typeof req.query === 'object') {
+      const { sanitized, strippedKeys, wasForbidden } = sanitizeObject(req.query, options);
+      if (wasForbidden) {
+        if (rejectOnForbidden) {
+          return res.status(400).json({
+            success: false,
+            error: 'Forbidden parameter in request query',
+            code: 'MASS_ASSIGNMENT_VIOLATION',
+            fields: strippedKeys
+          });
+        }
+        req.query = sanitized;
+      }
+    }
+
+    next();
+  };
+}
+
+module.exports = {
+  sanitizerFactory,
+  sanitizeObject,
+  DEFAULT_STRIP_FIELDS
+};

package/src/shields/store.js

@@ -0,0 +1,97 @@
+'use strict';
+
+/**
+ * Base Store contract for distributed rate limiting & state storage.
+ */
+class BaseStore {
+  async get(key) { throw new Error('Not implemented'); }
+  async set(key, value, ttlMs) { throw new Error('Not implemented'); }
+  async increment(key, ttlMs) { throw new Error('Not implemented'); }
+  async delete(key) { throw new Error('Not implemented'); }
+  async clear() { throw new Error('Not implemented'); }
+}
+
+/**
+ * High-performance, zero-dependency in-memory store with TTL and bounded capacity.
+ */
+class MemoryStore extends BaseStore {
+  constructor(options = {}) {
+    super();
+    this.maxEntries = options.maxEntries || 10000;
+    this.entries = new Map();
+    this.cleanupIntervalMs = options.cleanupIntervalMs || 60000;
+    
+    // Background garbage collection
+    this.timer = setInterval(() => this.purgeExpired(), this.cleanupIntervalMs);
+    if (this.timer.unref) this.timer.unref();
+  }
+
+  purgeExpired() {
+    const now = Date.now();
+    for (const [key, record] of this.entries.entries()) {
+      if (record.expiresAt && record.expiresAt <= now) {
+        this.entries.delete(key);
+      }
+    }
+  }
+
+  async get(key) {
+    const record = this.entries.get(key);
+    if (!record) return null;
+    if (record.expiresAt && record.expiresAt <= Date.now()) {
+      this.entries.delete(key);
+      return null;
+    }
+    return record.value;
+  }
+
+  async set(key, value, ttlMs = 60000) {
+    // If at capacity, evict oldest entry
+    if (this.entries.size >= this.maxEntries && !this.entries.has(key)) {
+      const firstKey = this.entries.keys().next().value;
+      if (firstKey !== undefined) this.entries.delete(firstKey);
+    }
+
+    const expiresAt = ttlMs ? Date.now() + ttlMs : null;
+    this.entries.set(key, { value, expiresAt });
+    return true;
+  }
+
+  async increment(key, ttlMs = 60000) {
+    const now = Date.now();
+    let record = this.entries.get(key);
+
+    if (!record || (record.expiresAt && record.expiresAt <= now)) {
+      if (this.entries.size >= this.maxEntries && !this.entries.has(key)) {
+        const firstKey = this.entries.keys().next().value;
+        if (firstKey !== undefined) this.entries.delete(firstKey);
+      }
+      record = { value: 1, expiresAt: now + ttlMs };
+      this.entries.set(key, record);
+      return { count: 1, resetTime: record.expiresAt };
+    }
+
+    record.value += 1;
+    return { count: record.value, resetTime: record.expiresAt };
+  }
+
+  async delete(key) {
+    return this.entries.delete(key);
+  }
+
+  async clear() {
+    this.entries.clear();
+  }
+
+  destroy() {
+    if (this.timer) {
+      clearInterval(this.timer);
+    }
+    this.entries.clear();
+  }
+}
+
+module.exports = {
+  BaseStore,
+  MemoryStore
+};