From ad8ea6e4e8c49a03e5dc8212f6d43bbbb0cbab10 Mon Sep 17 00:00:00 2001 From: Ho Nguyen Pham Date: Sat, 28 Oct 2023 21:26:50 +0200 Subject: [PATCH] refactor and add zkapp compatibility --- package.json | 2 +- .../Committee.test.ts} | 81 ++++++++++-- src/{CommitteeMember.ts => dkg/Committee.ts} | 124 +++++------------- src/dkg/Requestor.ts | 57 ++++++++ src/dkg/index.ts | 4 + src/index.ts | 12 +- src/utils/CustomScalar.test.ts | 14 ++ src/utils/CustomScalar.ts | 55 ++++++++ src/{ => utils}/DynamicArray.test.ts | 32 ++--- src/{ => utils}/DynamicArray.ts | 6 +- src/{ => utils}/Elgamal.test.ts | 0 src/{ => utils}/Elgamal.ts | 0 src/utils/index.ts | 5 + 13 files changed, 262 insertions(+), 130 deletions(-) rename src/{CommitteeMember.test.ts => dkg/Committee.test.ts} (64%) rename src/{CommitteeMember.ts => dkg/Committee.ts} (64%) create mode 100644 src/dkg/Requestor.ts create mode 100644 src/dkg/index.ts create mode 100644 src/utils/CustomScalar.test.ts create mode 100644 src/utils/CustomScalar.ts rename src/{ => utils}/DynamicArray.test.ts (68%) rename src/{ => utils}/DynamicArray.ts (98%) rename src/{ => utils}/Elgamal.test.ts (100%) rename src/{ => utils}/Elgamal.ts (100%) create mode 100644 src/utils/index.ts diff --git a/package.json b/package.json index 8992508..3a79850 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@auxo-dev/dkg-libs", - "version": "0.1.2", + "version": "0.1.3", "description": "", "author": "", "license": "Apache-2.0", diff --git a/src/CommitteeMember.test.ts b/src/dkg/Committee.test.ts similarity index 64% rename from src/CommitteeMember.test.ts rename to src/dkg/Committee.test.ts index ce3a660..b424874 100644 --- a/src/CommitteeMember.test.ts +++ b/src/dkg/Committee.test.ts @@ -1,10 +1,22 @@ -import { Field, Group, Provable, PrivateKey, PublicKey, Scalar } from 'o1js'; -import * as Committee from './CommitteeMember'; +import { + Field, + Group, + method, + Provable, + PrivateKey, + PublicKey, + Scalar, + SmartContract, + state, + State, + Reducer, +} from 'o1js'; +import * as Committee from './Committee.js'; +import * as Requestor from './Requestor.js'; -describe('CommitteeMember', () => { +describe('Committee', () => { let T = 3; let N = 5; - let keyId = Field(0); let committees: { privateKey: PrivateKey; index: number; @@ -35,12 +47,6 @@ describe('CommitteeMember', () => { beforeAll(async () => { for (let i = 0; i < N; i++) { let privateKey = PrivateKey.random(); - // let committeeMember = { - // publicKey: privateKey.toPublicKey(), - // index: i + 1, - // round1Contribution: undefined, - // round2Contribution: undefined, - // }; let secretPolynomial = Committee.generateRandomPolynomial(T, N); committees.push({ privateKey: privateKey, @@ -59,6 +65,7 @@ describe('CommitteeMember', () => { ); committees[i].round1Contribution = round1Contribution; round1Contributions.push(round1Contribution); + Provable.runAndCheck(() => round1Contribution); } publicKey = Committee.calculatePublicKey(round1Contributions); // Provable.log(publicKey); @@ -74,6 +81,7 @@ describe('CommitteeMember', () => { ); committees[i].round2Contribution = round2Contribution; round2Contributions.push(round2Contribution); + Provable.runAndCheck(() => round2Contribution); } // Provable.log(round2Contributions); // round2Contributions.map((e) => console.log(e.data)); @@ -81,7 +89,7 @@ describe('CommitteeMember', () => { it('Should accumulate encryption', async () => { for (let i = 0; i < plainVectors.length; i++) { - let encryptedVector = Committee.encryptVector( + let encryptedVector = Requestor.generateEncryption( Committee.calculatePublicKey(round1Contributions), plainVectors[i] ); @@ -89,7 +97,7 @@ describe('CommitteeMember', () => { M.push(encryptedVector.M); } - let accumulatedEncryption = Committee.accumulateEncryption(R, M); + let accumulatedEncryption = Requestor.accumulateEncryption(R, M); sumR = accumulatedEncryption.sumR; sumM = accumulatedEncryption.sumM; // Provable.log(sumR, sumM); @@ -103,7 +111,15 @@ describe('CommitteeMember', () => { index == committees[listIndex[i] - 1].index - 1 ? prev : round2Data.push( - curr.data[committees[listIndex[i] - 1].index - 1] + { + c: curr.c.values[ + committees[listIndex[i] - 1].index - 1 + ].toScalar(), + U: curr.U.values[ + committees[listIndex[i] - 1].index - 1 + ], + } + // curr.data[committees[listIndex[i] - 1].index - 1] ), {} ); @@ -113,9 +129,10 @@ describe('CommitteeMember', () => { round2Data, sumR ); + Provable.runAndCheck(() => tallyContribution); committees[listIndex[i] - 1].tallyContribution = tallyContribution; tallyContributions.push(tallyContribution); - D.push(tallyContribution.D); + D.push(tallyContribution.D.values.slice(0, T)); } // Provable.log(tallyContributions); // tallyContributions.map((e) => Provable.log(e.D)); @@ -132,4 +149,40 @@ describe('CommitteeMember', () => { expect(resultVector[i].y).toEqual(point.y); } }); + + it('Should be used in Smart Contract', async () => { + class TestRound1Contribution extends SmartContract { + reducer = Reducer({ actionType: Committee.Round1Contribution }); + @state(Field) keyId = State(); + @method test(): Field { + return Field(0); + } + } + + class TestRound2Contribution extends SmartContract { + reducer = Reducer({ actionType: Committee.Round2Contribution }); + @state(Field) keyId = State(); + @method test(): Field { + return Field(0); + } + } + + class TestTallyContribution extends SmartContract { + reducer = Reducer({ actionType: Committee.TallyContribution }); + @state(Field) keyId = State(); + @method test(): Field { + return Field(0); + } + } + + console.log('Compile test round 1...'); + await TestRound1Contribution.compile(); + console.log('DONE!'); + console.log('Compile test round 2...'); + await TestRound2Contribution.compile(); + console.log('DONE!'); + console.log('Compile test tally...'); + await TestTallyContribution.compile(); + console.log('DONE!'); + }); }); diff --git a/src/CommitteeMember.ts b/src/dkg/Committee.ts similarity index 64% rename from src/CommitteeMember.ts rename to src/dkg/Committee.ts index 0fb495b..26bcc90 100644 --- a/src/CommitteeMember.ts +++ b/src/dkg/Committee.ts @@ -1,18 +1,7 @@ -import { - Bool, - Encryption, - Field, - Group, - MerkleWitness, - Poseidon, - PrivateKey, - Provable, - PublicKey, - Scalar, - Struct, - UInt32, -} from 'o1js'; -import * as ElgamalECC from './Elgamal'; +import { Group, PrivateKey, PublicKey, Scalar, Struct } from 'o1js'; +import * as ElgamalECC from '../utils/Elgamal.js'; +import { DynamicArray } from '../utils/DynamicArray.js'; +import { CustomScalar } from '../utils/CustomScalar.js'; export { SecretPolynomial, @@ -28,39 +17,41 @@ export { getTallyContribution, getLagrangeCoefficient, getResultVector, - encryptVector, - accumulateEncryption, }; +const GroupDynamicArray = DynamicArray(Group, 32); +const ScalarDynamicArray = DynamicArray(CustomScalar, 32); + type SecretPolynomial = { a: Scalar[]; C: Group[]; f: Scalar[]; }; -type Round1Contribution = { - C: Group[]; -}; - type Round2Data = { - c: bigint; + c: Scalar; U: Group; }; -type Round2Contribution = { - data: Round2Data[]; -}; +class Round1Contribution extends Struct({ + C: GroupDynamicArray, +}) {} -type TallyContribution = { - D: Group[]; -}; +class Round2Contribution extends Struct({ + c: ScalarDynamicArray, + U: GroupDynamicArray, +}) {} + +class TallyContribution extends Struct({ + D: GroupDynamicArray, +}) {} function calculatePublicKey( round1Contributions: Round1Contribution[] ): PublicKey { let result = Group.zero; for (let i = 0; i < round1Contributions.length; i++) { - result = result.add(round1Contributions[i].C[0]); + result = result.add(round1Contributions[i].C.values[0]); } return PublicKey.fromGroup(result); } @@ -88,7 +79,8 @@ function generateRandomPolynomial(T: number, N: number): SecretPolynomial { } function getRound1Contribution(secret: SecretPolynomial): Round1Contribution { - return { C: secret.C }; + let provableC = GroupDynamicArray.from(secret.C); + return { C: provableC }; } function getRound2Contribution( @@ -97,24 +89,26 @@ function getRound2Contribution( round1Contributions: Round1Contribution[] ): Round2Contribution { let data = new Array(secret.f.length); + let c = new Array(secret.f.length); + let U = new Array(secret.f.length); for (let i = 0; i < data.length; i++) { if (i + 1 == index) { - data[i] = { - U: Group.zero, - c: 0n, - }; + c[i] = Scalar.from(0n); + U[i] = Group.zero; } else { let encryption = ElgamalECC.encrypt( secret.f[i].toBigInt(), - PublicKey.fromGroup(round1Contributions[i].C[0]) + PublicKey.fromGroup(round1Contributions[i].C.values[0]) ); - data[i] = { - U: encryption.U, - c: encryption.c, - }; + c[i] = Scalar.from(encryption.c); + U[i] = encryption.U; } } - return { data }; + let provablec = ScalarDynamicArray.from( + c.map((e) => CustomScalar.fromScalar(e)) + ); + let provableU = GroupDynamicArray.from(U); + return { c: provablec, U: provableU }; } function getTallyContribution( @@ -126,7 +120,7 @@ function getTallyContribution( let decryptions: Scalar[] = round2Data.map((data) => Scalar.from( ElgamalECC.decrypt( - data.c, + data.c.toBigInt(), data.U, PrivateKey.fromBigInt(secret.a[0].toBigInt()) ).m @@ -141,7 +135,7 @@ function getTallyContribution( for (let i = 0; i < R.length; i++) { D[i] = R[i].scale(ski); } - return { D }; + return { D: GroupDynamicArray.from(D) }; } function getLagrangeCoefficient(listIndex: number[]): Scalar[] { @@ -184,49 +178,3 @@ function getResultVector( } return result; } - -function encryptVector( - publicKey: PublicKey, - vector: bigint[] -): { - r: Scalar[]; - R: Group[]; - M: Group[]; -} { - let dimension = vector.length; - let r = new Array(dimension); - let R = new Array(dimension); - let M = new Array(dimension); - for (let i = 0; i < dimension; i++) { - let random = Scalar.random(); - r[i] = random; - R[i] = Group.generator.scale(random); - M[i] = - vector[i] > 0n - ? Group.generator - .scale(Scalar.from(vector[i])) - .add(publicKey.toGroup().scale(random)) - : Group.zero.add(publicKey.toGroup().scale(random)); - } - return { r, R, M }; -} - -function accumulateEncryption( - R: Group[][], - M: Group[][] -): { sumR: Group[]; sumM: Group[] } { - let quantity = R.length; - let dimension = R[0].length ?? 0; - let sumR = new Array(dimension); - let sumM = new Array(dimension); - sumR.fill(Group.zero); - sumM.fill(Group.zero); - - for (let i = 0; i < quantity; i++) { - for (let j = 0; j < dimension; j++) { - sumR[j] = sumR[j].add(R[i][j]); - sumM[j] = sumM[j].add(M[i][j]); - } - } - return { sumR, sumM }; -} diff --git a/src/dkg/Requestor.ts b/src/dkg/Requestor.ts new file mode 100644 index 0000000..b5247c7 --- /dev/null +++ b/src/dkg/Requestor.ts @@ -0,0 +1,57 @@ +import { Group, PublicKey, Scalar } from 'o1js'; + +export { calculatePublicKey, generateEncryption, accumulateEncryption }; + +function calculatePublicKey(contributedPublicKeys: Group[]): PublicKey { + let result = Group.zero; + for (let i = 0; i < contributedPublicKeys.length; i++) { + result = result.add(contributedPublicKeys[i]); + } + return PublicKey.fromGroup(result); +} + +function generateEncryption( + publicKey: PublicKey, + vector: bigint[] +): { + r: Scalar[]; + R: Group[]; + M: Group[]; +} { + let dimension = vector.length; + let r = new Array(dimension); + let R = new Array(dimension); + let M = new Array(dimension); + for (let i = 0; i < dimension; i++) { + let random = Scalar.random(); + r[i] = random; + R[i] = Group.generator.scale(random); + M[i] = + vector[i] > 0n + ? Group.generator + .scale(Scalar.from(vector[i])) + .add(publicKey.toGroup().scale(random)) + : Group.zero.add(publicKey.toGroup().scale(random)); + } + return { r, R, M }; +} + +function accumulateEncryption( + R: Group[][], + M: Group[][] +): { sumR: Group[]; sumM: Group[] } { + let quantity = R.length; + let dimension = R[0].length ?? 0; + let sumR = new Array(dimension); + let sumM = new Array(dimension); + sumR.fill(Group.zero); + sumM.fill(Group.zero); + + for (let i = 0; i < quantity; i++) { + for (let j = 0; j < dimension; j++) { + sumR[j] = sumR[j].add(R[i][j]); + sumM[j] = sumM[j].add(M[i][j]); + } + } + return { sumR, sumM }; +} diff --git a/src/dkg/index.ts b/src/dkg/index.ts new file mode 100644 index 0000000..ba71f13 --- /dev/null +++ b/src/dkg/index.ts @@ -0,0 +1,4 @@ +import * as Committee from './Committee.js'; +import * as Requestor from './Requestor.js'; + +export { Committee, Requestor }; diff --git a/src/index.ts b/src/index.ts index 2dde311..2950f52 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,10 +1,4 @@ -import * as CommitteeMember from './CommitteeMember'; -import * as DynamicArray from './DynamicArray'; -import * as ElgamalECC from './Elgamal'; +import * as DKG from './dkg/index.js'; +import * as Utils from './utils/index.js'; -export { CommitteeMember, DynamicArray, ElgamalECC }; -// module.exports = { -// CommitteeMember: CommitteeMember, -// DynamicArray: DynamicArray, -// ElgamalECC: ElgamalECC, -// }; +export { DKG, Utils }; diff --git a/src/utils/CustomScalar.test.ts b/src/utils/CustomScalar.test.ts new file mode 100644 index 0000000..b7e5add --- /dev/null +++ b/src/utils/CustomScalar.test.ts @@ -0,0 +1,14 @@ +import { PrivateKey, Provable, Scalar } from 'o1js'; +import { CustomScalar } from './CustomScalar'; + +describe('CustomScalar', () => { + it('Should convert correctly', async () => { + let original = Scalar.random(); + console.log('Original:', original); + let encoded = CustomScalar.fromScalar(original); + // console.log('Encoded:', encoded); + let decoded = encoded.toScalar(); + // console.log('Decoded:', decoded); + expect(original).toEqual(decoded); + }); +}); diff --git a/src/utils/CustomScalar.ts b/src/utils/CustomScalar.ts new file mode 100644 index 0000000..c398433 --- /dev/null +++ b/src/utils/CustomScalar.ts @@ -0,0 +1,55 @@ +import { Field, Scalar, Struct } from 'o1js'; + +export class CustomScalar extends Struct({ + head: Field, + tail: Field, +}) { + static fromScalar(scalar: Scalar): CustomScalar { + let bits = scalar.toFields().map((e) => e.toBits()[0]); + // console.log('Bits:', bits); + return new CustomScalar({ + head: Field.fromBits(bits.slice(0, 127)), + tail: Field.fromBits(bits.slice(127)), + }); + // return new CustomScalar({ + // head: Field(0), + // tail: Field(0), + // }); + } + + static fromFields(fields: Field[]): CustomScalar { + return new CustomScalar({ + head: fields[0], + tail: fields[1], + }); + } + + static toScalar(scalar: CustomScalar): Scalar { + return Scalar.fromBits( + scalar.head + .toBits() + .slice(0, 127) + .concat(scalar.tail.toBits().slice(0, 128)) + ); + } + + static toFields(value: { head: Field; tail: Field }): Field[] { + return [value.head].concat([value.tail]); + } + + static sizeInFields(): number { + return 2; + } + + toScalar(): Scalar { + return CustomScalar.toScalar(this); + } + + toFields(): Field[] { + return CustomScalar.toFields(this); + } + + fromFields(fields: Field[]): CustomScalar { + return CustomScalar.fromFields(fields); + } +} diff --git a/src/DynamicArray.test.ts b/src/utils/DynamicArray.test.ts similarity index 68% rename from src/DynamicArray.test.ts rename to src/utils/DynamicArray.test.ts index b485187..3ca4fac 100644 --- a/src/DynamicArray.test.ts +++ b/src/utils/DynamicArray.test.ts @@ -3,6 +3,7 @@ import { FlexibleProvable, FlexibleProvablePure, Group, + Mina, Provable, Reducer, Scalar, @@ -10,12 +11,18 @@ import { Struct, } from 'o1js'; import { DynamicArray } from './DynamicArray'; +import { CustomScalar } from './CustomScalar'; +import { + Round1Contribution, + Round2Contribution, + TallyContribution, +} from '../dkg/Committee'; describe('DynamicArray', () => { const MAX_HEIGHT = 2 ** 5; class DynamicFieldArray extends DynamicArray(Field, MAX_HEIGHT) {} class DynamicGroupArray extends DynamicArray(Group, MAX_HEIGHT) {} - class DynamicScalarArray extends DynamicArray(Scalar, MAX_HEIGHT) {} + class DynamicScalarArray extends DynamicArray(CustomScalar, MAX_HEIGHT) {} it('Should be provable', async () => { Provable.runAndCheck(() => { @@ -42,21 +49,16 @@ describe('DynamicArray', () => { Provable.log(groupArray); // Scalar - Not working because Scalar takes up 255 Fields => Stack overflow - // let scalarValues = Provable.Array(Scalar, 1).fromFields( - // Scalar.from(Scalar.ORDER).toFields() - // ); - // let scalarArray = Provable.witness( - // DynamicScalarArray, - // () => new DynamicScalarArray(scalarValues) - // ); + let scalarValues = Provable.Array(CustomScalar, 1).fromFields( + CustomScalar.fromScalar(Scalar.random()).toFields() + ); + Provable.log(scalarValues); + let scalarArray = Provable.witness( + DynamicScalarArray, + () => new DynamicScalarArray(scalarValues) + ); // Provable.log(Scalar.sizeInFields()); - // Provable.log(scalarArray); + Provable.log(scalarArray); }); }); - - xit('Should be used in Smart Contract', async () => { - class Test extends SmartContract { - reducer = Reducer({ actionType: DynamicFieldArray }); - } - }); }); diff --git a/src/DynamicArray.ts b/src/utils/DynamicArray.ts similarity index 98% rename from src/DynamicArray.ts rename to src/utils/DynamicArray.ts index e6e039f..49efda4 100644 --- a/src/DynamicArray.ts +++ b/src/utils/DynamicArray.ts @@ -35,9 +35,9 @@ function DynamicArray(type: ProvablePure, maxLength: number) { length: Field, values: Provable.Array(type, maxLength), }) { - // static from(values: T[]): _DynamicArray { - // return new _DynamicArray(values); - // } + static from(values: T[]): _DynamicArray { + return new _DynamicArray(values); + } static empty(length?: Field): _DynamicArray { const arr = new _DynamicArray(); diff --git a/src/Elgamal.test.ts b/src/utils/Elgamal.test.ts similarity index 100% rename from src/Elgamal.test.ts rename to src/utils/Elgamal.test.ts diff --git a/src/Elgamal.ts b/src/utils/Elgamal.ts similarity index 100% rename from src/Elgamal.ts rename to src/utils/Elgamal.ts diff --git a/src/utils/index.ts b/src/utils/index.ts new file mode 100644 index 0000000..f7671ed --- /dev/null +++ b/src/utils/index.ts @@ -0,0 +1,5 @@ +import { CustomScalar } from './CustomScalar.js'; +import { DynamicArray } from './DynamicArray.js'; +import * as Elgamal from './Elgamal.js'; + +export { CustomScalar, DynamicArray, Elgamal };