From 0319ca5568431eb91ad13c76d63141fd97093df9 Mon Sep 17 00:00:00 2001
From: Levent DEMIR <levent.demir@zama.ai>
Date: Wed, 27 Dec 2023 13:41:53 +0100
Subject: [PATCH] feature(decrypt): call KMS for decryption

IMPORTANT: the decrypt call should be replaced in most cases by cmux
It is still called in Governor example in fhevm repository.
To make tests passing we keep it for now.
In the future, it will be replaced by async decryption
For this early first version of KMS this is acceptable.
---
 fhevm/precompiles.go | 43 ++++++++++++++++++++++++++++++++++++++++++-
 fhevm/tfhe.go        |  8 ++++++--
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/fhevm/precompiles.go b/fhevm/precompiles.go
index 44d53a9..8740785 100644
--- a/fhevm/precompiles.go
+++ b/fhevm/precompiles.go
@@ -2025,11 +2025,52 @@ func decryptRun(environment EVMEnvironment, caller common.Address, addr common.A
 	} else if !optReqResult {
 		return nil, ErrExecutionReverted
 	}
-	plaintext, err := decryptValue(ct.ciphertext)
+
+	var fheType kms.FheType
+	switch ct.ciphertext.fheUintType {
+	case FheUint8:
+		fheType = kms.FheType_Euint8
+	case FheUint16:
+		fheType = kms.FheType_Euint16
+	case FheUint32:
+		fheType = kms.FheType_Euint32
+	}
+
+	pubKey := input[32:64]
+
+	// TODO: generate merkle proof for some data
+	proof := &kms.Proof{
+		Height:              4,
+		MerklePatriciaProof: []byte{},
+	}
+
+	decryptionRequest := &kms.DecryptionRequest{
+		FheType:    fheType,
+		Ciphertext: ct.ciphertext.serialization,
+		Request:    pubKey, // TODO: change according to the structure of `Request`
+		Proof:      proof,
+	}
+
+	conn, err := grpc.Dial(kms.KmsEndpointAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return nil, errors.New("kms unreachable")
+	}
+	defer conn.Close()
+
+	ep := kms.NewKmsEndpointClient(conn)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	res, err := ep.Decrypt(ctx, decryptionRequest)
 	if err != nil {
 		logger.Error("decrypt failed", "err", err)
 		return nil, err
 	}
+
+	var plaintext = uint64(res.Plaintext)
+	logger.Info("decrypt success", "plaintext", plaintext)
+
 	// Always return a 32-byte big-endian integer.
 	ret := make([]byte, 32)
 	bigIntValue := big.NewInt(0)
diff --git a/fhevm/tfhe.go b/fhevm/tfhe.go
index dd08e05..1eb705b 100644
--- a/fhevm/tfhe.go
+++ b/fhevm/tfhe.go
@@ -1574,7 +1574,7 @@ func InitGlobalKeysFromFiles(keysDir string) error {
 	var cksPath = path.Join(keysDir, "cks")
 	cksBytes, err := os.ReadFile(cksPath)
 	if err != nil {
-		return err
+		fmt.Println("INFO: cks not loaded from: " + keysDir)
 	}
 	var pksPath = path.Join(keysDir, "pks")
 	pksBytes, err := os.ReadFile(pksPath)
@@ -1587,7 +1587,11 @@ func InitGlobalKeysFromFiles(keysDir string) error {
 	pksHash = crypto.Keccak256Hash(pksBytes)
 	pks = C.deserialize_compact_public_key(toBufferView(pksBytes))
 
-	cks = C.deserialize_client_key(toBufferView(cksBytes))
+	// cks will be handled by the KMS from now on
+	// TODO: completely remove after KMS is well tested
+	if len(cksBytes) > 0 {
+		cks = C.deserialize_client_key(toBufferView(cksBytes))
+	}
 
 	initCiphertextSizes()