@kya-os/agentshield-nextjs
Version:
Next.js middleware for AgentShield AI agent detection
310 lines (307 loc) • 10.3 kB
JavaScript
;
var server = require('next/server');
var agentshieldShared = require('@kya-os/agentshield-shared');
// src/edge-wasm-middleware.ts
var wasmExports = null;
var wasmInitPromise = null;
var heap = new Array(128).fill(void 0);
heap.push(void 0, null, true, false);
var heap_next = heap.length;
function addHeapObject(obj) {
if (heap_next === heap.length) heap.push(heap.length + 1);
const idx = heap_next;
heap_next = heap[idx];
heap[idx] = obj;
return idx;
}
function getObject(idx) {
return heap[idx];
}
function dropObject(idx) {
if (idx < 132) return;
heap[idx] = heap_next;
heap_next = idx;
}
function takeObject(idx) {
const ret = getObject(idx);
dropObject(idx);
return ret;
}
var cachedTextDecoder = new TextDecoder("utf-8", { ignoreBOM: true, fatal: true });
var cachedTextEncoder = new TextEncoder();
var cachedUint8ArrayMemory0 = null;
function getUint8ArrayMemory0() {
if (cachedUint8ArrayMemory0 === null || cachedUint8ArrayMemory0.byteLength === 0) {
cachedUint8ArrayMemory0 = new Uint8Array(wasmExports.memory.buffer);
}
return cachedUint8ArrayMemory0;
}
function getStringFromWasm0(ptr, len) {
ptr = ptr >>> 0;
return cachedTextDecoder.decode(getUint8ArrayMemory0().subarray(ptr, ptr + len));
}
var WASM_VECTOR_LEN = 0;
function passStringToWasm0(arg, malloc, realloc) {
if (realloc === void 0) {
const buf = cachedTextEncoder.encode(arg);
const ptr2 = malloc(buf.length, 1) >>> 0;
getUint8ArrayMemory0().subarray(ptr2, ptr2 + buf.length).set(buf);
WASM_VECTOR_LEN = buf.length;
return ptr2;
}
let len = arg.length;
let ptr = malloc(len, 1) >>> 0;
const mem = getUint8ArrayMemory0();
let offset = 0;
for (; offset < len; offset++) {
const code = arg.charCodeAt(offset);
if (code > 127) break;
mem[ptr + offset] = code;
}
if (offset !== len) {
if (offset !== 0) {
arg = arg.slice(offset);
}
ptr = realloc(ptr, len, len = offset + arg.length * 3, 1) >>> 0;
const view = getUint8ArrayMemory0().subarray(ptr + offset, ptr + len);
const ret = cachedTextEncoder.encodeInto(arg, view);
offset += ret.written;
}
WASM_VECTOR_LEN = offset;
return ptr;
}
var cachedDataViewMemory0 = null;
function getDataViewMemory0() {
if (cachedDataViewMemory0 === null || cachedDataViewMemory0.buffer !== wasmExports.memory.buffer) {
cachedDataViewMemory0 = new DataView(wasmExports.memory.buffer);
}
return cachedDataViewMemory0;
}
async function initializeEdgeWasm(wasmModule) {
if (wasmInitPromise) {
return wasmInitPromise;
}
wasmInitPromise = (async () => {
try {
const imports = {
"./agentshield_wasm_bg.js": {
__wbindgen_object_drop_ref: function(arg0) {
dropObject(arg0);
},
__wbindgen_string_new: function(arg0, arg1) {
const ret = getStringFromWasm0(arg0, arg1);
return addHeapObject(ret);
},
__wbindgen_string_get: function(arg0, arg1) {
const obj = getObject(arg1);
const ret = typeof obj === "string" ? obj : void 0;
const ptr1 = ret ? passStringToWasm0(
ret,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
) : 0;
const len1 = WASM_VECTOR_LEN;
getDataViewMemory0().setInt32(arg0 + 4 * 1, len1, true);
getDataViewMemory0().setInt32(arg0 + 4 * 0, ptr1, true);
},
__wbindgen_throw: function(arg0, arg1) {
throw new Error(getStringFromWasm0(arg0, arg1));
}
},
wbg: {
__wbindgen_object_drop_ref: function(arg0) {
dropObject(arg0);
},
__wbindgen_string_new: function(arg0, arg1) {
const ret = getStringFromWasm0(arg0, arg1);
return addHeapObject(ret);
},
__wbindgen_string_get: function(arg0, arg1) {
const obj = getObject(arg1);
const ret = typeof obj === "string" ? obj : void 0;
const ptr1 = ret ? passStringToWasm0(
ret,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
) : 0;
const len1 = WASM_VECTOR_LEN;
getDataViewMemory0().setInt32(arg0 + 4 * 1, len1, true);
getDataViewMemory0().setInt32(arg0 + 4 * 0, ptr1, true);
},
__wbindgen_throw: function(arg0, arg1) {
throw new Error(getStringFromWasm0(arg0, arg1));
}
}
};
const instance = await WebAssembly.instantiate(wasmModule, imports);
wasmExports = instance.exports;
if (wasmExports.__wbindgen_start) {
wasmExports.__wbindgen_start();
}
console.log(
"\u2705 AgentShield: WASM module loaded successfully in Edge Runtime (95-100% confidence enabled)"
);
} catch (error) {
console.error("\u274C AgentShield: Failed to initialize WASM in Edge Runtime:", error);
throw error;
}
})();
return wasmInitPromise;
}
async function detectWithWasm(metadata) {
if (!wasmExports) {
throw new Error("WASM not initialized. Call initializeEdgeWasm first.");
}
try {
const userAgent = metadata.userAgent || "";
const ipAddress = metadata.ipAddress || "";
const headers = JSON.stringify(metadata.headers);
const timestamp = metadata.timestamp;
const url = metadata.url || "";
const method = metadata.method || "";
const clientFingerprint = metadata.clientFingerprint || "";
const ptr0 = passStringToWasm0(
userAgent,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len0 = WASM_VECTOR_LEN;
const ptr1 = passStringToWasm0(
ipAddress,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len1 = WASM_VECTOR_LEN;
const ptr2 = passStringToWasm0(
headers,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len2 = WASM_VECTOR_LEN;
const ptr3 = passStringToWasm0(
timestamp,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len3 = WASM_VECTOR_LEN;
const ptr4 = passStringToWasm0(
url,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len4 = WASM_VECTOR_LEN;
const ptr5 = passStringToWasm0(
method,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len5 = WASM_VECTOR_LEN;
const ptr6 = passStringToWasm0(
clientFingerprint,
wasmExports.__wbindgen_malloc,
wasmExports.__wbindgen_realloc
);
const len6 = WASM_VECTOR_LEN;
const metadataPtr = wasmExports.jsrequestmetadata_new(
ptr0,
len0,
ptr1,
len1,
ptr2,
len2,
ptr3,
len3,
ptr4,
len4,
ptr5,
len5,
ptr6,
len6
);
const resultPtr = wasmExports.detect_agent(metadataPtr);
const result = takeObject(resultPtr);
wasmExports.__wbg_jsrequestmetadata_free(metadataPtr, 0);
const parsedResult = typeof result === "string" ? JSON.parse(result) : result;
return {
isAgent: parsedResult.is_agent || false,
confidence: parsedResult.confidence || 0,
agent: parsedResult.agent,
verificationMethod: parsedResult.verification_method || "wasm",
riskLevel: parsedResult.risk_level || "low",
timestamp: parsedResult.timestamp || metadata.timestamp,
reasons: parsedResult.reasons || []
};
} catch (error) {
console.error("WASM detection failed:", error);
return {
isAgent: false,
confidence: 0,
verificationMethod: "pattern",
riskLevel: "low",
timestamp: metadata.timestamp
};
}
}
function createEdgeWasmMiddleware(config) {
const {
wasmModule,
onAgentDetected,
blockOnHighConfidence = false,
confidenceThreshold: configThreshold = 0.9,
skipPaths = [],
blockedResponse = {
status: 403,
message: "AI agent access restricted",
headers: { "Content-Type": "application/json" }
}
} = config;
const confidenceThreshold = agentshieldShared.normalizeConfidence(configThreshold, "confidenceThreshold");
const initPromise = initializeEdgeWasm(wasmModule);
return async function middleware(request) {
const path = request.nextUrl.pathname;
if (skipPaths.some((skip) => path.startsWith(skip))) {
return server.NextResponse.next();
}
try {
await initPromise;
const metadata = {
userAgent: request.headers.get("user-agent") || void 0,
ipAddress: request.headers.get("x-forwarded-for") || request.headers.get("x-real-ip") || void 0,
headers: Object.fromEntries(request.headers.entries()),
timestamp: (/* @__PURE__ */ new Date()).toISOString()
};
const result = await detectWithWasm(metadata);
if (onAgentDetected && result.isAgent) {
await onAgentDetected(result);
}
if (blockOnHighConfidence && result.isAgent && result.confidence >= confidenceThreshold) {
return server.NextResponse.json(
{
error: blockedResponse.message,
agent: result.agent,
confidence: agentshieldShared.toPercent(result.confidence)
// Convert to 0-100 for display
},
{
status: blockedResponse.status || 403,
headers: blockedResponse.headers || {}
}
);
}
const response = server.NextResponse.next();
if (result.isAgent) {
response.headers.set("X-Agent-Detected", result.agent || "unknown");
response.headers.set("X-Agent-Confidence", String(agentshieldShared.toPercent(result.confidence)));
response.headers.set("X-Agent-Verification", result.verificationMethod);
}
return response;
} catch (error) {
console.error("Edge WASM middleware error:", error);
return server.NextResponse.next();
}
};
}
exports.createEdgeWasmMiddleware = createEdgeWasmMiddleware;
exports.initializeEdgeWasm = initializeEdgeWasm;
//# sourceMappingURL=edge-wasm-middleware.js.map
//# sourceMappingURL=edge-wasm-middleware.js.map