agentshield-sdk 13.5.0 → 14.0.0

package/src/middleware.js

@@ -14,11 +14,87 @@ const { createShieldError } = require('./errors');
 /** Coerce any value to a scannable string. */
 const textify = (val) => typeof val === 'string' ? val : (val != null ? JSON.stringify(val) : '');
 
+/**
+ * Default maximum body size (in bytes) enforced by expressMiddleware
+ * when `options.maxBodySize` is not provided. Defaults to 1 MB.
+ */
+const DEFAULT_MAX_BODY_SIZE = 1 * 1024 * 1024;
+
+/**
+ * Computes the approximate size in bytes of a parsed request body.
+ * - String: exact UTF-8 byte length
+ * - Buffer: exact length
+ * - Object: JSON.stringify length (fallback)
+ *
+ * @param {*} body
+ * @returns {number}
+ */
+const computeBodySize = (body) => {
+  if (body == null) return 0;
+  if (Buffer.isBuffer(body)) return body.length;
+  if (typeof body === 'string') return Buffer.byteLength(body, 'utf8');
+  if (typeof body === 'object') {
+    try {
+      return JSON.stringify(body).length;
+    } catch (_) {
+      return 0;
+    }
+  }
+  return 0;
+};
+
+/**
+ * Attaches a cumulative byte-counter to the raw request stream and aborts
+ * the request with 413 once the configured limit is exceeded. This runs
+ * in addition to the post-parse body size check so attackers cannot
+ * bypass the limit by streaming a huge payload before the body parser
+ * buffers it.
+ *
+ * @param {import('http').IncomingMessage} req
+ * @param {import('http').ServerResponse} res
+ * @param {number} limit
+ * @returns {boolean} True if the stream watcher was attached.
+ */
+const attachRawSizeGuard = (req, res, limit) => {
+  if (!req || typeof req.on !== 'function') return false;
+  // Already read/parsed — nothing to guard.
+  if (req._agentShieldRawGuardAttached) return false;
+  req._agentShieldRawGuardAttached = true;
+
+  let received = 0;
+  const onData = (chunk) => {
+    received += chunk ? chunk.length : 0;
+    if (received > limit) {
+      req.removeListener('data', onData);
+      try {
+        if (typeof req.pause === 'function') req.pause();
+        if (!res.headersSent) {
+          res.status(413).json({
+            error: 'Payload Too Large',
+            message: `Request body exceeds maximum allowed size of ${limit} bytes`,
+            maxBodySize: limit
+          });
+        }
+        if (typeof req.destroy === 'function') req.destroy();
+      } catch (_) {
+        // Swallow — the response has already been sent or the socket closed.
+      }
+    }
+  };
+  req.on('data', onData);
+  return true;
+};
+
 /**
  * Creates an Express/Connect-style middleware that scans request bodies
  * for AI-specific threats before they reach your agent endpoint.
  *
+ * Enforces a configurable body-size limit (default 1MB) so callers do
+ * not need to configure body-parser separately. Oversized payloads are
+ * rejected with HTTP 413 before any scanning takes place.
+ *
  * @param {object} [config] - AgentShield configuration.
+ * @param {number} [config.maxBodySize=1048576] - Maximum accepted request body size in bytes.
  * @returns {Function} Express middleware function.
  *
  * @example
@@ -27,7 +103,7 @@ const textify = (val) => typeof val === 'string' ? val : (val != null ? JSON.str
  *
  * const app = express();
  * app.use(express.json());
- * app.use(expressMiddleware({ blockOnThreat: true, blockThreshold: 'high' }));
+ * app.use(expressMiddleware({ blockOnThreat: true, blockThreshold: 'high', maxBodySize: 512 * 1024 }));
  *
  * app.post('/agent', (req, res) => {
  *   // req.agentShield contains scan results
@@ -39,13 +115,33 @@ const textify = (val) => typeof val === 'string' ? val : (val != null ? JSON.str
  */
 const expressMiddleware = (config = {}) => {
   const shield = new AgentShield({ blockOnThreat: true, ...config });
+  const maxBodySize = Number.isFinite(config.maxBodySize) && config.maxBodySize > 0
+    ? config.maxBodySize
+    : DEFAULT_MAX_BODY_SIZE;
+
+  console.log('[Agent Shield] Middleware body size limit: %dKB. Configure options.maxBodySize to override.', Math.round(maxBodySize / 1024));
 
   return (req, res, next) => {
+    // Attach raw-stream guard for unparsed requests so attackers cannot
+    // bypass the post-parse size check with huge streamed payloads.
+    attachRawSizeGuard(req, res, maxBodySize);
+
     if (!req.body) {
       req.agentShield = { status: 'safe', threats: [], blocked: false };
       return next();
     }
 
+    // Enforce body-size limit before scanning to avoid DoS via huge inputs.
+    const bodySize = computeBodySize(req.body);
+    if (bodySize > maxBodySize) {
+      return res.status(413).json({
+        error: 'Payload Too Large',
+        message: `Request body (${bodySize} bytes) exceeds maximum allowed size of ${maxBodySize} bytes`,
+        maxBodySize,
+        receivedSize: bodySize
+      });
+    }
+
     // Extract text from common request body shapes
     const text = extractTextFromBody(req.body);
 
@@ -306,4 +402,13 @@ const shieldMiddleware = (config = {}) => {
   };
 };
 
-module.exports = { expressMiddleware, wrapAgent, shieldTools, extractTextFromBody, rateLimitMiddleware, shieldMiddleware };
+module.exports = {
+  expressMiddleware,
+  wrapAgent,
+  shieldTools,
+  extractTextFromBody,
+  rateLimitMiddleware,
+  shieldMiddleware,
+  computeBodySize,
+  DEFAULT_MAX_BODY_SIZE
+};

package/src/native-scanner.js

@@ -0,0 +1,104 @@
+'use strict';
+
+/**
+ * Agent Shield — Native Rust Scanner Bridge
+ *
+ * Provides a transparent bridge to the Rust-core pattern matching engine
+ * compiled via NAPI-RS. When the native module is available, scans run
+ * through Rust's RegexSet for O(n) multi-pattern matching — typically
+ * 5-10x faster than the pure-JS scanner on long inputs.
+ *
+ * Falls back silently to the pure-JS scanner if the native module is
+ * not compiled or unavailable for the current platform.
+ *
+ * Build the native module:
+ *   cd rust-core && cargo build --release --features node
+ *   cp target/release/libagent_shield_core.so agent-shield-core.node  # Linux
+ *   cp target/release/libagent_shield_core.dylib agent-shield-core.node  # macOS
+ *
+ * @module native-scanner
+ */
+
+const path = require('path');
+
+let nativeModule = null;
+let nativeAvailable = false;
+
+const NATIVE_PATHS = [
+  path.join(__dirname, '..', 'rust-core', 'agent-shield-core.node'),
+  path.join(__dirname, '..', 'rust-core', 'target', 'release', 'agent-shield-core.node'),
+  path.join(__dirname, '..', 'native', 'agent-shield-core.node'),
+];
+
+for (const p of NATIVE_PATHS) {
+  try {
+    nativeModule = require(p);
+    nativeAvailable = true;
+    console.log('[Agent Shield] Native Rust scanner loaded from: ' + path.basename(p));
+    break;
+  } catch {
+    // Not available at this path, try next
+  }
+}
+
+/**
+ * Returns true if the native Rust scanner is available.
+ * @returns {boolean}
+ */
+function isNativeAvailable() {
+  return nativeAvailable;
+}
+
+/**
+ * Scan text using the native Rust engine.
+ * Returns null if native is not available (caller should fall back to JS).
+ *
+ * @param {string} text - Text to scan.
+ * @returns {object|null} ScanResult or null if native unavailable.
+ */
+function nativeScan(text) {
+  if (!nativeAvailable || !text || typeof text !== 'string') return null;
+  try {
+    const json = nativeModule.scanText(text);
+    return JSON.parse(json);
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Batch scan multiple texts using the native Rust engine.
+ *
+ * @param {string[]} texts - Array of texts to scan.
+ * @returns {object[]|null} Array of ScanResults or null if native unavailable.
+ */
+function nativeScanBatch(texts) {
+  if (!nativeAvailable || !Array.isArray(texts)) return null;
+  try {
+    const json = nativeModule.scanBatch(texts.filter(t => typeof t === 'string'));
+    return JSON.parse(json);
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Get all patterns from the native Rust engine.
+ *
+ * @returns {object[]|null} Array of patterns or null if native unavailable.
+ */
+function nativeGetPatterns() {
+  if (!nativeAvailable) return null;
+  try {
+    return JSON.parse(nativeModule.getPatterns());
+  } catch {
+    return null;
+  }
+}
+
+module.exports = {
+  isNativeAvailable,
+  nativeScan,
+  nativeScanBatch,
+  nativeGetPatterns,
+};