Add JWK parsing limits and align TS implementation with Rust

packages/pq-key-encoder/ts/src/jwk.ts

@@ -3,6 +3,9 @@ import type { AlgorithmName, JwkExportOptions, KeyData, PQJwk } from './types';
 import { decodeBase64Url, encodeBase64Url } from './utils/base64';
 import { assertKeyData, getAlgorithmInfo } from './utils/validation';
 
+const MAX_JSON_SIZE = 65_536;
+const MAX_JSON_FIELDS = 32;
+
 /** Validate that the key type is public or private. */
 function assertKeyType(keyType: KeyData['type']): void {
   if (keyType !== 'public' && keyType !== 'private') {
@@ -39,6 +42,64 @@ function requirePublicKey(algorithm: AlgorithmName, publicKey: unknown): Uint8Ar
   return publicKey;
 }
 
+/**
+ * Extract top-level keys from a JSON object string.
+ * Properly handles string escaping and nesting depth.
+ */
+function extractTopLevelKeys(json: string): string[] {
+  const keys: string[] = [];
+  let i = 0;
+  const len = json.length;
+
+  // Find opening brace
+  while (i < len && json[i] !== '{') i++;
+  if (i >= len) return keys;
+  i++; // skip '{'
+
+  let depth = 0;
+
+  while (i < len) {
+    const ch = json[i];
+
+    if (ch === '"') {
+      // Parse string
+      const start = i + 1;
+      i++; // skip opening quote
+      while (i < len && json[i] !== '"') {
+        if (json[i] === '\\') i++; // skip escaped character
+        i++;
+      }
+      const end = i;
+      i++; // skip closing quote
+
+      // If at top level and next non-ws char is ':', this is a key
+      if (depth === 0) {
+        let j = i;
+        while (j < len && (json[j] === ' ' || json[j] === '\t' || json[j] === '\n' || json[j] === '\r')) j++;
+        if (j < len && json[j] === ':') {
+          const raw = json.slice(start, end);
+          try {
+            keys.push(JSON.parse(`"${raw}"`));
+          } catch {
+            keys.push(raw);
+          }
+        }
+      }
+    } else if (ch === '{' || ch === '[') {
+      depth++;
+      i++;
+    } else if (ch === '}' || ch === ']') {
+      if (depth === 0) break; // end of root object
+      depth--;
+      i++;
+    } else {
+      i++;
+    }
+  }
+
+  return keys;
+}
+
 /** Convert key data to a PQ JWK. */
 export function toJWK(key: KeyData, options: JwkExportOptions = {}): PQJwk {
   assertKeyType(key.type);
@@ -49,31 +110,43 @@ export function toJWK(key: KeyData, options: JwkExportOptions = {}): PQJwk {
       throw new InvalidInputError('includePrivate is not valid for public keys.');
     }
     const encoded = encodeBase64Url(key.bytes);
-    return {
+    const jwk: PQJwk = {
       kty: 'PQC',
       alg: key.alg,
       x: encoded,
     };
+    if (options.kid !== undefined) {
+      jwk.kid = options.kid;
+    }
+    return jwk;
   }
 
   const publicKey = requirePublicKey(key.alg, options.publicKey);
 
   const publicEncoded = encodeBase64Url(publicKey);
   if (!options.includePrivate) {
-    return {
+    const jwk: PQJwk = {
       kty: 'PQC',
       alg: key.alg,
       x: publicEncoded,
     };
+    if (options.kid !== undefined) {
+      jwk.kid = options.kid;
+    }
+    return jwk;
   }
 
   const privateEncoded = encodeBase64Url(key.bytes);
-  return {
+  const jwk: PQJwk = {
     kty: 'PQC',
     alg: key.alg,
     x: publicEncoded,
     d: privateEncoded,
   };
+  if (options.kid !== undefined) {
+    jwk.kid = options.kid;
+  }
+  return jwk;
 }
 
 /** Parse a PQ JWK into raw key data. */
@@ -82,11 +155,12 @@ export function fromJWK(jwk: PQJwk): KeyData {
     throw new InvalidInputError('jwk must be an object.');
   }
 
-  const { kty, alg, x, d } = jwk as {
+  const { kty, alg, x, d, kid } = jwk as {
     kty?: unknown;
     alg?: unknown;
     x?: unknown;
     d?: unknown;
+    kid?: unknown;
   };
 
   if (kty !== 'PQC') {
@@ -101,6 +175,9 @@ export function fromJWK(jwk: PQJwk): KeyData {
   if (d !== undefined && typeof d !== 'string') {
     throw new InvalidInputError('JWK d must be a string when provided.');
   }
+  if (kid !== undefined && typeof kid !== 'string') {
+    throw new InvalidInputError('JWK kid must be a string when provided.');
+  }
 
   const algorithm = parseAlgorithmName(alg);
   const isPrivate = typeof d === 'string';
@@ -122,3 +199,58 @@ export function fromJWK(jwk: PQJwk): KeyData {
     bytes: publicBytes,
   };
 }
+
+/** Serialize key data to a JWK JSON string. */
+export function toJWKString(key: KeyData, options: JwkExportOptions = {}): string {
+  const jwk = toJWK(key, options);
+  return JSON.stringify(jwk);
+}
+
+/**
+ * Parse a JWK JSON string into raw key data.
+ *
+ * Enforces resource limits matching the Rust implementation:
+ * - Input size capped at 64 KiB
+ * - Field count capped at 32
+ * - Duplicate known fields are rejected
+ */
+export function fromJWKString(json: string): KeyData {
+  if (typeof json !== 'string') {
+    throw new InvalidInputError('json must be a string.');
+  }
+
+  const trimmed = json.trim();
+  if (trimmed.length > MAX_JSON_SIZE) {
+    throw new InvalidInputError('JWK input exceeds maximum size.');
+  }
+
+  // Extract top-level keys to enforce field count and duplicate detection
+  const keys = extractTopLevelKeys(trimmed);
+  if (keys.length > MAX_JSON_FIELDS) {
+    throw new InvalidInputError('Too many fields in JWK object.');
+  }
+
+  const knownFields = ['kty', 'alg', 'x', 'd', 'kid'];
+  for (const field of knownFields) {
+    let count = 0;
+    for (const key of keys) {
+      if (key === field) count++;
+    }
+    if (count > 1) {
+      throw new InvalidInputError(`Duplicate '${field}' field.`);
+    }
+  }
+
+  let parsed: unknown;
+  try {
+    parsed = JSON.parse(trimmed);
+  } catch {
+    throw new InvalidInputError('Invalid JSON.');
+  }
+
+  if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
+    throw new InvalidInputError('JWK must be a JSON object.');
+  }
+
+  return fromJWK(parsed as PQJwk);
+}

packages/pq-key-encoder/ts/src/types.ts

@@ -32,4 +32,5 @@ export type PQJwk = PQPublicJwk | PQPrivateJwk;
 export type JwkExportOptions = {
   includePrivate?: boolean;
   publicKey?: Uint8Array;
+  kid?: string;
 };

packages/pq-key-encoder/ts/src/utils/base64.ts

@@ -2,6 +2,12 @@ import { InvalidEncodingError } from '../errors';
 
 const BASE64_PATTERN = /^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/;
 
+const BASE64_VALUES: Record<string, number> = {};
+const BASE64_ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/';
+for (let i = 0; i < BASE64_ALPHABET.length; i += 1) {
+  BASE64_VALUES[BASE64_ALPHABET[i]] = i;
+}
+
 declare const Buffer:
   | {
       from(data: Uint8Array): { toString(encoding: 'base64'): string };
@@ -32,6 +38,21 @@ function binaryStringToBytes(input: string): Uint8Array {
   return bytes;
 }
 
+/** Validate that trailing bits in base64 are zero (RFC 4648 §3.5). */
+function validateTrailingBits(normalized: string): void {
+  const unpadded = normalized.replace(/=+$/, '');
+  const remainder = unpadded.length % 4;
+  if (remainder === 2) {
+    if ((BASE64_VALUES[unpadded[unpadded.length - 1]] & 0x0f) !== 0) {
+      throw new InvalidEncodingError('Non-zero trailing bits in base64.');
+    }
+  } else if (remainder === 3) {
+    if ((BASE64_VALUES[unpadded[unpadded.length - 1]] & 0x03) !== 0) {
+      throw new InvalidEncodingError('Non-zero trailing bits in base64.');
+    }
+  }
+}
+
 /** Normalize base64 input by stripping whitespace and padding. */
 export function normalizeBase64(input: string): string {
   const cleaned = stripWhitespace(input);
@@ -75,6 +96,7 @@ export function decodeBase64(input: string): Uint8Array {
   if (normalized.length === 0) {
     return new Uint8Array();
   }
+  validateTrailingBits(normalized);
   if (typeof globalThis.atob === 'function') {
     return binaryStringToBytes(globalThis.atob(normalized));
   }