fix: proof verification when offset != 0

Author: adiwajshing

js/package-lock.json

@@ -46,6 +46,7 @@
     }
   },
   "dependencies": {
+    "@stablelib/chacha20poly1305": "^1.0.0",
     "js-base64": "^3.7.7"
   },
   "devDependencies": {

js/package.json

@@ -1,10 +1,14 @@
+import { ChaCha20Poly1305 } from '@stablelib/chacha20poly1305'
+import { subtle } from 'crypto'
+import { AlgorithmConfig, EncryptionAlgorithm } from './types'
 import { bitsToUint8Array, bitsToUintArray, toUint8Array, toUintArray, uint8ArrayToBits, uintArrayToBits } from './utils'
 
 // commit hash for this repo
 export const GIT_COMMIT_HASH = 'd82fab41fa4033aa13feda3374f80a9df7af52b2'
 
-export const CONFIG = {
+export const CONFIG: { [E in EncryptionAlgorithm]: AlgorithmConfig } = {
 	'chacha20': {
+		index: 0,
 		chunkSize: 16,
 		bitsPerWord: 32,
 		keySizeBytes: 32,
@@ -15,14 +19,20 @@ export const CONFIG = {
 		// chacha20 circuit uses LE encoding
 		isLittleEndian: true,
 		uint8ArrayToBits: (arr: Uint8Array) => (
-			uintArrayToBits(toUintArray(arr))
+			uintArrayToBits(toUintArray(arr)).flat()
 		),
 		bitsToUint8Array: (bits: number[]) => {
 			const arr = bitsToUintArray(bits)
 			return toUint8Array(arr)
 		},
+		encrypt({ key, iv, in: data }) {
+			const cipher = new ChaCha20Poly1305(key)
+			const ciphertext = cipher.seal(iv, data)
+			return ciphertext.slice(0, data.length)
+		},
 	},
 	'aes-256-ctr': {
+		index: 1,
 		chunkSize: 64,
 		bitsPerWord: 8,
 		keySizeBytes: 32,
@@ -34,8 +44,10 @@ export const CONFIG = {
 		isLittleEndian: false,
 		uint8ArrayToBits,
 		bitsToUint8Array,
+		encrypt: makeAesCtr(256),
 	},
 	'aes-128-ctr': {
+		index: 2,
 		chunkSize: 64,
 		bitsPerWord: 8,
 		keySizeBytes: 16,
@@ -47,5 +59,23 @@ export const CONFIG = {
 		isLittleEndian: false,
 		uint8ArrayToBits,
 		bitsToUint8Array,
+		encrypt: makeAesCtr(128),
+	}
+}
+
+function makeAesCtr(keyLenBits: number): AlgorithmConfig['encrypt'] {
+	return async({ key, iv, in: inp }) => {
+		const keyImp = await subtle.importKey(
+			'raw', key,
+			{ name: 'AES-GCM', length: keyLenBits },
+			false,
+			['encrypt']
+		)
+		const buff = await subtle.encrypt(
+			{ name: 'AES-GCM', iv },
+			keyImp,
+			inp
+		)
+		return new Uint8Array(buff).slice(0, inp.length)
 	}
 }

js/src/config.ts

@@ -65,7 +65,8 @@ export function makeGnarkZkOperator({
 			return lib
 		}
 
-		const { id, ext } = ALGS_MAP[algorithm]
+		const { ext } = ALGS_MAP[algorithm]
+		const { index: id } = CONFIG[algorithm]
 		const [pk, r1cs] = await Promise.all([
 			fetcher.fetch('gnark', `pk.${ext}`, logger),
 			fetcher.fetch('gnark', `r1cs.${ext}`, logger),

js/src/gnark/operator.ts

@@ -13,14 +13,11 @@ export type GnarkLib = {
 }
 
 export const ALGS_MAP: {
-	[key in EncryptionAlgorithm]: {
-		id: number
-		ext: string
-	}
+	[key in EncryptionAlgorithm]: { ext: string }
 } = {
-	'chacha20': { id: 0, ext: 'chacha20' },
-	'aes-128-ctr': { id: 1, ext: 'aes128' },
-	'aes-256-ctr': { id: 2, ext: 'aes256' },
+	'chacha20': { ext: 'chacha20' },
+	'aes-128-ctr': { ext: 'aes128' },
+	'aes-256-ctr': { ext: 'aes256' },
 }
 
 // golang uses different arch names
@@ -99,8 +96,7 @@ export function strToUint8Array(str: string) {
 	return new TextEncoder().encode(str)
 }
 
-export
-function generateGnarkWitness(cipher: EncryptionAlgorithm, input) {
+export function generateGnarkWitness(cipher: EncryptionAlgorithm, input) {
 	const {
 		bitsToUint8Array,
 		isLittleEndian

js/src/gnark/utils.ts

@@ -42,7 +42,8 @@ export function makeSnarkJsZKOperator({
 	let wc: Promise<unknown> | undefined
 
 	return {
-		async generateWitness(input, logger) {
+		// eslint-disable-next-line @typescript-eslint/no-unused-vars
+		async generateWitness({ out, ...input }, logger) {
 			circuitWasm ||= getCircuitWasm()
 			wc ||= (async() => {
 				if(!snarkjs.wtns.getWtnsCalculator) {

js/src/snarkjs/operator.ts

@@ -83,20 +83,12 @@ describe.each(ALL_ALGOS)('%s Lib Tests', (algorithm) => {
 			)
 			const publicInput = { ciphertext, iv: iv, offset: 0 }
 
-
 			const proof = await generateProof({
 				algorithm,
 				privateInput,
 				publicInput,
 				operator
 			})
-			// ensure the ZK decrypted data matches the plaintext
-			expect(
-				proof.plaintext
-					.slice(0, plaintext.length)
-			).toEqual(
-				plaintext
-			)
 			// client will send proof to witness
 			// witness would verify proof
 			await verifyProof({ proof, publicInput, operator })
@@ -106,8 +98,6 @@ describe.each(ALL_ALGOS)('%s Lib Tests', (algorithm) => {
 			const totalPlaintext = new Uint8Array(randomBytes(chunkSizeBytes * 5))
 			// use a chunk in the middle
 			const offset = 2
-			const plaintext = totalPlaintext
-				.subarray(chunkSizeBytes * offset, chunkSizeBytes * (offset + 1))
 
 			const iv = Buffer.alloc(12, 3)
 			const privateInput: PrivateInput = {
@@ -130,13 +120,8 @@ describe.each(ALL_ALGOS)('%s Lib Tests', (algorithm) => {
 				publicInput,
 				operator
 			})
-			// ensure the ZK decrypted data matches the plaintext
-			expect(
-				proof.plaintext
-					.slice(0, plaintext.length)
-			).toEqual(
-				plaintext
-			)
+
+			await verifyProof({ proof, publicInput, operator })
 		})
 
 		it('should fail to verify incorrect data', async() => {
@@ -146,7 +131,6 @@ describe.each(ALL_ALGOS)('%s Lib Tests', (algorithm) => {
 				key: Buffer.alloc(keySizeBytes, 2),
 			}
 
-
 			const iv = Buffer.alloc(12, 3)
 			const ciphertext = encryptData(
 				algorithm,

js/src/tests/lib.test.ts

@@ -58,8 +58,7 @@ export type GenerateProofOpts = {
 	 */
 	algorithm: EncryptionAlgorithm
 	/**
-	 * private input to the circuit
-	 * * will include the key, iv, and counter
+	 * private input to the circuit (i.e. key)
 	 */
 	privateInput: PrivateInput
 	/**
@@ -75,18 +74,42 @@ export type GenerateProofOpts = {
 }
 
 export type VerifyProofOpts = {
-	/** JSON proof generated by "generateProof" fn */
+	/**
+	 * JSON proof generated by "generateProof" fn
+	 * with the plaintext
+	 * */
 	proof: Proof
 	publicInput: PublicInput
 	operator: ZKOperator
 	logger?: Logger
 }
 
+export type AlgorithmConfig = {
+	index: number
+	chunkSize: number
+	bitsPerWord: number
+	keySizeBytes: number
+	ivSizeBytes: number
+	startCounter: number
+	blocksPerChunk: number
+	isLittleEndian: boolean
+	uint8ArrayToBits: (arr: Uint8Array) => number[]
+	bitsToUint8Array: (bits: number[]) => Uint8Array
+	/**
+	 * Encrypt some ciphertext with the given key and IV
+	 */
+	encrypt(opts: {
+		key: Uint8Array
+		iv: Uint8Array
+		in: Uint8Array
+	}): Promise<Uint8Array> | Uint8Array
+}
+
 type ZKProof = { [_: string]: unknown } | string
 
 type ZKProofOutput = {
 	proof: ZKProof
-	publicSignals: number[]
+	publicSignals?: number[]
 }
 
 type ZKInputItem = number[] | number[][]
@@ -96,6 +119,7 @@ type ZKProofInput = {
 	nonce: ZKInputItem
 	counter: ZKInputItem
 	in: ZKInputItem
+	out: ZKInputItem
 }
 
 /**
@@ -137,7 +161,6 @@ export type PrivateInput = {
 export type PublicInput = {
 	/** the ciphertext to decrypt */
 	ciphertext: Uint8Array
-
 	/** 192 bit IV for the ciphertext decryption */
 	iv: Uint8Array
 	/**