Skip to content

Commit

Permalink
zklogin ts sdk (MystenLabs#13399)
Browse files Browse the repository at this point in the history
## 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
  • Loading branch information
Jordan-Mysten authored Aug 15, 2023
1 parent 980bc05 commit 5d34439
Show file tree
Hide file tree
Showing 16 changed files with 373 additions and 272 deletions.
5 changes: 5 additions & 0 deletions .changeset/nine-bags-run.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@mysten/zklogin': patch
---

Initial experimental zklogin SDK
2 changes: 1 addition & 1 deletion .eslintrc.js
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,7 @@ module.exports = {
},
},
{
files: ['sdk/ledgerjs-hw-app-sui/**/*', 'apps/wallet/**/*'],
files: ['sdk/ledgerjs-hw-app-sui/**/*', 'apps/wallet/**/*', 'sdk/zklogin/**/*'],
rules: {
// ledgerjs-hw-app-sui and wallet use Buffer
'no-restricted-globals': ['off'],
Expand Down
3 changes: 2 additions & 1 deletion apps/wallet/configs/ts/tsconfig.common.json
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,8 @@
"@mysten/kiosk": ["../../sdk/kiosk/src/index.ts"],
"@mysten/bcs": ["../../sdk/bcs/src/"],
"@mysten/ledgerjs-hw-app-sui": ["../../sdk/ledgerjs-hw-app-sui/src/Sui.ts"],
"@mysten/wallet-standard": ["../../sdk/wallet-standard/src/"]
"@mysten/wallet-standard": ["../../sdk/wallet-standard/src/"],
"@mysten/zklogin": ["../../sdk/zklogin/src/"]
}
},
"include": ["../../src", "../../tests"],
Expand Down
1 change: 1 addition & 0 deletions apps/wallet/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@
"@mysten/ledgerjs-hw-app-sui": "workspace:*",
"@mysten/sui.js": "workspace:*",
"@mysten/wallet-standard": "workspace:*",
"@mysten/zklogin": "workspace:*",
"@noble/hashes": "^1.3.1",
"@radix-ui/react-checkbox": "^1.0.4",
"@radix-ui/react-collapsible": "^1.0.3",
Expand Down
7 changes: 3 additions & 4 deletions apps/wallet/src/background/accounts/zk/ZkAccount.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,16 @@ import {
toSerializedSignature,
} from '@mysten/sui.js/cryptography';
import { fromB64, toB64 } from '@mysten/sui.js/utils';
import { computeZkAddress, zkBcs } from '@mysten/zklogin';
import { blake2b } from '@noble/hashes/blake2b';
import { toBigIntBE } from 'bigint-buffer';
import { decodeJwt } from 'jose';
import { bcs } from './bcs';
import { getCurrentEpoch } from './current-epoch';
import { type ZkProvider } from './providers';
import {
type PartialZkSignature,
createPartialZKSignature,
fetchPin,
getAddress,
prepareZKLogin,
zkLogin,
} from './utils';
Expand Down Expand Up @@ -117,7 +116,7 @@ export class ZkAccount
};
return {
type: 'zk',
address: await getAddress({
address: computeZkAddress({
claimName: 'sub',
claimValue: decodedJWT.sub,
iss: decodedJWT.iss,
Expand Down Expand Up @@ -229,7 +228,7 @@ export class ZkAccount
signatureScheme: keyPair.getKeyScheme(),
publicKey: keyPair.getPublicKey(),
});
const bytes = bcs
const bytes = zkBcs
.ser(
'ZkSignature',
{
Expand Down
202 changes: 3 additions & 199 deletions apps/wallet/src/background/accounts/zk/utils.ts
Original file line number Diff line number Diff line change
@@ -1,189 +1,20 @@
// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

import { SIGNATURE_SCHEME_TO_FLAG } from '@mysten/sui.js/cryptography';
import { Ed25519Keypair } from '@mysten/sui.js/keypairs/ed25519';
import { SUI_ADDRESS_LENGTH, normalizeSuiAddress } from '@mysten/sui.js/utils';
import { blake2b } from '@noble/hashes/blake2b';
import { bytesToHex, randomBytes } from '@noble/hashes/utils';
import { generateNonce } from '@mysten/zklogin';
import { randomBytes } from '@noble/hashes/utils';
import { toBigIntBE } from 'bigint-buffer';
import { base64url } from 'jose';
import {
poseidon1,
poseidon2,
poseidon3,
poseidon4,
poseidon5,
poseidon6,
poseidon7,
poseidon8,
poseidon9,
poseidon10,
poseidon11,
poseidon12,
poseidon13,
poseidon14,
poseidon15,
poseidon16,
} from 'poseidon-lite';
import Browser from 'webextension-polyfill';
import { bcs } from './bcs';
import { type ZkProvider, zkProviderDataMap } from './providers';
import { fetchWithSentry } from '_src/shared/utils';

const maxKeyClaimNameLength = 40;
const maxKeyClaimValueLength = 256;
const packWidth = 248;
const nonceLength = 27;
const poseidonNumToHashFN = [
undefined,
poseidon1,
poseidon2,
poseidon3,
poseidon4,
poseidon5,
poseidon6,
poseidon7,
poseidon8,
poseidon9,
poseidon10,
poseidon11,
poseidon12,
poseidon13,
poseidon14,
poseidon15,
poseidon16,
];
type bit = 0 | 1;

interface UserInfo {
name: string;
value: string;
}

function poseidonHash(inputs: (string | number | bigint)[]): bigint {
const hashFN = poseidonNumToHashFN[inputs.length];
if (hashFN) {
return hashFN(inputs);
} else if (inputs.length <= 32) {
const hash1 = poseidonHash(inputs.slice(0, 16));
const hash2 = poseidonHash(inputs.slice(16));
return poseidonHash([hash1, hash2]);
} else {
throw new Error(`Yet to implement: Unable to hash a vector of length ${inputs.length}`);
}
}

function padWithZeroes<T>(inArr: T[], outCount: number) {
if (inArr.length > outCount) {
throw new Error('inArr is big enough');
}
const extraZeroes = outCount - inArr.length;
const arrPadded = inArr.concat(Array(extraZeroes).fill(0));
return arrPadded;
}

function bigIntArray2Bits(arr: bigint[], intSize: number): bit[] {
return arr.reduce((bitArray: bit[], n) => {
const binaryString = n.toString(2).padStart(intSize, '0');
const bitValues = binaryString.split('').map((bit) => (bit === '1' ? 1 : 0));
return bitArray.concat(bitValues);
}, []);
}

function arrayChunk<T>(array: T[], chunkSize: number): T[][] {
return Array(Math.ceil(array.length / chunkSize))
.fill(undefined)
.map((_, index) => index * chunkSize)
.map((begin) => array.slice(begin, begin + chunkSize));
}

// Pack into an array of chunks each outWidth bits
function pack(inArr: bigint[], inWidth: number, outWidth: number): bigint[] {
const bits = bigIntArray2Bits(inArr, inWidth);
const extraBits = bits.length % outWidth === 0 ? 0 : outWidth - (bits.length % outWidth);
const bitsPadded = bits.concat(Array(extraBits).fill(0));
if (bitsPadded.length % outWidth !== 0) {
throw new Error('Invalid logic');
}
const packed = arrayChunk(bitsPadded, outWidth).map((chunk) => BigInt('0b' + chunk.join('')));
return packed;
}

// Pack into exactly outCount chunks of outWidth bits each
function pack2(inArr: bigint[], inWidth: number, outWidth: number, outCount: number): bigint[] {
const packed = pack(inArr, inWidth, outWidth);
if (packed.length > outCount) {
throw new Error('packed is big enough');
}
return packed.concat(Array(outCount - packed.length).fill(0));
}

async function mapToField(input: bigint[], inWidth: number) {
if (packWidth % 8 !== 0) {
throw new Error('packWidth must be a multiple of 8');
}
const numElements = Math.ceil((input.length * inWidth) / packWidth);
const packed = pack2(input, inWidth, packWidth, numElements);
return poseidonHash(packed);
}

// Pads a stream of bytes and maps it to a field element
async function mapBytesToField(str: string, maxSize: number) {
if (str.length > maxSize) {
throw new Error(`String ${str} is longer than ${maxSize} chars`);
}
// Note: Padding with zeroes is safe because we are only using this function to map human-readable sequence of bytes.
// So the ASCII values of those characters will never be zero (null character).
const strPadded = padWithZeroes(
str.split('').map((c) => BigInt(c.charCodeAt(0))),
maxSize,
);
return mapToField(strPadded, 8);
}

async function genAddressSeed(pin: bigint, { name, value }: UserInfo) {
if (name.length > maxKeyClaimNameLength) {
throw new Error('Name is too long');
}
if (value.length > maxKeyClaimValueLength) {
throw new Error('Value is too long');
}
return poseidonHash([
await mapBytesToField(name, maxKeyClaimNameLength),
await mapBytesToField(value, maxKeyClaimValueLength),
poseidonHash([pin]),
]);
}

// use custom function because when having a width 20 the library
// returns different results between native and browser computations
function toBufferBE(num: bigint, width: number) {
const hex = num.toString(16);
return Buffer.from(hex.padStart(width * 2, '0').slice(-width * 2), 'hex');
}

function generateNonce(ephemeralPublicKey: bigint, maxEpoch: number, randomness: bigint) {
const eph_public_key_0 = ephemeralPublicKey / 2n ** 128n;
const eph_public_key_1 = ephemeralPublicKey % 2n ** 128n;
const bigNum = poseidonHash([eph_public_key_0, eph_public_key_1, maxEpoch, randomness]);
const Z = toBufferBE(bigNum, 20);
const nonce = base64url.encode(Z);
if (nonce.length !== nonceLength) {
throw new Error(`Length of nonce ${nonce} (${nonce.length}) is not equal to ${nonceLength}`);
}
return nonce;
}

export function prepareZKLogin(currentEpoch: number) {
const maxEpoch = currentEpoch + 2;
const ephemeralKeyPair = new Ed25519Keypair();
const randomness = toBigIntBE(Buffer.from(randomBytes(16)));
const nonce = generateNonce(
toBigIntBE(Buffer.from(ephemeralKeyPair.getPublicKey().toRawBytes())),
maxEpoch,
randomness,
);
const nonce = generateNonce(ephemeralKeyPair.getPublicKey(), maxEpoch, randomness);
return {
ephemeralKeyPair,
randomness,
Expand All @@ -192,33 +23,6 @@ export function prepareZKLogin(currentEpoch: number) {
};
}

export async function getAddress({
claimName,
claimValue,
userPin,
iss,
aud,
}: {
claimName: string;
claimValue: string;
userPin: bigint;
iss: string;
aud: string;
}) {
const addressSeedBytes = toBufferBE(
await genAddressSeed(userPin, { name: claimName, value: claimValue }),
32,
);
const addressParamBytes = bcs.ser('AddressParams', { iss, aud }).toBytes();
const tmp = new Uint8Array(1 + addressSeedBytes.length + addressParamBytes.length);
tmp.set([SIGNATURE_SCHEME_TO_FLAG['Zk']]);
tmp.set(addressParamBytes, 1);
tmp.set(addressSeedBytes, 1 + addressParamBytes.length);
return normalizeSuiAddress(
bytesToHex(blake2b(tmp, { dkLen: 32 })).slice(0, SUI_ADDRESS_LENGTH * 2),
);
}

export async function zkLogin({
provider,
nonce,
Expand Down
Loading

0 comments on commit 5d34439

Please sign in to comment.