From ae84dc636b0503341cd92678113a012e552eece7 Mon Sep 17 00:00:00 2001
From: Petar Ivanov <29689712+dartdart26@users.noreply.github.com>
Date: Wed, 5 Jun 2024 18:19:49 +0300
Subject: [PATCH] fix: invalid input to C.u256_from_big_endian_bytes

An invalid input in bigIntToU256 was causing a segfault. Also add a test
for success and failure due to big inputs.
---
 fhevm/tfhe/tfhe_test.go     | 24 ++++++++++++++++++++++++
 fhevm/tfhe/tfhe_wrappers.go | 18 +++++++-----------
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/fhevm/tfhe/tfhe_test.go b/fhevm/tfhe/tfhe_test.go
index ed0fa8f..7ecb4ef 100644
--- a/fhevm/tfhe/tfhe_test.go
+++ b/fhevm/tfhe/tfhe_test.go
@@ -25,6 +25,30 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+func TestBigIntToCU256(t *testing.T) {
+	i := big.NewInt(0)
+	b := make([]byte, 32)
+	b[0] = 42
+	b[31] = 43
+	i.SetBytes(b)
+	_, err := bigIntToU256(i)
+	if err != nil {
+		t.Fatalf("bigIntToU256 must have succeeded")
+	}
+}
+
+func TestBigIntToCU256BigInput(t *testing.T) {
+	i := big.NewInt(0)
+	b := make([]byte, 33)
+	b[0] = 42
+	b[32] = 43
+	i.SetBytes(b)
+	_, err := bigIntToU256(i)
+	if err == nil {
+		t.Fatalf("bigIntToU256 must have failed on big input")
+	}
+}
+
 func TfheEncryptDecrypt(t *testing.T, fheUintType FheUintType) {
 	var val big.Int
 	switch fheUintType {
diff --git a/fhevm/tfhe/tfhe_wrappers.go b/fhevm/tfhe/tfhe_wrappers.go
index 2fc2c8c..1eb39b4 100644
--- a/fhevm/tfhe/tfhe_wrappers.go
+++ b/fhevm/tfhe/tfhe_wrappers.go
@@ -104,23 +104,19 @@ func EncryptAndSerializeCompact(value uint64, fheUintType FheUintType) []byte {
 	return ser
 }
 
-// bigIntToU256 uses u256_from_big_endian_bytes to convert big.Int to U256
+// Convert a big.Int `value` to C.U256.
+// Note: If `value` is bigger than 32 bytes an error is returned.
 func bigIntToU256(value *big.Int) (*C.U256, error) {
-	// Convert big.Int to 32-byte big-endian slice
-	bytes := value.Bytes()
-	if len(bytes) > 32 {
-		return nil, fmt.Errorf("big.Int too large for U256")
+	if len(value.Bytes()) > 32 {
+		return nil, fmt.Errorf("big.Int value is more than 32 bytes")
 	}
-	paddedBytes := make([]byte, 32-len(bytes)) // Padding
-	paddedBytes = append(paddedBytes, bytes...)
-
+	bytes := make([]byte, 32)
+	value.FillBytes(bytes)
 	var result C.U256
-
-	_, err := C.u256_from_big_endian_bytes((*C.uint8_t)(unsafe.Pointer(&paddedBytes[0])), C.size_t(32), &result)
+	_, err := C.u256_from_big_endian_bytes((*C.uint8_t)(unsafe.Pointer(&bytes[0])), C.size_t(32), &result)
 	if err != nil {
 		return nil, fmt.Errorf("failed to convert big.Int to U256: %v", err)
 	}
-
 	return &result, nil
 }