From 7084f9342d67c45d982f8d29570e1af3ceea1c1f Mon Sep 17 00:00:00 2001
From: kenta-mori3322 <kenta.mori32@gmail.com>
Date: Mon, 27 May 2024 03:20:41 +0000
Subject: [PATCH] chore: add tee.decrypt

---
 fhevm/tee_crypto.go      |  4 ++--
 fhevm/tee_crypto_test.go |  5 ++++-
 tee/tee_mock.go          | 32 ++++++++++++++++++++++++++++++
 tee/tee_mock_test.go     | 42 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/fhevm/tee_crypto.go b/fhevm/tee_crypto.go
index 7a2dc4e..4f19bf2 100644
--- a/fhevm/tee_crypto.go
+++ b/fhevm/tee_crypto.go
@@ -137,7 +137,7 @@ func teeEvaluateRemainingOptimisticRequires(environment EVMEnvironment) (bool, e
 		var cumulative *tfhe.TfheCiphertext = requires[0]
 		var err error
 		for i := 1; i < length; i++ {
-			cumulative, err = cumulative.Bitand(requires[i])
+			cumulative, err = tee.BitAnd(cumulative, requires[i])
 			if err != nil {
 				environment.GetLogger().Error("evaluateRemainingOptimisticRequires bitand failed", "err", err)
 				return false, err
@@ -153,7 +153,7 @@ func teeEvaluateRemainingOptimisticRequires(environment EVMEnvironment) (bool, e
 		ret := make([]byte, 32)
 		copy(ret[32-len(plaintext):], plaintext)
 
-		retVal := *new(big.Int).SetBytes(ret)
+		retVal := new(big.Int).SetBytes(ret)
 		return retVal.Uint64() != 0, err
 	}
 	return true, nil
diff --git a/fhevm/tee_crypto_test.go b/fhevm/tee_crypto_test.go
index 46f7bcb..9a74a43 100644
--- a/fhevm/tee_crypto_test.go
+++ b/fhevm/tee_crypto_test.go
@@ -47,7 +47,7 @@ func TestTeeDecryptRun(t *testing.T) {
 	})
 }
 
-func TestTeeOptimisticRequire(t *testing.T) {
+func TestTeeOptimisticRequireSuccess(t *testing.T) {
 	signature := "teeOptimisticRequire(uint256)"
 	rapid.Check(t, func(t *rapid.T) {
 		testcases := []struct {
@@ -81,7 +81,10 @@ func TestTeeOptimisticRequire(t *testing.T) {
 			t.Fatalf("incorrect result, expected=%d, got=0", 1)
 		}
 	})
+}
 
+func TestTeeOptimisticRequireFail(t *testing.T) {
+	signature := "teeOptimisticRequire(uint256)"
 	rapid.Check(t, func(t *rapid.T) {
 		testcases := []struct {
 			typ     tfhe.FheUintType
diff --git a/tee/tee_mock.go b/tee/tee_mock.go
index 8b59065..a6998b6 100644
--- a/tee/tee_mock.go
+++ b/tee/tee_mock.go
@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -118,3 +119,34 @@ func Decrypt(ct *tfhe.TfheCiphertext) (TeePlaintext, error) {
 
 	return plaintext, nil
 }
+
+// TODO:
+// Need to be reviewed by Amaury and Elmer as we already have teeBitAndRun in tee_bit.go
+// As I couldn't find BitOperator with 2 ciphertext in tee_bit.go, defining a new function here.
+// Decrypt two ciphertexts in parameter and do bitAnd operation and then reencrypt it.
+// Return a new ciphertext.
+func BitAnd(lt *tfhe.TfheCiphertext, rt *tfhe.TfheCiphertext) (*tfhe.TfheCiphertext, error) {
+	lp, err := Decrypt(lt)
+	if err != nil {
+		return nil, err
+	}
+
+	rp, err := Decrypt(rt)
+	if err != nil {
+		return nil, err
+	}
+
+	l := big.NewInt(0).SetBytes(lp.Value).Uint64()
+	r := big.NewInt(0).SetBytes(rp.Value).Uint64()
+	result := l & r
+
+	// Re-encrypt
+	resultBz := []byte{byte(result)}
+	teePlaintext := NewTeePlaintext(resultBz, tfhe.FheUint8, common.Address{})
+	ct, err := Encrypt(teePlaintext)
+	if err != nil {
+		return nil, err
+	}
+
+	return &ct, nil
+}
diff --git a/tee/tee_mock_test.go b/tee/tee_mock_test.go
index 90581c6..67d445b 100644
--- a/tee/tee_mock_test.go
+++ b/tee/tee_mock_test.go
@@ -2,6 +2,7 @@ package tee_test
 
 import (
 	"bytes"
+	"math/big"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -65,3 +66,44 @@ func TestUniqueCiphertexts(t *testing.T) {
 		}
 	})
 }
+
+func TestBitAndCiphertext(t *testing.T) {
+	rapid.Check(t, func(t *rapid.T) {
+		//
+		one := big.NewInt(1).Bytes()
+		zero := big.NewInt(0).Bytes()
+		teePlaintextOne := tee.NewTeePlaintext(one, tfhe.FheUint8, common.Address{})
+		teePlaintextZero := tee.NewTeePlaintext(zero, tfhe.FheUint8, common.Address{})
+		ctOne, err := tee.Encrypt(teePlaintextOne)
+		if err != nil {
+			t.Fatal(err)
+		}
+		ctZero, err := tee.Encrypt(teePlaintextZero)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Encrypt twice the same plaintext
+		ctResult, err := tee.BitAnd(&ctOne, &ctZero)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		result, err := tee.Decrypt(ctResult)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		plaintext := result.Value
+		// Always return a 32-byte big-endian integer.
+		ret := make([]byte, 32)
+		copy(ret[32-len(plaintext):], plaintext)
+
+		retVal := new(big.Int).SetBytes(ret)
+
+		// Make sure the ciphertexts are different
+		if retVal.Int64() != big.NewInt(0).Int64() {
+			t.Fatalf("expected different result, got %d", retVal.Int64())
+		}
+	})
+}