Merge pull request ethereum#56 from zama-ai/petar/dynamic-ciphertext-…

…size-and-msg-modulus

Determine ciphertext size and msg modulus from keys

core/vm/contracts.go

@@ -1343,8 +1343,8 @@ func (e *verifyCiphertext) Run(accessibleState PrecompileAccessibleState, caller
 	}
 	var ctBytes []byte
 	if !tomlConfig.Zk.Verify {
-		// For testing: if input size <= `ciphertextSize`, treat the whole input as ciphertext.
-		ctBytes = input[0:minInt(ciphertextSize, len(input))]
+		// For testing: if input size <= `fheCiphertextSize`, treat the whole input as ciphertext.
+		ctBytes = input[0:minInt(fheCiphertextSize, len(input))]
 	} else {
 		ctBytes = verifyZkProof(input)
 		if ctBytes == nil {
@@ -1809,16 +1809,14 @@ func (e *fheRand) Run(accessibleState PrecompileAccessibleState, caller common.A
 		return nil, err
 	}
 
-	// TODO: The amount of bytes depends on the ciphertext parameters. For now, get 1 random byte only
-	// and do mod 8 for use with MSG3/CARRY3 plaintexts.
-	randBytes := make([]byte, 1)
 	// XOR a byte array of 0s with the stream from the cipher and receive the result in the same array.
+	randBytes := make([]byte, 8)
 	cipher.XORKeyStream(randBytes, randBytes)
-	randValue := uint64(randBytes[0]) % 8
 
-	// Trivially encrypt the random value.
+	// Trivially encrypt the random integer.
+	randInt := binary.BigEndian.Uint64(randBytes) % fheMessageModulus
 	randCt := new(tfheCiphertext)
-	randCt.trivialEncrypt(randValue)
+	randCt.trivialEncrypt(randInt)
 	verifiedCiphertext := &verifiedCiphertext{
 		depth:      accessibleState.Interpreter().evm.depth,
 		ciphertext: randCt,

core/vm/instructions.go

@@ -612,7 +612,7 @@ func persistIfVerifiedCiphertext(val common.Hash, protectedStorage common.Addres
 	if metadataInt.IsZero() {
 		// If no metadata, it means this ciphertext itself hasn't been persisted to protected storage yet. We do that as part of SSTORE.
 		metadata.refCount = 1
-		metadata.length = uint64(ciphertextSize)
+		metadata.length = uint64(fheCiphertextSize)
 		ciphertextSlot := newInt(val.Bytes())
 		ciphertextSlot.AddUint64(ciphertextSlot, 1)
 		ctPart32 := make([]byte, 32)

core/vm/instructions_test.go

@@ -713,7 +713,7 @@ func (c testCallerAddress) Address() common.Address {
 func newTestScopeConext() *ScopeContext {
 	c := new(ScopeContext)
 	c.Memory = NewMemory()
-	c.Memory.Resize(ciphertextSize * 3)
+	c.Memory.Resize(uint64(fheCiphertextSize) * 3)
 	c.Stack = newstack()
 	c.Contract = NewContract(testCallerAddress{}, testContractAddress{}, big.NewInt(10), 100000)
 	return c

core/vm/tfhe.go

@@ -153,6 +153,15 @@ void* trivial_encrypt(void* sks, uint64_t value) {
 	return ct;
 }
 
+size_t get_message_modulus(void* cks) {
+	size_t modulus = 0;
+
+	const int r = shortint_client_key_get_message_modulus(cks, &modulus);
+	assert(r == 0);
+
+	return modulus;
+}
+
 */
 import "C"
 import (
@@ -168,8 +177,6 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 )
 
-const ciphertextSize = 8248
-
 func toBufferView(in []byte) C.BufferView {
 	return C.BufferView{
 		pointer: (*C.uchar)(unsafe.Pointer(&in[0])),
@@ -185,6 +192,12 @@ func homeDir() string {
 	return home
 }
 
+// The TFHE ciphertext size, in bytes.
+var fheCiphertextSize int
+
+// The TFHE message modulus. Extracted from the `cks`.
+var fheMessageModulus uint64
+
 var sks unsafe.Pointer
 var cks unsafe.Pointer
 var networkKeysDir string
@@ -221,6 +234,14 @@ func init() {
 	sks = C.deserialize_server_key(toBufferView(sks_bytes))
 	cks = C.deserialize_client_key(toBufferView(cks_bytes))
 
+	// Use trivial encryption to determine the ciphertext size for the used parameters.
+	// Note: parameters are embedded in the client `cks` key.
+	ct := new(tfheCiphertext)
+	ct.trivialEncrypt(1)
+	fheCiphertextSize = len(ct.serialize())
+
+	fheMessageModulus = uint64(C.get_message_modulus(cks))
+
 	go runGc()
 }
 
@@ -261,7 +282,7 @@ func (ct *tfheCiphertext) makeRandom() {
 	if ct.initialized() {
 		panic("cannot make an existing ciphertext random")
 	}
-	ct.serialization = make([]byte, ciphertextSize)
+	ct.serialization = make([]byte, fheCiphertextSize)
 	rand.Read(ct.serialization)
 	ct.random = true
 }

core/vm/tfhe_test.go

@@ -184,7 +184,7 @@ func TestTfheTrivialAndEncryptedAdd(t *testing.T) {
 func TestTfheTrivialSerializeSize(t *testing.T) {
 	ct := new(tfheCiphertext)
 	ct.trivialEncrypt(2)
-	if len(ct.serialize()) != ciphertextSize {
+	if len(ct.serialize()) != fheCiphertextSize {
 		t.Fatalf("serialization of trivially encrypted unexpected size")
 	}
 }