Skip to content

Commit

Permalink
more encryption tests
Browse files Browse the repository at this point in the history
  • Loading branch information
XuyangSong committed Nov 22, 2024
1 parent 281076c commit 65d246c
Showing 3 changed files with 57 additions and 24 deletions.
38 changes: 17 additions & 21 deletions native/cairo_prover/src/encryption.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::error::CairoError;
use crate::{error::CairoError, utils::bytes_to_felt_vec};
use starknet_crypto::{poseidon_hash, poseidon_hash_many};
use starknet_curve::curve_params::GENERATOR;
use starknet_types_core::{
@@ -74,20 +74,18 @@ impl Ciphertext {
// Add encrypt_nonce
cipher.push(*encrypt_nonce);

Ok(cipher.into())
let ret: [Felt; CIPHERTEXT_NUM] = cipher
.try_into()
.map_err(|_| CairoError::InvalidCiphertextLength)?;

Ok(Self(ret))
}

pub fn decrypt(&self, sk: &Felt) -> Result<Vec<Felt>, CairoError> {
let cipher_text = self.inner();
let cipher_len = cipher_text.len();
if cipher_len != CIPHERTEXT_NUM {
return Err(CairoError::InvalidCiphertextLength);
}

let mac = cipher_text[CIPHERTEXT_MAC];
let pk_x = cipher_text[CIPHERTEXT_PK_X];
let pk_y = cipher_text[CIPHERTEXT_PK_Y];
let encrypt_nonce = cipher_text[CIPHERTEXT_NONCE];
let mac = self.inner()[CIPHERTEXT_MAC];
let pk_x = self.inner()[CIPHERTEXT_PK_X];
let pk_y = self.inner()[CIPHERTEXT_PK_Y];
let encrypt_nonce = self.inner()[CIPHERTEXT_NONCE];

if let Ok(pk) = AffinePoint::new(pk_x, pk_y) {
// Generate the secret key
@@ -104,7 +102,7 @@ impl Ciphertext {

// Decrypt
let mut msg = vec![];
for cipher_element in &cipher_text[0..PLAINTEXT_NUM] {
for cipher_element in &self.inner()[0..PLAINTEXT_NUM] {
let msg_element = *cipher_element - poseidon_state;
msg.push(msg_element);
poseidon_state = poseidon_hash(*cipher_element, secret_key_x);
@@ -119,15 +117,13 @@ impl Ciphertext {
Err(CairoError::InvalidPublicKey)
}
}
}

impl From<Vec<Felt>> for Ciphertext {
fn from(input_vec: Vec<Felt>) -> Self {
Ciphertext(
input_vec
.try_into()
.expect("public input with incorrect length"),
)
pub fn from_bytes(input_vec: Vec<Vec<u8>>) -> Result<Self, CairoError> {
let cipher_felt = bytes_to_felt_vec(input_vec)?;
let cipher: [Felt; CIPHERTEXT_NUM] = cipher_felt
.try_into()
.map_err(|_| CairoError::InvalidCiphertextLength)?;
Ok(Self(cipher))
}
}

4 changes: 2 additions & 2 deletions native/cairo_prover/src/lib.rs
Original file line number Diff line number Diff line change
@@ -443,13 +443,13 @@ fn encrypt(
#[rustler::nif]
fn decrypt(cihper: Vec<Vec<u8>>, sk: Vec<u8>) -> NifResult<Vec<Vec<u8>>> {
// Decode messages
let cipher_felt = bytes_to_felt_vec(cihper)?;
let cipher = Ciphertext::from_bytes(cihper)?;

// Decode sk
let sk_felt = bytes_to_felt(sk)?;

// Encrypt
let plaintext = Ciphertext::from(cipher_felt).decrypt(&sk_felt)?;
let plaintext = cipher.decrypt(&sk_felt)?;
let plaintext_bytes = plaintext.iter().map(|x| x.to_bytes_be().to_vec()).collect();

Ok(plaintext_bytes)
39 changes: 38 additions & 1 deletion test/cairo_encryption.exs
Original file line number Diff line number Diff line change
@@ -36,8 +36,45 @@ defmodule NifTest do
assert Cairo.get_output(public_input) == expected_cipher

# decryption
plaintext = Cairo.decrypt(expected_cipher, felt_bytes_1)
plaintext = Cairo.decrypt(expected_cipher, sk)

assert plaintext == expected_plaintext

# decryption: wrong sk
assert {:error, "Invalid DH key"} = Cairo.decrypt(expected_cipher, felt_bytes_0)
end

test "cairo_encryption_invalid_input_test" do
felt_bytes = List.duplicate(1, 32)
plaintext = List.duplicate(felt_bytes, 10)
pk = Cairo.get_public_key(felt_bytes)
invalid_pk = List.duplicate(1, 64)
sk = felt_bytes
nonce = felt_bytes

assert {:error, "Invalid finite field: 32 bytes needed"} =
Cairo.encrypt([[]], pk, sk, nonce)

assert {:error, "Invalid Point"} = Cairo.encrypt(plaintext, [], sk, nonce)

assert {:error, "Invalid Point"} =
Cairo.encrypt(plaintext, invalid_pk, sk, nonce)

assert {:error, "Invalid finite field: 32 bytes needed"} =
Cairo.encrypt(plaintext, pk, [], nonce)

assert {:error, "Invalid finite field: 32 bytes needed"} =
Cairo.encrypt(plaintext, pk, sk, [])

cipher = List.duplicate(felt_bytes, 14)

assert {:error, "Invalid finite field: 32 bytes needed"} =
Cairo.decrypt([[]], sk)

assert {:error, "The length of ciphertext is not correct"} =
Cairo.decrypt([felt_bytes], sk)

assert {:error, "Invalid finite field: 32 bytes needed"} =
Cairo.decrypt(cipher, [])
end
end

0 comments on commit 65d246c

Please sign in to comment.