diff --git a/src/contracts/elGamalHE.ts b/src/contracts/elGamalHE.ts index 57416dea..0d819da0 100644 --- a/src/contracts/elGamalHE.ts +++ b/src/contracts/elGamalHE.ts @@ -1,4 +1,4 @@ -import { SmartContract, assert, method, prop } from 'scrypt-ts' +import { SmartContract, assert, hash256, method, prop } from 'scrypt-ts' import { Point, SECP256K1 } from 'scrypt-ts-lib' export type CT = { @@ -19,7 +19,11 @@ export class ElGamalHE extends SmartContract { public add(toAdd: CT) { // Add encrypted value to the total sum. this.salarySum = ElGamalHE.addCT(this.salarySum, toAdd) - assert(true) + + const outputs = + this.buildStateOutput(this.ctx.utxo.value) + + this.buildChangeOutput() + assert(hash256(outputs) == this.ctx.hashOutputs, 'hashOutputs mismatch') } @method() diff --git a/src/contracts/paillierHE.ts b/src/contracts/paillierHE.ts index 9c16c725..ed98ce26 100644 --- a/src/contracts/paillierHE.ts +++ b/src/contracts/paillierHE.ts @@ -1,4 +1,4 @@ -import { SmartContract, assert, method, prop } from 'scrypt-ts' +import { SmartContract, assert, hash256, method, prop } from 'scrypt-ts' export class PaillierHE extends SmartContract { // max # of bits for e = ceil(log2(n)) @@ -21,13 +21,21 @@ export class PaillierHE extends SmartContract { @method() public add(toAdd: bigint) { this.x = PaillierHE.addCT(this.x, toAdd, this.nSquare) - assert(true) + + const outputs = + this.buildStateOutput(this.ctx.utxo.value) + + this.buildChangeOutput() + assert(hash256(outputs) == this.ctx.hashOutputs, 'hashOutputs mismatch') } @method() public mul(factor: bigint) { this.x = PaillierHE.mulCT(this.x, factor, this.nSquare) - assert(true) + + const outputs = + this.buildStateOutput(this.ctx.utxo.value) + + this.buildChangeOutput() + assert(hash256(outputs) == this.ctx.hashOutputs, 'hashOutputs mismatch') } @method() diff --git a/tests/paillierHE.test.ts b/tests/paillierHE.test.ts index 26ee22b1..939276fd 100644 --- a/tests/paillierHE.test.ts +++ b/tests/paillierHE.test.ts @@ -27,7 +27,7 @@ describe('Heavy: Test SmartContract `PaillierHE`', () => { let currentInstance = instance for (let i = 0; i < 5; ++i) { - const nextInstance = currentInstance.next() + let nextInstance = currentInstance.next() // Add encrypted amount (100) to the contract commulative value. const toAdd = publicKey.encrypt(100n) @@ -49,6 +49,8 @@ describe('Heavy: Test SmartContract `PaillierHE`', () => { currentInstance = nextInstance + nextInstance = currentInstance.next() + // Multiply encrypted amount. const k = 5n nextInstance.x = PaillierHE.mulCT(