@intflows/genkit-guard 0.0.9 → 0.0.11

package/scripts/test-concurrent-vault.js

@@ -0,0 +1,96 @@
+import assert from 'node:assert/strict';
+import { InMemoryPiiVaultStorage } from '../dist/index.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+
+const storage = new InMemoryPiiVaultStorage();
+const userCount = 100;
+
+const users = Array.from({ length: userCount }, (_, index) => ({
+  scopeId: `tenantA:user${index}`,
+  email: `user${index}@example.com`,
+}));
+
+const results = await Promise.all(users.map(async ({ scopeId, email }) => {
+  const tokenizer = new PiiTokenizer({ scopeId, storage });
+  const prompt = `Fetch metadata for the blob named 'sensitive-user-data-${email}.json' inside Azure Blob Storage.`;
+  const masked = await tokenizer.mask(prompt, [{ type: 'EMAIL', value: email }]);
+
+  const simulatedModelToolInput = {
+    blobName: masked.maskedText.match(/'([^']+)'/)?.[1],
+  };
+
+  const toolInput = {
+    blobName: await tokenizer.unmask(simulatedModelToolInput.blobName),
+  };
+
+  const simulatedModelResponse = `The blob belongs to ${masked.maskedText.match(/\[\[EMAIL_[^\]]+\]\]/)?.[0]}.`;
+  const response = await tokenizer.unmask(simulatedModelResponse);
+
+  return {
+    scopeId,
+    email,
+    prompt,
+    maskedText: masked.maskedText,
+    toolInput,
+    response,
+  };
+}));
+
+const placeholders = new Set(results.map((result) => result.maskedText));
+assert.equal(placeholders.size, userCount, 'each concurrent user should get a unique scoped placeholder');
+
+for (const result of results) {
+  assert.equal(result.toolInput.blobName, `sensitive-user-data-${result.email}.json`);
+  assert.equal(result.response, `The blob belongs to ${result.email}.`);
+}
+
+const first = results[0];
+const second = results[1];
+const firstTokenizer = new PiiTokenizer({ scopeId: first.scopeId, storage });
+assert.equal(
+  await firstTokenizer.unmask(second.maskedText),
+  second.maskedText,
+  'one user scope must not resolve another user scope placeholder'
+);
+
+const promptTokenizer = new PiiTokenizer({ scopeId: 'trace:request', storage });
+const promptMasked = await promptTokenizer.mask(
+  "Fetch metadata for the blob named 'sensitive-user-data-john.doe@example.com.json' inside Azure Blob Storage.",
+  [{ type: 'EMAIL', value: 'john.doe@example.com' }]
+);
+
+const toolResponseTokenizer = new PiiTokenizer({ scopeId: 'trace:request', storage });
+await toolResponseTokenizer.mask(
+  'john.doe@example.com',
+  [{ type: 'EMAIL', value: 'john.doe@example.com' }]
+);
+
+const finalModelResponse = `The metadata for the blob '${promptMasked.maskedText.match(/\[\[EMAIL_[^\]]+\]\]/)?.[0]}' has been successfully retrieved.`;
+let unmaskedFinalResponse = finalModelResponse;
+for (const tokenizer of [promptTokenizer, toolResponseTokenizer]) {
+  unmaskedFinalResponse = await tokenizer.unmask(unmaskedFinalResponse);
+}
+
+assert.equal(
+  unmaskedFinalResponse,
+  "The metadata for the blob 'john.doe@example.com' has been successfully retrieved.",
+  'final model output should unmask placeholders created by earlier middleware passes'
+);
+
+const firstMiddlewarePass = new PiiTokenizer({ storage });
+const firstPassMasked = await firstMiddlewarePass.mask(
+  "Fetch metadata for the blob named 'sensitive-user-data-jane.doe@example.com.json' inside Azure Blob Storage.",
+  [{ type: 'EMAIL', value: 'sensitive-user-data-jane.doe@example.com.json' }]
+);
+const laterMiddlewarePass = new PiiTokenizer({ storage });
+const crossContextResponse = `The metadata for the blob '${firstPassMasked.maskedText.match(/\[\[EMAIL_[^\]]+\]\]/)?.[0]}' has been successfully retrieved.`;
+
+await laterMiddlewarePass.importTokens(firstPassMasked.maskedText);
+
+assert.equal(
+  await laterMiddlewarePass.unmask(crossContextResponse),
+  "The metadata for the blob 'sensitive-user-data-jane.doe@example.com.json' has been successfully retrieved.",
+  'a later middleware context should resolve an earlier opaque token from the shared default vault'
+);
+
+console.log(`Concurrent PII vault simulation passed for ${userCount} isolated scopes.`);

package/scripts/test-redis-vault.js

