ai-shield-middleware 0.1.0

package/package.json

@@ -0,0 +1,43 @@
+{
+  "name": "ai-shield-middleware",
+  "version": "0.1.0",
+  "description": "AI Shield middleware for Express and Hono — route-level LLM protection",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    },
+    "./express": {
+      "types": "./dist/express.d.ts",
+      "import": "./dist/express.js"
+    },
+    "./hono": {
+      "types": "./dist/hono.d.ts",
+      "import": "./dist/hono.js"
+    }
+  },
+  "scripts": {
+    "build": "tsc",
+    "typecheck": "tsc --noEmit"
+  },
+  "peerDependencies": {
+    "express": ">=4.0.0",
+    "hono": ">=4.0.0"
+  },
+  "peerDependenciesMeta": {
+    "express": { "optional": true },
+    "hono": { "optional": true }
+  },
+  "dependencies": {
+    "ai-shield-core": "0.1.0"
+  },
+  "devDependencies": {
+    "@types/express": "^5.0.0",
+    "express": "^4.21.0",
+    "hono": "^4.7.0",
+    "typescript": "^5.7.0"
+  }
+}

package/src/express.ts

@@ -0,0 +1,94 @@
+import type { ScanResult } from "ai-shield-core";
+import type { ShieldMiddlewareConfig } from "./shared.js";
+import { defaultGetInput, scanRequest } from "./shared.js";
+
+// ============================================================
+// Express Middleware — AI Shield route guard
+// ============================================================
+
+// Minimal Express types (avoid hard dependency on @types/express)
+interface ExpressRequest {
+  body?: unknown;
+  path: string;
+  url: string;
+  method: string;
+  headers: Record<string, string | string[] | undefined>;
+}
+
+interface ExpressResponse {
+  status(code: number): ExpressResponse;
+  json(body: unknown): void;
+  locals: Record<string, unknown>;
+}
+
+type NextFunction = (err?: unknown) => void;
+type ExpressMiddleware = (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => void;
+
+/**
+ * Express middleware that scans request body for prompt injection and PII.
+ *
+ * @example
+ * ```ts
+ * import express from "express";
+ * import { shieldMiddleware } from "@ai-shield/middleware/express";
+ *
+ * const app = express();
+ * app.use(express.json());
+ *
+ * // Protect all /api/chat routes
+ * app.use("/api/chat", shieldMiddleware({
+ *   shield: { injection: { strictness: "high" } },
+ *   skipPaths: ["/api/chat/health"],
+ * }));
+ *
+ * // Access scan result in route handler
+ * app.post("/api/chat", (req, res) => {
+ *   const shieldResult = res.locals.shieldResult as ScanResult;
+ *   // shieldResult.sanitized has PII masked
+ * });
+ * ```
+ */
+export function shieldMiddleware(config: ShieldMiddlewareConfig = {}): ExpressMiddleware {
+  return (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
+    // Skip non-mutating methods
+    if (req.method === "GET" || req.method === "HEAD" || req.method === "OPTIONS") {
+      return next();
+    }
+
+    // Skip configured paths
+    if (config.skipPaths?.some((p) => req.path.startsWith(p))) {
+      return next();
+    }
+
+    const getInput = config.getInput ?? defaultGetInput;
+    const input = getInput(req.body);
+
+    // No scannable content — pass through
+    if (!input) {
+      return next();
+    }
+
+    const context = config.getContext?.(req) ?? {};
+    if (config.getAgentId) {
+      context.agentId = config.getAgentId(req);
+    }
+
+    // Async scan
+    void scanRequest(config, input, context)
+      .then(({ blocked, result, response }) => {
+        if (blocked && response) {
+          res.status(response.status).json(response.body);
+          return;
+        }
+
+        // Attach result to res.locals for downstream handlers
+        res.locals.shieldResult = result;
+        next();
+      })
+      .catch((err: unknown) => {
+        next(err);
+      });
+  };
+}
+
+export type { ShieldMiddlewareConfig, ScanResult };

package/src/hono.ts

@@ -0,0 +1,99 @@
+import type { ScanResult } from "ai-shield-core";
+import type { ShieldMiddlewareConfig } from "./shared.js";
+import { defaultGetInput, scanRequest } from "./shared.js";
+
+// ============================================================
+// Hono Middleware — AI Shield route guard
+// ============================================================
+
+// Minimal Hono types (avoid hard dependency)
+interface HonoContext {
+  req: {
+    method: string;
+    path: string;
+    url: string;
+    header(name: string): string | undefined;
+    json(): Promise<unknown>;
+    raw: { headers: Headers };
+  };
+  json(data: unknown, status?: number): Response;
+  set(key: string, value: unknown): void;
+  get(key: string): unknown;
+}
+
+type HonoNext = () => Promise<void>;
+type HonoMiddleware = (c: HonoContext, next: HonoNext) => Promise<Response | void>;
+
+/**
+ * Hono middleware that scans request body for prompt injection and PII.
+ *
+ * @example
+ * ```ts
+ * import { Hono } from "hono";
+ * import { shieldMiddleware } from "@ai-shield/middleware/hono";
+ *
+ * const app = new Hono();
+ *
+ * // Protect chat routes
+ * app.use("/api/chat/*", shieldMiddleware({
+ *   shield: { injection: { strictness: "high" } },
+ * }));
+ *
+ * app.post("/api/chat", async (c) => {
+ *   const shieldResult = c.get("shieldResult") as ScanResult;
+ *   // Use shieldResult.sanitized for PII-masked input
+ * });
+ * ```
+ */
+export function shieldMiddleware(config: ShieldMiddlewareConfig = {}): HonoMiddleware {
+  return async (c: HonoContext, next: HonoNext): Promise<Response | void> => {
+    // Skip non-mutating methods
+    if (c.req.method === "GET" || c.req.method === "HEAD" || c.req.method === "OPTIONS") {
+      return next();
+    }
+
+    // Skip configured paths
+    if (config.skipPaths?.some((p) => c.req.path.startsWith(p))) {
+      return next();
+    }
+
+    // Parse body
+    let body: unknown;
+    try {
+      body = await c.req.json();
+    } catch {
+      // No JSON body — pass through
+      return next();
+    }
+
+    const getInput = config.getInput ?? defaultGetInput;
+    const input = getInput(body);
+
+    if (!input) {
+      return next();
+    }
+
+    // Build context from headers
+    const headers: Record<string, string | string[] | undefined> = {};
+    c.req.raw.headers.forEach((value, key) => {
+      headers[key] = value;
+    });
+
+    const context = config.getContext?.({ headers, body }) ?? {};
+    if (config.getAgentId) {
+      context.agentId = config.getAgentId({ headers, path: c.req.path, url: c.req.url });
+    }
+
+    const { blocked, result, response } = await scanRequest(config, input, context);
+
+    if (blocked && response) {
+      return c.json(response.body, response.status);
+    }
+
+    // Attach result for downstream handlers
+    c.set("shieldResult", result);
+    return next();
+  };
+}
+
+export type { ShieldMiddlewareConfig, ScanResult };

package/src/index.ts

@@ -0,0 +1,13 @@
+// ============================================================
+// @ai-shield/middleware — Public API
+// ============================================================
+
+// Shared types and utilities
+export { defaultGetInput, defaultBlockedResponse } from "./shared.js";
+export type { ShieldMiddlewareConfig } from "./shared.js";
+
+// Framework-specific exports
+export { shieldMiddleware as expressShield } from "./express.js";
+export { shieldMiddleware as honoShield } from "./hono.js";
+
+export type { ScanResult } from "ai-shield-core";

package/src/shared.ts

@@ -0,0 +1,102 @@
+import type { AIShield, ShieldConfig, ScanContext, ScanResult } from "ai-shield-core";
+
+// ============================================================
+// Shared middleware logic — used by Express and Hono adapters
+// ============================================================
+
+export interface ShieldMiddlewareConfig {
+  /** AI Shield config */
+  shield?: ShieldConfig;
+  /** Pre-created AIShield instance (shared across routes) */
+  shieldInstance?: AIShield;
+  /** Extract agent ID from request */
+  getAgentId?: (req: { headers: Record<string, string | string[] | undefined>; path?: string; url?: string }) => string | undefined;
+  /** Extract scan context from request */
+  getContext?: (req: { headers: Record<string, string | string[] | undefined>; body?: unknown }) => ScanContext;
+  /** Extract text to scan from request body */
+  getInput?: (body: unknown) => string | null;
+  /** Custom blocked response */
+  onBlocked?: (result: ScanResult) => { status: number; body: unknown };
+  /** Called on warnings (non-blocking) */
+  onWarning?: (result: ScanResult) => void;
+  /** Skip scanning for certain paths */
+  skipPaths?: string[];
+}
+
+/** Default: extract text from common chat API body shapes */
+export function defaultGetInput(body: unknown): string | null {
+  if (!body || typeof body !== "object") return null;
+
+  const obj = body as Record<string, unknown>;
+
+  // Direct message field
+  if (typeof obj.message === "string") return obj.message;
+  if (typeof obj.input === "string") return obj.input;
+  if (typeof obj.prompt === "string") return obj.prompt;
+  if (typeof obj.text === "string") return obj.text;
+  if (typeof obj.content === "string") return obj.content;
+  if (typeof obj.query === "string") return obj.query;
+
+  // OpenAI-style messages array
+  if (Array.isArray(obj.messages)) {
+    const userMessages = (obj.messages as Array<{ role?: string; content?: string }>)
+      .filter((m) => m.role === "user" && typeof m.content === "string")
+      .map((m) => m.content as string);
+    if (userMessages.length > 0) return userMessages.join("\n");
+  }
+
+  return null;
+}
+
+/** Default blocked response */
+export function defaultBlockedResponse(result: ScanResult): { status: number; body: unknown } {
+  return {
+    status: 403,
+    body: {
+      error: "Request blocked by AI Shield",
+      decision: result.decision,
+      violations: result.violations.map((v) => ({
+        type: v.type,
+        message: v.message,
+      })),
+    },
+  };
+}
+
+/** Lazy-load AIShield instance */
+let _sharedShield: AIShield | null = null;
+let _shieldReady: Promise<AIShield> | null = null;
+
+export async function getOrCreateShield(config: ShieldMiddlewareConfig): Promise<AIShield> {
+  if (config.shieldInstance) return config.shieldInstance;
+  if (_sharedShield) return _sharedShield;
+  if (_shieldReady) return _shieldReady;
+
+  _shieldReady = import("ai-shield-core").then((mod) => {
+    _sharedShield = new mod.AIShield(config.shield ?? {});
+    return _sharedShield;
+  });
+
+  return _shieldReady;
+}
+
+/** Core scan logic shared between Express and Hono */
+export async function scanRequest(
+  config: ShieldMiddlewareConfig,
+  input: string,
+  context: ScanContext,
+): Promise<{ blocked: boolean; result: ScanResult; response?: { status: number; body: unknown } }> {
+  const shield = await getOrCreateShield(config);
+  const result = await shield.scan(input, context);
+
+  if (result.decision === "block") {
+    const onBlocked = config.onBlocked ?? defaultBlockedResponse;
+    return { blocked: true, result, response: onBlocked(result) };
+  }
+
+  if (result.decision === "warn") {
+    config.onWarning?.(result);
+  }
+
+  return { blocked: false, result };
+}

package/tsconfig.json

@@ -0,0 +1,8 @@
+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "outDir": "./dist",
+    "rootDir": "./src"
+  },
+  "include": ["src/**/*"]
+}