Skip to content

Commit

Permalink
test: adjustments for reencrypt functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@PacificYield
PacificYield committed Dec 6, 2024
1 parent 951ab7b commit 9f32e20
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 80 deletions.
28 changes: 11 additions & 17 deletions test/confidentialERC20/ConfidentialERC20.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,8 @@ describe("ConfidentialERC20", function () {
// Reencrypt Alice's balance
const balanceHandleAlice = await this.erc20.balanceOf(this.signers.alice);
const balanceAlice = await reencryptEuint64(
this.signers,
this.signers.alice,
this.fhevm,
"alice",
balanceHandleAlice,
this.contractAddress,
);
Expand Down Expand Up @@ -59,27 +58,26 @@ describe("ConfidentialERC20", function () {
// Reencrypt Alice's balance
const balanceHandleAlice = await this.erc20.balanceOf(this.signers.alice);
const balanceAlice = await reencryptEuint64(
this.signers,
this.signers.alice,
this.fhevm,
"alice",
balanceHandleAlice,
this.contractAddress,
);
expect(balanceAlice).to.equal(10000 - 1337);

// Reencrypt Bob's balance
const balanceHandleBob = await this.erc20.balanceOf(this.signers.bob);
const balanceBob = await reencryptEuint64(this.signers, this.fhevm, "bob", balanceHandleBob, this.contractAddress);
const balanceBob = await reencryptEuint64(this.signers.bob, this.fhevm, balanceHandleBob, this.contractAddress);
expect(balanceBob).to.equal(1337);

// on the other hand, Bob should be unable to read Alice's balance
await expect(
reencryptEuint64(this.signers, this.fhevm, "bob", balanceHandleAlice, this.contractAddress),
reencryptEuint64(this.signers.bob, this.fhevm, balanceHandleAlice, this.contractAddress),
).to.be.rejectedWith("User is not authorized to reencrypt this handle!");

// and should be impossible to call reencrypt if contractAddress === userAddress
await expect(
reencryptEuint64(this.signers, this.fhevm, "alice", balanceHandleAlice, this.signers.alice.address),
reencryptEuint64(this.signers.alice, this.fhevm, balanceHandleAlice, this.signers.alice.address),
).to.be.rejectedWith("userAddress should not be equal to contractAddress when requesting reencryption!");
});

Expand All @@ -99,17 +97,16 @@ describe("ConfidentialERC20", function () {

const balanceHandleAlice = await this.erc20.balanceOf(this.signers.alice);
const balanceAlice = await reencryptEuint64(
this.signers,
this.signers.alice,
this.fhevm,
"alice",
balanceHandleAlice,
this.contractAddress,
);
expect(balanceAlice).to.equal(1000);

// Reencrypt Bob's balance
const balanceHandleBob = await this.erc20.balanceOf(this.signers.bob);
const balanceBob = await reencryptEuint64(this.signers, this.fhevm, "bob", balanceHandleBob, this.contractAddress);
const balanceBob = await reencryptEuint64(this.signers.bob, this.fhevm, balanceHandleBob, this.contractAddress);
expect(balanceBob).to.equal(0);
});

Expand Down Expand Up @@ -142,17 +139,16 @@ describe("ConfidentialERC20", function () {
// Decrypt Alice's balance
const balanceHandleAlice = await this.erc20.balanceOf(this.signers.alice);
const balanceAlice = await reencryptEuint64(
this.signers,
this.signers.alice,
this.fhevm,
"alice",
balanceHandleAlice,
this.contractAddress,
);
expect(balanceAlice).to.equal(10000); // check that transfer did not happen, as expected

// Decrypt Bob's balance
const balanceHandleBob = await this.erc20.balanceOf(this.signers.bob);
const balanceBob = await reencryptEuint64(this.signers, this.fhevm, "bob", balanceHandleBob, this.contractAddress);
const balanceBob = await reencryptEuint64(this.signers.bob, this.fhevm, balanceHandleBob, this.contractAddress);
expect(balanceBob).to.equal(0); // check that transfer did not happen, as expected

const inputBob2 = this.fhevm.createEncryptedInput(this.contractAddress, this.signers.bob.address);
Expand All @@ -169,9 +165,8 @@ describe("ConfidentialERC20", function () {
// Decrypt Alice's balance
const balanceHandleAlice2 = await this.erc20.balanceOf(this.signers.alice);
const balanceAlice2 = await reencryptEuint64(
this.signers,
this.signers.alice,
this.fhevm,
"alice",
balanceHandleAlice2,
this.contractAddress,
);
Expand All @@ -180,9 +175,8 @@ describe("ConfidentialERC20", function () {
// Decrypt Bob's balance
const balanceHandleBob2 = await this.erc20.balanceOf(this.signers.bob);
const balanceBob2 = await reencryptEuint64(
this.signers,
this.signers.bob,
this.fhevm,
"bob",
balanceHandleBob2,
this.contractAddress,
);
Expand Down
102 changes: 40 additions & 62 deletions test/reencrypt.ts
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,6 @@
import { Signer } from "ethers";
import { FhevmInstance } from "fhevmjs/node";

import { ACCOUNT_NAMES } from "./constants";
import { Signers } from "./signers";

// Add type definition at the top of the file
type AccountName = (typeof ACCOUNT_NAMES)[number];

const EBOOL_T = 0;
const EUINT4_T = 1;
const EUINT8_T = 2;
Expand Down Expand Up @@ -36,165 +31,148 @@ export function verifyType(handle: bigint, expectedType: number) {
}

export async function reencryptEbool(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<boolean> {
verifyType(handle, EBOOL_T);
return (await reencryptHandle(signers, instances, user, handle, contractAddress)) === 1n;
return (await reencryptHandle(signer, instance, handle, contractAddress)) === 1n;
}

export async function reencryptEuint4(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT4_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEuint8(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT8_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEuint16(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT16_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEuint32(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT32_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEuint64(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT64_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEuint128(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT128_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEaddress(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<string> {
verifyType(handle, EUINT160_T);
const addressAsUint160: bigint = await reencryptHandle(signers, instances, user, handle, contractAddress);
const addressAsUint160: bigint = await reencryptHandle(signer, instance, handle, contractAddress);
const handleStr = "0x" + addressAsUint160.toString(16).padStart(40, "0");
return handleStr;
}

export async function reencryptEuint256(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EUINT256_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEbytes64(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EBYTES64_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEbytes128(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EBYTES128_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

export async function reencryptEbytes256(
signers: Signers,
instances: FhevmInstance,
user: AccountName,
signer: Signer,
instance: FhevmInstance,
handle: bigint,
contractAddress: string,
): Promise<bigint> {
verifyType(handle, EBYTES256_T);
return reencryptHandle(signers, instances, user, handle, contractAddress);
return reencryptHandle(signer, instance, handle, contractAddress);
}

/**
* @dev This function is to reencrypt handles.
* It does not verify types.
*/
async function reencryptHandle(
signers: Signers,
signer: Signer,
instance: FhevmInstance,
user: AccountName,
handle: bigint,
contractAddress: string,
): Promise<any> {
const { publicKey: publicKey, privateKey: privateKey } = instance.generateKeypair();
const eip712 = instance.createEIP712(publicKey, contractAddress);
const signature = await signers[user as keyof Signers].signTypedData(
eip712.domain,
{ Reencrypt: eip712.types.Reencrypt },
eip712.message,
);
const signature = await signer.signTypedData(eip712.domain, { Reencrypt: eip712.types.Reencrypt }, eip712.message);

const reencryptedHandle = await instance.reencrypt(
handle,
privateKey,
publicKey,
signature.replace("0x", ""),
contractAddress,
signers[user as keyof Signers].address,
await signer.getAddress(),
);

return reencryptedHandle;
Expand Down
2 changes: 1 addition & 1 deletion test/signers.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ export interface Signers {
[K in AccountNames]: HardhatEthersSigner;
}

let signers: Signers = {} as Signers;
const signers: Signers = {} as Signers;

export const initSigners = async (): Promise<void> => {
if (Object.entries(signers).length === 0) {
Expand Down

0 comments on commit 9f32e20

Please sign in to comment.