@@ -0,0 +1,153 @@
+import assert from 'node:assert/strict';
+import net from 'node:net';
+import { createRedisPiiVaultStorage } from '../dist/index.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+
+class RespRedisClient {
+  constructor(host = '127.0.0.1', port = 6379) {
+    this.host = host;
+    this.port = port;
+  }
+
+  async hGet(key, field) {
+    return this.command('HGET', key, field);
+  }
+
+  async hSet(key, field, value) {
+    return this.command('HSET', key, field, value);
+  }
+
+  async hGetAll(key) {
+    const values = await this.command('HGETALL', key);
+    const out = {};
+    for (let i = 0; i < values.length; i += 2) {
+      out[values[i]] = values[i + 1];
+    }
+    return out;
+  }
+
+  async expire(key, seconds) {
+    return this.command('EXPIRE', key, String(seconds));
+  }
+
+  async del(key) {
+    return this.command('DEL', key);
+  }
+
+  command(...parts) {
+    return new Promise((resolve, reject) => {
+      const socket = net.createConnection({ host: this.host, port: this.port });
+      let buffer = Buffer.alloc(0);
+
+      socket.setTimeout(3000);
+      socket.once('error', reject);
+      socket.once('timeout', () => {
+        socket.destroy();
+        reject(new Error(`Timed out connecting to Redis at ${this.host}:${this.port}`));
+      });
+      socket.once('connect', () => {
+        socket.write(encodeCommand(parts));
+      });
+      socket.on('data', (chunk) => {
+        buffer = Buffer.concat([buffer, chunk]);
+        try {
+          const [value, offset] = parseResp(buffer, 0);
+          if (offset <= buffer.length) {
+            socket.end();
+            resolve(value);
+          }
+        } catch (error) {
+          if (error.message !== 'Incomplete RESP response') {
+            socket.destroy();
+            reject(error);
+          }
+        }
+      });
+    });
+  }
+}
+
+const redis = new RespRedisClient(process.env.REDIS_HOST ?? '127.0.0.1', Number(process.env.REDIS_PORT ?? 6379));
+const keyPrefix = `genkit-guard:test:${Date.now()}`;
+const storage = createRedisPiiVaultStorage(redis, { keyPrefix, ttlSeconds: 60 });
+
+try {
+  const userA = new PiiTokenizer({ scopeId: 'tenant:userA', storage });
+  const userB = new PiiTokenizer({ scopeId: 'tenant:userB', storage });
+
+  const maskedA = await userA.mask('Email alice@example.com', [
+    { type: 'EMAIL', value: 'alice@example.com' },
+  ]);
+  const maskedB = await userB.mask('Email bob@example.com', [
+    { type: 'EMAIL', value: 'bob@example.com' },
+  ]);
+
+  assert.equal(await userA.unmask(maskedA.maskedText), 'Email alice@example.com');
+  assert.equal(await userB.unmask(maskedB.maskedText), 'Email bob@example.com');
+  assert.equal(await userA.unmask(maskedB.maskedText), maskedB.maskedText);
+
+  const laterPass = new PiiTokenizer({ scopeId: 'tenant:userA:later', storage });
+  await laterPass.importTokens(maskedA.maskedText);
+  assert.equal(await laterPass.unmask(maskedA.maskedText), 'Email alice@example.com');
+
+  console.log('Live Redis PII vault test passed.');
+} finally {
+  await redis.del(`${keyPrefix}:tokens`).catch(() => {});
+  await redis.del(`${keyPrefix}:scope:tenant:userA`).catch(() => {});
+  await redis.del(`${keyPrefix}:scope:tenant:userB`).catch(() => {});
+  await redis.del(`${keyPrefix}:scope:tenant:userA:later`).catch(() => {});
+}
+
+function encodeCommand(parts) {
+  return `*${parts.length}\r\n${parts.map((part) => {
+    const value = String(part);
+    return `$${Buffer.byteLength(value)}\r\n${value}\r\n`;
+  }).join('')}`;
+}
+
+function parseResp(buffer, offset) {
+  if (offset >= buffer.length) {
+    throw new Error('Incomplete RESP response');
+  }
+
+  const prefix = String.fromCharCode(buffer[offset]);
+  if (prefix === '+') return parseSimple(buffer, offset);
+  if (prefix === '-') {
+    const [message] = parseSimple(buffer, offset);
+    throw new Error(message);
+  }
+  if (prefix === ':') {
+    const [value, next] = parseSimple(buffer, offset);
+    return [Number(value), next];
+  }
+  if (prefix === '$') return parseBulk(buffer, offset);
+  if (prefix === '*') return parseArray(buffer, offset);
+  throw new Error(`Unsupported RESP prefix: ${prefix}`);
+}
+
+function parseSimple(buffer, offset) {
+  const end = buffer.indexOf('\r\n', offset);
+  if (end === -1) throw new Error('Incomplete RESP response');
+  return [buffer.toString('utf8', offset + 1, end), end + 2];
+}
+
+function parseBulk(buffer, offset) {
+  const [lengthText, valueStart] = parseSimple(buffer, offset);
+  const length = Number(lengthText);
+  if (length === -1) return [null, valueStart];
+  const valueEnd = valueStart + length;
+  if (buffer.length < valueEnd + 2) throw new Error('Incomplete RESP response');
+  return [buffer.toString('utf8', valueStart, valueEnd), valueEnd + 2];
+}
+
+function parseArray(buffer, offset) {
+  const [lengthText, start] = parseSimple(buffer, offset);
+  const values = [];
+  let next = start;
+  for (let i = 0; i < Number(lengthText); i++) {
+    const [value, newOffset] = parseResp(buffer, next);
+    values.push(value);
+    next = newOffset;
+  }
+  return [values, next];
+}

