From e179abd0fd43cdfc4d04fa5705652f0206b5a2e9 Mon Sep 17 00:00:00 2001
From: Deepak Maram <deepak@mystenlabs.com>
Date: Mon, 12 Aug 2024 10:49:07 -0400
Subject: [PATCH] [zkLogin] Check that bigints fit into the field for extra
 safety (#16891)

## Description

Turns out we don't check anywhere that the inputs fit into the field and
neither does the poseidon-lite library (it just truncates elements and
fits them into the field).

Instrumenting poseidon seems like the simplest way to do this. The other
option (which I have not taken) is to instrument as soon as we receive
some inputs..

## Test Plan

How did you test the new or updated feature?

---
If your changes are not user-facing and do not break anything, you can
skip the following section. Otherwise, please briefly describe what has
changed under the Release Notes section.

### 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
---
 .changeset/cuddly-days-wash.md    |  5 +++++
 sdk/zklogin/src/poseidon.ts       | 10 ++++++++++
 sdk/zklogin/test/poseidon.test.ts | 25 +++++++++++++++++++++++++
 3 files changed, 40 insertions(+)
 create mode 100644 .changeset/cuddly-days-wash.md
 create mode 100644 sdk/zklogin/test/poseidon.test.ts

diff --git a/.changeset/cuddly-days-wash.md b/.changeset/cuddly-days-wash.md
new file mode 100644
index 0000000000000..00efc86baca81
--- /dev/null
+++ b/.changeset/cuddly-days-wash.md
@@ -0,0 +1,5 @@
+---
+'@mysten/zklogin': patch
+---
+
+Additional check for bigints
diff --git a/sdk/zklogin/src/poseidon.ts b/sdk/zklogin/src/poseidon.ts
index f51ce896137cd..c002176e7df3a 100644
--- a/sdk/zklogin/src/poseidon.ts
+++ b/sdk/zklogin/src/poseidon.ts
@@ -39,7 +39,17 @@ const poseidonNumToHashFN = [
 	poseidon16,
 ];
 
+export const BN254_FIELD_SIZE =
+	21888242871839275222246405745257275088548364400416034343698204186575808495617n;
+
 export function poseidonHash(inputs: (number | bigint | string)[]): bigint {
+	inputs.forEach((x) => {
+		const b = BigInt(x);
+		if (b < 0 || b >= BN254_FIELD_SIZE) {
+			throw new Error(`Element ${b} not in the BN254 field`);
+		}
+	});
+
 	const hashFN = poseidonNumToHashFN[inputs.length - 1];
 
 	if (hashFN) {
diff --git a/sdk/zklogin/test/poseidon.test.ts b/sdk/zklogin/test/poseidon.test.ts
new file mode 100644
index 0000000000000..7499b7164e5b4
--- /dev/null
+++ b/sdk/zklogin/test/poseidon.test.ts
@@ -0,0 +1,25 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+import { expect, test } from 'vitest';
+
+import { BN254_FIELD_SIZE, poseidonHash } from '../src/poseidon';
+
+test('can hash single input', () => {
+	const result = poseidonHash([123]);
+	expect(result).toBeTypeOf('bigint');
+});
+
+test('can hash multiple inputs', () => {
+	const result = poseidonHash([1, 2, 3, 4, 5]);
+	expect(result).toBeTypeOf('bigint');
+});
+
+test('throws error for invalid input', () => {
+	expect(() => poseidonHash([-1])).toThrowError('Element -1 not in the BN254 field');
+});
+
+test('throws error for invalid input greater than BN254_FIELD_SIZE', () => {
+	expect(() => poseidonHash([BN254_FIELD_SIZE])).toThrowError(
+		'Element 21888242871839275222246405745257275088548364400416034343698204186575808495617 not in the BN254 field',
+	);
+});