return non-legacy bytes from legacy addresses ZKLoginPublicIdentifiers

sdk/typescript/src/zklogin/publickey.ts

@@ -1,13 +1,16 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-import { fromBase64, toBase64 } from '@mysten/bcs';
+import { fromBase64, toBase64, toHex } from '@mysten/bcs';
+import { blake2b } from '@noble/hashes/blake2b';
+import { bytesToHex } from '@noble/hashes/utils';
 
 import { PublicKey } from '../cryptography/publickey.js';
 import type { PublicKeyInitData } from '../cryptography/publickey.js';
 import { SIGNATURE_SCHEME_TO_FLAG } from '../cryptography/signature-scheme.js';
 import { SuiGraphQLClient } from '../graphql/client.js';
 import { graphql } from '../graphql/schemas/2024.4/index.js';
+import { normalizeSuiAddress, SUI_ADDRESS_LENGTH } from '../utils/sui-types.js';
 import { extractClaimValue } from './jwt-utils.js';
 import { parseZkLoginSignature } from './signature.js';
 import { toBigEndianBytes, toPaddedBigEndianBytes } from './utils.js';
@@ -18,6 +21,7 @@ import { toBigEndianBytes, toPaddedBigEndianBytes } from './utils.js';
 export class ZkLoginPublicIdentifier extends PublicKey {
 	#data: Uint8Array;
 	#client?: SuiGraphQLClient;
+	#legacyAddress: boolean;
 
 	/**
 	 * Create a new ZkLoginPublicIdentifier object
@@ -35,6 +39,11 @@ export class ZkLoginPublicIdentifier extends PublicKey {
 		} else {
 			this.#data = Uint8Array.from(value);
 		}
+		this.#legacyAddress = this.#data.length !== this.#data[0] + 1 + 32;
+
+		if (this.#legacyAddress) {
+			this.#data = normalizeZkLoginPublicKeyBytes(this.#data);
+		}
 	}
 
 	/**
@@ -44,6 +53,20 @@ export class ZkLoginPublicIdentifier extends PublicKey {
 		return super.equals(publicKey);
 	}
 
+	override toSuiAddress(): string {
+		if (this.#legacyAddress) {
+			const legacyBytes = normalizeZkLoginPublicKeyBytes(this.#data, true);
+			const addressBytes = new Uint8Array(legacyBytes.length + 1);
+			addressBytes[0] = this.flag();
+			addressBytes.set(legacyBytes, 1);
+			return normalizeSuiAddress(
+				bytesToHex(blake2b(addressBytes, { dkLen: 32 })).slice(0, SUI_ADDRESS_LENGTH * 2),
+			);
+		}
+
+		return super.toSuiAddress();
+	}
+
 	/**
 	 * Return the byte array representation of the zkLogin public identifier
 	 */
@@ -107,6 +130,7 @@ export function toZkLoginPublicIdentifier(
 	const addressSeedBytesBigEndian = options?.legacyAddress
 		? toBigEndianBytes(addressSeed, 32)
 		: toPaddedBigEndianBytes(addressSeed, 32);
+
 	const issBytes = new TextEncoder().encode(iss);
 	const tmp = new Uint8Array(1 + issBytes.length + addressSeedBytesBigEndian.length);
 	tmp.set([issBytes.length], 0);
@@ -134,6 +158,18 @@ const VerifyZkLoginSignatureQuery = graphql(`
 	}
 `);
 
+function normalizeZkLoginPublicKeyBytes(bytes: Uint8Array, legacyAddress = false) {
+	const issByteLength = bytes[0] + 1;
+	const addressSeed = BigInt(`0x${toHex(bytes.slice(issByteLength))}`);
+	const seedBytes = legacyAddress
+		? toBigEndianBytes(addressSeed, 32)
+		: toPaddedBigEndianBytes(addressSeed, 32);
+	const data = new Uint8Array(issByteLength + seedBytes.length);
+	data.set(bytes.slice(0, issByteLength), 0);
+	data.set(seedBytes, issByteLength);
+	return data;
+}
+
 async function graphqlVerifyZkLoginSignature({
 	address,
 	bytes,