package/scripts/test-storage.js

@@ -0,0 +1,78 @@
+import assert from 'node:assert/strict';
+import {
+  createPiiVaultStorage,
+  createRedisPiiVaultStorage,
+  InMemoryPiiVaultStorage,
+} from '../dist/index.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+
+class FakeRedis {
+  hashes = new Map();
+  expirations = new Map();
+
+  async hGet(key, field) {
+    return this.hashes.get(key)?.[field] ?? null;
+  }
+
+  async hSet(key, field, value) {
+    const hash = this.hashes.get(key) ?? {};
+    hash[field] = value;
+    this.hashes.set(key, hash);
+  }
+
+  async hGetAll(key) {
+    return this.hashes.get(key) ?? {};
+  }
+
+  async expire(key, seconds) {
+    this.expirations.set(key, seconds);
+  }
+}
+
+const memory = new InMemoryPiiVaultStorage();
+await memory.set('scope-a', '[[EMAIL_token_0]]', 'a@example.com');
+assert.equal(await memory.get('scope-a', '[[EMAIL_token_0]]'), 'a@example.com');
+assert.equal(await memory.get('scope-b', '[[EMAIL_token_0]]'), undefined);
+assert.equal(await memory.getByToken('[[EMAIL_token_0]]'), 'a@example.com');
+
+const custom = createPiiVaultStorage({
+  async get(scopeId, token) {
+    return scopeId === 'custom' && token === '[[EMAIL_custom_0]]' ? 'custom@example.com' : undefined;
+  },
+  async set() {},
+  async entries(scopeId) {
+    return scopeId === 'custom'
+      ? [{ token: '[[EMAIL_custom_0]]', value: 'custom@example.com' }]
+      : [];
+  },
+});
+
+const customTokenizer = new PiiTokenizer({ scopeId: 'custom', storage: custom });
+assert.equal(await customTokenizer.unmask('Hi [[EMAIL_custom_0]]'), 'Hi custom@example.com');
+
+const redis = new FakeRedis();
+const redisStorage = createRedisPiiVaultStorage(redis, {
+  keyPrefix: 'test:pii',
+  ttlSeconds: 60,
+});
+
+const userATokenizer = new PiiTokenizer({ scopeId: 'tenant:userA', storage: redisStorage });
+const userBTokenizer = new PiiTokenizer({ scopeId: 'tenant:userB', storage: redisStorage });
+
+const userAMasked = await userATokenizer.mask('Email alice@example.com', [
+  { type: 'EMAIL', value: 'alice@example.com' },
+]);
+const userBMasked = await userBTokenizer.mask('Email bob@example.com', [
+  { type: 'EMAIL', value: 'bob@example.com' },
+]);
+
+assert.notEqual(userAMasked.maskedText, userBMasked.maskedText);
+assert.equal(await userATokenizer.unmask(userAMasked.maskedText), 'Email alice@example.com');
+assert.equal(await userBTokenizer.unmask(userBMasked.maskedText), 'Email bob@example.com');
+assert.equal(await userATokenizer.unmask(userBMasked.maskedText), userBMasked.maskedText);
+
+const laterPassTokenizer = new PiiTokenizer({ scopeId: 'tenant:userA:later', storage: redisStorage });
+await laterPassTokenizer.importTokens(userAMasked.maskedText);
+assert.equal(await laterPassTokenizer.unmask(userAMasked.maskedText), 'Email alice@example.com');
+
+console.log('PII vault storage tests passed.');

package/scripts/test-types.ts

@@ -0,0 +1,21 @@
+import { guard, createRedisPiiVaultStorage, type RedisPiiVaultClient } from '../src/index.js';
+
+const redis: RedisPiiVaultClient = {
+  async hGet() {
+    return undefined;
+  },
+  async hSet() {},
+  async hGetAll() {
+    return {};
+  },
+};
+
+guard({
+  pii: {
+    reversible: true,
+    vault: {
+      storage: createRedisPiiVaultStorage(redis),
+      scopeId: (req: any, ctx: any) => ctx?.auth?.sessionId ?? req?.metadata?.requestId,
+    },
+  },
+});