feat: Add JWT string validator (#3893)

* feat: add JWT string validator - Add z.string().jwt() validator for checking JWT format - Add optional algorithm validation with z.string().jwt({ alg: string }) - Implement in both main and Deno versions - Add comprehensive test coverage - Use atob() for browser compatibility * fix: rename algorithm to alg --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
colinhacks · Dec 10, 2024 · b68c05f · b68c05f
1 parent c1dd537
commit b68c05f
Show file tree

Hide file tree

Showing 6 changed files with 148 additions and 0 deletions.
diff --git a/deno/lib/ZodError.ts b/deno/lib/ZodError.ts
@@ -105,6 +105,7 @@ export type StringValidation =
   | "ip"
   | "cidr"
   | "base64"
+  | "jwt"
   | "base64url"
   | { includes: string; position?: number }
   | { startsWith: string }

diff --git a/deno/lib/__tests__/string.test.ts b/deno/lib/__tests__/string.test.ts
@@ -247,6 +247,44 @@ for (const str of invalidBase64URLStrings) {
   });
 }
 
+test("jwt validations", () => {
+  const jwt = z.string().jwt();
+  const jwtWithAlg = z.string().jwt({ alg: "HS256" });
+
+  // Valid JWTs
+  const validHeader = Buffer.from(JSON.stringify({ typ: "JWT", alg: "HS256" })).toString('base64url');
+  const validPayload = Buffer.from("{}").toString('base64url');
+  const validSignature = "signature";
+  const validJWT = `${validHeader}.${validPayload}.${validSignature}`;
+
+  expect(() => jwt.parse(validJWT)).not.toThrow();
+  expect(() => jwtWithAlg.parse(validJWT)).not.toThrow();
+
+  // Invalid format
+  expect(() => jwt.parse("invalid")).toThrow();
+  expect(() => jwt.parse("invalid.invalid")).toThrow();
+  expect(() => jwt.parse("invalid.invalid.invalid")).toThrow();
+
+  // Invalid header
+  const invalidHeader = Buffer.from("{}").toString('base64url');
+  const invalidHeaderJWT = `${invalidHeader}.${validPayload}.${validSignature}`;
+  expect(() => jwt.parse(invalidHeaderJWT)).toThrow();
+
+  // Wrong algorithm
+  const wrongAlgHeader = Buffer.from(JSON.stringify({ typ: "JWT", alg: "RS256" })).toString('base64url');
+  const wrongAlgJWT = `${wrongAlgHeader}.${validPayload}.${validSignature}`;
+  expect(() => jwtWithAlg.parse(wrongAlgJWT)).toThrow();
+
+  // Custom error message
+  const customMsg = "Invalid JWT token";
+  const jwtWithMsg = z.string().jwt({ message: customMsg });
+  try {
+    jwtWithMsg.parse("invalid");
+  } catch (error) {
+    expect((error as z.ZodError).issues[0].message).toBe(customMsg);
+  }
+});
+
 test("url validations", () => {
   const url = z.string().url();
   url.parse("http://google.com");

diff --git a/deno/lib/types.ts b/deno/lib/types.ts
@@ -604,6 +604,7 @@ export type ZodStringCheck =
   | { kind: "trim"; message?: string }
   | { kind: "toLowerCase"; message?: string }
   | { kind: "toUpperCase"; message?: string }
+  | { kind: "jwt"; alg?: string; message?: string }
   | {
       kind: "datetime";
       offset: boolean;
@@ -641,6 +642,7 @@ const ulidRegex = /^[0-9A-HJKMNP-TV-Z]{26}$/i;
 const uuidRegex =
   /^[0-9a-fA-F]{8}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{12}$/i;
 const nanoidRegex = /^[a-z0-9_-]{21}$/i;
+const jwtRegex = /^[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+\.[A-Za-z0-9-_]*$/;
 const durationRegex =
   /^[-+]?P(?!$)(?:(?:[-+]?\d+Y)|(?:[-+]?\d+[.,]\d+Y$))?(?:(?:[-+]?\d+M)|(?:[-+]?\d+[.,]\d+M$))?(?:(?:[-+]?\d+W)|(?:[-+]?\d+[.,]\d+W$))?(?:(?:[-+]?\d+D)|(?:[-+]?\d+[.,]\d+D$))?(?:T(?=[\d+-])(?:(?:[-+]?\d+H)|(?:[-+]?\d+[.,]\d+H$))?(?:(?:[-+]?\d+M)|(?:[-+]?\d+[.,]\d+M$))?(?:[-+]?\d+(?:[.,]\d+)?S)?)??$/;
 
@@ -739,6 +741,25 @@ function isValidIP(ip: string, version?: IpVersion) {
   return false;
 }
 
+function isValidJWT(jwt: string, alg?: string): boolean {
+  if (!jwtRegex.test(jwt)) return false;
+  try {
+    const [header] = jwt.split(".");
+    // Convert base64url to base64
+    const base64 = header
+      .replace(/-/g, "+")
+      .replace(/_/g, "/")
+      .padEnd(header.length + ((4 - (header.length % 4)) % 4), "=");
+    const decoded = JSON.parse(atob(base64));
+    if (typeof decoded !== "object" || decoded === null) return false;
+    if (!decoded.typ || !decoded.alg) return false;
+    if (alg && decoded.alg !== alg) return false;
+    return true;
+  } catch {
+    return false;
+  }
+}
+
 function isValidCidr(ip: string, version?: IpVersion) {
   if ((version === "v4" || !version) && ipv4CidrRegex.test(ip)) {
     return true;
@@ -1012,6 +1033,16 @@ export class ZodString extends ZodType<string, ZodStringDef, string> {
           });
           status.dirty();
         }
+      } else if (check.kind === "jwt") {
+        if (!isValidJWT(input.data, check.alg)) {
+          ctx = this._getOrReturnCtx(input, ctx);
+          addIssueToContext(ctx, {
+            validation: "jwt",
+            code: ZodIssueCode.invalid_string,
+            message: check.message,
+          });
+          status.dirty();
+        }
       } else if (check.kind === "cidr") {
         if (!isValidCidr(input.data, check.version)) {
           ctx = this._getOrReturnCtx(input, ctx);
@@ -1105,6 +1136,10 @@ export class ZodString extends ZodType<string, ZodStringDef, string> {
     return this._addCheck({ kind: "base64url", ...errorUtil.errToObj(message) });
   }
 
+  jwt(options?: { alg?: string; message?: string }) {
+    return this._addCheck({ kind: "jwt", ...errorUtil.errToObj(options) });
+  }
+
   ip(options?: string | { version?: IpVersion; message?: string }) {
     return this._addCheck({ kind: "ip", ...errorUtil.errToObj(options) });
   }

diff --git a/src/ZodError.ts b/src/ZodError.ts
@@ -105,6 +105,7 @@ export type StringValidation =
   | "ip"
   | "cidr"
   | "base64"
+  | "jwt"
   | "base64url"
   | { includes: string; position?: number }
   | { startsWith: string }

diff --git a/src/__tests__/string.test.ts b/src/__tests__/string.test.ts
@@ -246,6 +246,44 @@ for (const str of invalidBase64URLStrings) {
   });
 }
 
+test("jwt validations", () => {
+  const jwt = z.string().jwt();
+  const jwtWithAlg = z.string().jwt({ alg: "HS256" });
+
+  // Valid JWTs
+  const validHeader = Buffer.from(JSON.stringify({ typ: "JWT", alg: "HS256" })).toString('base64url');
+  const validPayload = Buffer.from("{}").toString('base64url');
+  const validSignature = "signature";
+  const validJWT = `${validHeader}.${validPayload}.${validSignature}`;
+
+  expect(() => jwt.parse(validJWT)).not.toThrow();
+  expect(() => jwtWithAlg.parse(validJWT)).not.toThrow();
+
+  // Invalid format
+  expect(() => jwt.parse("invalid")).toThrow();
+  expect(() => jwt.parse("invalid.invalid")).toThrow();
+  expect(() => jwt.parse("invalid.invalid.invalid")).toThrow();
+
+  // Invalid header
+  const invalidHeader = Buffer.from("{}").toString('base64url');
+  const invalidHeaderJWT = `${invalidHeader}.${validPayload}.${validSignature}`;
+  expect(() => jwt.parse(invalidHeaderJWT)).toThrow();
+
+  // Wrong algorithm
+  const wrongAlgHeader = Buffer.from(JSON.stringify({ typ: "JWT", alg: "RS256" })).toString('base64url');
+  const wrongAlgJWT = `${wrongAlgHeader}.${validPayload}.${validSignature}`;
+  expect(() => jwtWithAlg.parse(wrongAlgJWT)).toThrow();
+
+  // Custom error message
+  const customMsg = "Invalid JWT token";
+  const jwtWithMsg = z.string().jwt({ message: customMsg });
+  try {
+    jwtWithMsg.parse("invalid");
+  } catch (error) {
+    expect((error as z.ZodError).issues[0].message).toBe(customMsg);
+  }
+});
+
 test("url validations", () => {
   const url = z.string().url();
   url.parse("http://google.com");

diff --git a/src/types.ts b/src/types.ts
@@ -604,6 +604,7 @@ export type ZodStringCheck =
   | { kind: "trim"; message?: string }
   | { kind: "toLowerCase"; message?: string }
   | { kind: "toUpperCase"; message?: string }
+  | { kind: "jwt"; alg?: string; message?: string }
   | {
       kind: "datetime";
       offset: boolean;
@@ -641,6 +642,7 @@ const ulidRegex = /^[0-9A-HJKMNP-TV-Z]{26}$/i;
 const uuidRegex =
   /^[0-9a-fA-F]{8}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{12}$/i;
 const nanoidRegex = /^[a-z0-9_-]{21}$/i;
+const jwtRegex = /^[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+\.[A-Za-z0-9-_]*$/;
 const durationRegex =
   /^[-+]?P(?!$)(?:(?:[-+]?\d+Y)|(?:[-+]?\d+[.,]\d+Y$))?(?:(?:[-+]?\d+M)|(?:[-+]?\d+[.,]\d+M$))?(?:(?:[-+]?\d+W)|(?:[-+]?\d+[.,]\d+W$))?(?:(?:[-+]?\d+D)|(?:[-+]?\d+[.,]\d+D$))?(?:T(?=[\d+-])(?:(?:[-+]?\d+H)|(?:[-+]?\d+[.,]\d+H$))?(?:(?:[-+]?\d+M)|(?:[-+]?\d+[.,]\d+M$))?(?:[-+]?\d+(?:[.,]\d+)?S)?)??$/;
 
@@ -739,6 +741,25 @@ function isValidIP(ip: string, version?: IpVersion) {
   return false;
 }
 
+function isValidJWT(jwt: string, alg?: string): boolean {
+  if (!jwtRegex.test(jwt)) return false;
+  try {
+    const [header] = jwt.split(".");
+    // Convert base64url to base64
+    const base64 = header
+      .replace(/-/g, "+")
+      .replace(/_/g, "/")
+      .padEnd(header.length + ((4 - (header.length % 4)) % 4), "=");
+    const decoded = JSON.parse(atob(base64));
+    if (typeof decoded !== "object" || decoded === null) return false;
+    if (!decoded.typ || !decoded.alg) return false;
+    if (alg && decoded.alg !== alg) return false;
+    return true;
+  } catch {
+    return false;
+  }
+}
+
 function isValidCidr(ip: string, version?: IpVersion) {
   if ((version === "v4" || !version) && ipv4CidrRegex.test(ip)) {
     return true;
@@ -1012,6 +1033,16 @@ export class ZodString extends ZodType<string, ZodStringDef, string> {
           });
           status.dirty();
         }
+      } else if (check.kind === "jwt") {
+        if (!isValidJWT(input.data, check.alg)) {
+          ctx = this._getOrReturnCtx(input, ctx);
+          addIssueToContext(ctx, {
+            validation: "jwt",
+            code: ZodIssueCode.invalid_string,
+            message: check.message,
+          });
+          status.dirty();
+        }
       } else if (check.kind === "cidr") {
         if (!isValidCidr(input.data, check.version)) {
           ctx = this._getOrReturnCtx(input, ctx);
@@ -1105,6 +1136,10 @@ export class ZodString extends ZodType<string, ZodStringDef, string> {
     return this._addCheck({ kind: "base64url", ...errorUtil.errToObj(message) });
   }
 
+  jwt(options?: { alg?: string; message?: string }) {
+    return this._addCheck({ kind: "jwt", ...errorUtil.errToObj(options) });
+  }
+
   ip(options?: string | { version?: IpVersion; message?: string }) {
     return this._addCheck({ kind: "ip", ...errorUtil.errToObj(options) });
   }