[zkLogin] Check various lengths before address derivation (MystenLabs…

…#13992)

## Description 

Describe the changes or additions included in this PR.

## Test Plan 

How did you test the new or updated feature?

---
If your changes are not user-facing and not a breaking change, you can
skip the following section. Otherwise, please indicate what changed, and
then add to the Release Notes section as highlighted during the release
process.

### Type of Change (Check all that apply)

- [ ] protocol change
- [ ] user-visible impact
- [ ] breaking change for a client SDKs
- [ ] breaking change for FNs (FN binary must upgrade)
- [ ] breaking change for validators or node operators (must upgrade
binaries)
- [ ] breaking change for on-chain data layout
- [ ] necessitate either a data wipe or data migration

### Release notes

.changeset/green-baboons-itch.md

@@ -0,0 +1,5 @@
+---
+'@mysten/zklogin': patch
+---
+
+Introduce precise key-value pair parsing that matches the circuit

sdk/zklogin/package.json

@@ -29,7 +29,8 @@
 		"eslint:check": "eslint --max-warnings=0 .",
 		"eslint:fix": "pnpm run eslint:check --fix",
 		"lint": "pnpm run eslint:check && pnpm run prettier:check",
-		"lint:fix": "pnpm run eslint:fix && pnpm run prettier:fix"
+		"lint:fix": "pnpm run eslint:fix && pnpm run prettier:fix",
+		"test": "vitest"
 	},
 	"repository": {
 		"type": "git",
@@ -42,13 +43,15 @@
 	"devDependencies": {
 		"@mysten/build-scripts": "workspace:*",
 		"@types/node": "^20.4.2",
-		"typescript": "^5.1.6"
+		"typescript": "^5.1.6",
+		"vitest": "^0.33.0"
 	},
 	"dependencies": {
 		"@mysten/bcs": "workspace:*",
 		"@mysten/sui.js": "workspace:*",
 		"@noble/hashes": "^1.3.1",
 		"jose": "^4.14.4",
+		"jsonc-parser": "^3.2.0",
 		"poseidon-lite": "^0.2.0",
 		"tsx": "^3.12.7"
 	}

sdk/zklogin/src/address.ts

@@ -2,25 +2,38 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { computeZkLoginAddressFromSeed } from '@mysten/sui.js/zklogin';
-import { decodeJwt } from 'jose';
+import { base64url, decodeJwt } from 'jose';
 
+import { lengthChecks } from './checks';
+import { JSONProcessor } from './jsonprocessor.js';
 import { genAddressSeed } from './utils.js';
 
 export function jwtToAddress(jwt: string, userSalt: bigint) {
 	const decodedJWT = decodeJwt(jwt);
-	if (!decodedJWT.sub || !decodedJWT.iss || !decodedJWT.aud) {
-		throw new Error('Missing jwt data');
+	if (!decodedJWT.iss) {
+		throw new Error('Missing iss');
 	}
 
-	if (Array.isArray(decodedJWT.aud)) {
-		throw new Error('Not supported aud. Aud is an array, string was expected.');
+	const keyClaimName = 'sub';
+	const [header, payload] = jwt.split('.');
+	const decoded_payload = base64url.decode(payload).toString();
+	const processor = new JSONProcessor(decoded_payload);
+	const keyClaimDetails = processor.process(keyClaimName); // throws an error if key claim name is not found
+	if (typeof keyClaimDetails.value !== 'string') {
+		throw new Error('Key claim value must be a string');
 	}
+	const audDetails = processor.process('aud');
+	if (typeof audDetails.value !== 'string') {
+		throw new Error('Aud claim value must be a string');
+	}
+
+	lengthChecks(header, payload, keyClaimName, processor);
 
 	return computeZkLoginAddress({
 		userSalt,
-		claimName: 'sub',
-		claimValue: decodedJWT.sub,
-		aud: decodedJWT.aud,
+		claimName: keyClaimName,
+		claimValue: keyClaimDetails.value,
+		aud: audDetails.value,
 		iss: decodedJWT.iss,
 	});
 }

sdk/zklogin/src/checks.ts

@@ -0,0 +1,103 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+import { JSONProcessor } from './jsonprocessor.js';
+import { MAX_AUD_VALUE_LENGTH, MAX_KEY_CLAIM_VALUE_LENGTH } from './utils.js';
+
+const MAX_HEADER_LEN_B64 = 248;
+const MAX_PADDED_UNSIGNED_JWT_LEN = 64 * 25;
+const MAX_EXTENDED_KEY_CLAIM_LEN = 126;
+const MAX_EXTENDED_EV_LEN = 53;
+const MAX_EXTENDED_NONCE_LEN = 44;
+const MAX_EXTENDED_AUD_LEN = 160;
+const MAX_EXTENDED_ISS_LEN_B64 = 224;
+
+export function lengthChecks(
+	header: string,
+	payload: string,
+	keyClaimName: string,
+	processor: JSONProcessor,
+) {
+	/// Is the header length small enough
+	const header_len = header.length;
+	if (header_len > MAX_HEADER_LEN_B64) {
+		throw new Error(`Header is too long`);
+	}
+
+	/// Is the combined length of the header and payload small enough
+	const unsigned_jwt = header + '.' + payload;
+	const L = unsigned_jwt.length * 8;
+	const K = (512 + 448 - ((L % 512) + 1)) % 512;
+	if ((L + 1 + K + 64) % 512 !== 0) {
+		throw new Error('This should never happen');
+	}
+
+	// The SHA2 padding is 1 followed by K zeros, followed by the length of the message
+	const padded_unsigned_jwt_len = (L + 1 + K + 64) / 8;
+
+	// The padded unsigned JWT must be less than the max_padded_unsigned_jwt_len
+	if (padded_unsigned_jwt_len > MAX_PADDED_UNSIGNED_JWT_LEN) {
+		throw new Error(`The JWT is too long`);
+	}
+
+	const keyClaimDetails = processor.process(keyClaimName); // throws an error if key claim name is not found
+	const keyClaimValue = processor.getRawClaimValue(keyClaimName);
+	const keyClaimValueLen = keyClaimValue.length;
+	if (keyClaimValueLen > MAX_KEY_CLAIM_VALUE_LENGTH) {
+		throw new Error('Key claim value is too long');
+	}
+	// Note: Key claim name length is being checked in genAddressSeed.
+
+	/// Are the extended claims small enough (key claim, email_verified)
+	const extendedKeyClaimLen = keyClaimDetails.ext_claim.length;
+	if (extendedKeyClaimLen > MAX_EXTENDED_KEY_CLAIM_LEN) {
+		throw new Error(`Extended key claim length is too long`);
+	}
+
+	if (keyClaimName === 'email') {
+		const evClaimDetails = processor.process('email_verified');
+		const value = evClaimDetails.value;
+		if (!(value === true || value === 'true')) {
+			throw new Error(`Unexpected email_verified claim value ${value}`);
+		}
+		const extEVClaimLen = evClaimDetails.ext_claim.length;
+		if (extEVClaimLen > MAX_EXTENDED_EV_LEN) {
+			throw new Error('Extended email_verified claim length is too long');
+		}
+	}
+
+	/// Check that nonce extended nonce length is as expected.
+	const nonce_claim_details = processor.process('nonce');
+	const nonce_value_len = nonce_claim_details.offsets.value_length;
+	if (nonce_value_len !== 27) {
+		throw new Error(`Nonce value length is not 27`);
+	}
+	const extended_nonce_claim_len = nonce_claim_details.ext_claim.length;
+	if (extended_nonce_claim_len < 38) {
+		throw new Error(`Extended nonce claim is too short`);
+	}
+	if (extended_nonce_claim_len > MAX_EXTENDED_NONCE_LEN) {
+		throw new Error('Extended nonce claim is too long');
+	}
+
+	/// 5. Check if aud value is small enough
+	const aud_claim_details = processor.process('aud');
+	const aud_value = processor.getRawClaimValue('aud');
+	const aud_value_len = aud_value.length;
+	if (aud_value_len > MAX_AUD_VALUE_LENGTH) {
+		throw new Error(`aud is too long`);
+	}
+
+	const extended_aud_claim_len = aud_claim_details.ext_claim.length;
+	if (extended_aud_claim_len > MAX_EXTENDED_AUD_LEN) {
+		throw new Error(`Extended aud is too long`);
+	}
+
+	/// 6. Check if iss is small enough
+	const iss_claim_details = processor.process('iss');
+	// A close upper bound of the length of the extended iss claim (in base64)
+	const iss_claim_len_b64 = 4 * (1 + Math.floor(iss_claim_details.offsets.ext_length / 3));
+	if (iss_claim_len_b64 > MAX_EXTENDED_ISS_LEN_B64) {
+		throw new Error(`Extended iss is too long`);
+	}
+}

sdk/zklogin/src/jsonprocessor.ts

@@ -0,0 +1,181 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+import { JSONVisitor, ParseError, ParseErrorCode, visit } from 'jsonc-parser';
+
+// JSON parsing code inspired from https://github.com/microsoft/node-jsonc-parser/blob/main/src/test/json.test.ts#L69
+interface VisitorCallback {
+	id: keyof JSONVisitor;
+	offset: number;
+	length: number;
+	// Not expecting any claim that is not a string or a boolean (email_verified is sometimes a boolean).
+	// Ensuring that key claim and aud are strings is done in getRawClaimValue
+	arg?: string | boolean | number;
+}
+
+interface VisitorError extends ParseError {
+	startLine: number;
+	startCharacter: number;
+}
+
+export interface ClaimDetails {
+	name: string; // e.g., "sub"
+	// Not expecting any claim that is not a string or a boolean (boolean for email_verified)...
+	value: string | boolean | number; // e.g., "1234567890"
+	ext_claim: string; // e.g., "sub": "1234567890",
+	offsets: {
+		start: number; // start index
+		colon: number; // index of the colon (within the ext_claim)
+		value: number; // index of the value (within the ext_claim)
+		value_length: number; // length of the value
+		name_length: number; // length of the name
+		ext_length: number; // ext_claim.length
+	};
+}
+
+export class JSONProcessor {
+	decoded_payload: string;
+	processed: Record<string, ClaimDetails>;
+	events: VisitorCallback[];
+
+	constructor(decoded_payload: string) {
+		this.decoded_payload = decoded_payload;
+		this.events = this.visit();
+		this.processed = {};
+	}
+
+	visit(): VisitorCallback[] {
+		const errors: VisitorError[] = [];
+		const actuals: VisitorCallback[] = [];
+		const noArgHolder = (id: keyof JSONVisitor) => (offset: number, length: number) =>
+			actuals.push({ id, offset, length });
+		const oneArgHolder =
+			(id: keyof JSONVisitor) => (arg: string | boolean, offset: number, length: number) =>
+				actuals.push({ id, offset, length, arg });
+		visit(
+			this.decoded_payload,
+			{
+				onObjectBegin: noArgHolder('onObjectBegin'),
+				onObjectProperty: oneArgHolder('onObjectProperty'),
+				onObjectEnd: noArgHolder('onObjectEnd'),
+				onLiteralValue: oneArgHolder('onLiteralValue'),
+				onSeparator: oneArgHolder('onSeparator'), // triggers on both : and ,
+				onArrayBegin: noArgHolder('onArrayBegin'),
+				// Of all the events, the ones that we do not listen to are
+				//  onArrayEnd (as onArrayBegin allows us to throw errors if arrays are seen)
+				//  and onComment (as we disallow comments anyway)
+				onError: (
+					error: ParseErrorCode,
+					offset: number,
+					length: number,
+					startLine: number,
+					startCharacter: number,
+				) => {
+					errors.push({ error, offset, length, startLine, startCharacter });
+				},
+			},
+			{
+				disallowComments: true,
+			},
+		);
+		if (errors.length > 0) {
+			console.error(JSON.stringify(errors));
+			throw new Error(`Parse errors encountered`);
+		}
+		return actuals;
+	}
+
+	process(name: string): ClaimDetails {
+		if (Object.prototype.hasOwnProperty.call(this.processed, name)) {
+			return this.processed[name];
+		}
+
+		const name_event_idx = this.events.findIndex(
+			(e) => e.id === 'onObjectProperty' && e.arg === name,
+		);
+		if (name_event_idx === -1) {
+			throw new Error('Claim ' + name + ' not found');
+		}
+
+		const name_event = this.events[name_event_idx];
+
+		const colon_event_idx = name_event_idx + 1;
+		const colon_event = this.events[colon_event_idx];
+		if (
+			this.events[colon_event_idx].id !== 'onSeparator' ||
+			this.events[colon_event_idx].arg !== ':'
+		) {
+			throw new Error(`Unexpected error: Colon not found`);
+		}
+
+		const value_event_idx = colon_event_idx + 1;
+		const value_event = this.events[value_event_idx];
+		if (value_event.id !== 'onLiteralValue') {
+			throw new Error(`Unexpected JSON value type: ${value_event.id}`);
+		}
+
+		const ext_claim_end_event_idx = value_event_idx + 1;
+		const ext_claim_end_event = this.events[ext_claim_end_event_idx];
+		if (ext_claim_end_event.id !== 'onSeparator' && ext_claim_end_event.id !== 'onObjectEnd') {
+			throw new Error(`Unexpected ext_claim_end_event ${ext_claim_end_event.id}`);
+		}
+
+		if (value_event.arg === undefined) {
+			throw new Error(`Undefined type for ${name}`);
+		}
+		if (
+			typeof value_event.arg !== 'string' &&
+			typeof value_event.arg !== 'boolean' &&
+			typeof value_event.arg !== 'number'
+		) {
+			throw new Error(`Unexpected type for ${name}: ${typeof value_event.arg}`);
+		}
+		this.processed[name] = {
+			name: name,
+			value: value_event.arg,
+			ext_claim: this.decoded_payload.slice(name_event.offset, ext_claim_end_event.offset + 1),
+			offsets: {
+				start: name_event.offset,
+				colon: colon_event.offset - name_event.offset,
+				value: value_event.offset - name_event.offset,
+				value_length: value_event.length,
+				name_length: name_event.length,
+				ext_length: ext_claim_end_event.offset - name_event.offset + 1,
+			},
+		};
+		return this.processed[name];
+	}
+
+	/**
+	 * Returns the claim value exactly as it appears in the JWT.
+	 * So, if it has escapes, no unescaping is done.
+	 * Assumes that the claim value is a string.
+	 * (reasonable as aud and common key claims like sub, email and username are JSON strings)
+	 *
+	 * @param name claim name
+	 * @returns claim value as it appears in the JWT without the quotes. The quotes are omitted to faciliate address derivation.
+	 *
+	 * NOTE: This function is only used to obtain claim values for address generation.
+	 * Do not use it elsewhere unless you know what you're doing.
+	 */
+	getRawClaimValue(name: string): string {
+		if (!Object.prototype.hasOwnProperty.call(this.processed, name)) {
+			throw new Error('Claim ' + name + ' not processed');
+		}
+		const details = this.processed[name];
+		if (typeof details.value !== 'string') {
+			throw new Error(`Claim ${name} does not have a string value.`);
+		}
+
+		const value_index = details.offsets.value + details.offsets.start;
+		const value_length = details.offsets.value_length;
+		if (this.decoded_payload[value_index] !== '"') {
+			throw new Error(`Claim ${name} does not have a string value.`);
+		}
+		if (this.decoded_payload[value_index + value_length - 1] !== '"') {
+			throw new Error(`Claim ${name} does not have a string value.`);
+		}
+
+		const raw_value = this.decoded_payload.slice(value_index + 1, value_index + value_length - 1); // omit the quotes
+		return raw_value;
+	}
+}