1#![allow(incomplete_features)]
  2#![feature(generic_const_exprs)]
  3
  4pub mod circuit;
  5
  6use std::array;
  7
  8use plonky2::field::types::Field;
  9
 10/// Feistel-type cipher for finite fields, where XORing is replaced
 11/// with addition, cf. https://en.wikipedia.org/wiki/Feistel_cipher or
 12/// §3.5.1 of Stinson's *Cryptography* (3rd edition).
 13pub fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize, F: Field>(
 14    state: [F; 2 * STATE_HALF_LEN],
 15    key_schedule: &[[F; KEY_LEN]],
 16    f: &impl Fn(Vec<F>) -> [F; STATE_HALF_LEN],
 17) -> [F; 2 * STATE_HALF_LEN] {
 18    if key_schedule.is_empty() {
 19        state
 20    } else {
 21        let k = key_schedule[0];
 22        let l = &state[..STATE_HALF_LEN];
 23        let r = &state[STATE_HALF_LEN..];
 24
 25        let new_l = r;
 26        let round_offset = f([r, &k].concat());
 27        let new_r: [_; STATE_HALF_LEN] = array::from_fn(|i| l[i] + round_offset[i]);
 28
 29        feistel_cipher(
 30            array::from_fn(|i| {
 31                if i < STATE_HALF_LEN {
 32                    new_l[i]
 33                } else {
 34                    new_r[i - STATE_HALF_LEN]
 35                }
 36            }),
 37            &key_schedule[1..],
 38            f,
 39        )
 40    }
 41}
 42
 43/// Feistel-type inverse cipher for finite fields, where XORing is
 44/// replaced with subtraction.
 45pub fn feistel_inv_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize, F: Field>(
 46    state: [F; 2 * STATE_HALF_LEN],
 47    reverse_key_schedule: &[[F; KEY_LEN]],
 48    f: &impl Fn(Vec<F>) -> [F; STATE_HALF_LEN],
 49) -> [F; 2 * STATE_HALF_LEN] {
 50    if reverse_key_schedule.is_empty() {
 51        state
 52    } else {
 53        let k = reverse_key_schedule[0];
 54        let l = &state[..STATE_HALF_LEN];
 55        let r = &state[STATE_HALF_LEN..];
 56
 57        let new_r = l;
 58        let round_offset = f([l, &k].concat());
 59        let new_l: [_; STATE_HALF_LEN] = array::from_fn(|i| r[i] - round_offset[i]);
 60
 61        feistel_inv_cipher(
 62            array::from_fn(|i| {
 63                if i < STATE_HALF_LEN {
 64                    new_l[i]
 65                } else {
 66                    new_r[i - STATE_HALF_LEN]
 67                }
 68            }),
 69            &reverse_key_schedule[1..],
 70            f,
 71        )
 72    }
 73}
 74
 75#[cfg(test)]
 76mod tests {
 77    use plonky2::{
 78        field::{
 79            goldilocks_field::GoldilocksField as F,
 80            types::{Field, Field64},
 81        },
 82        hash::{
 83            hashing::hash_n_to_hash_no_pad, keccak::KeccakPermutation,
 84            poseidon::PoseidonPermutation,
 85        },
 86    };
 87    use rand::{Rng, rngs::OsRng};
 88
 89    use crate::{feistel_cipher, feistel_inv_cipher};
 90
 91    pub(crate) fn random_fields<const N: usize>() -> [F; N] {
 92        std::array::from_fn(|_| F::from_canonical_u64(OsRng.gen_range(0..F::ORDER)))
 93    }
 94
 95    #[test]
 96    fn feistel_keccak_poseidon_roundtrip() {
 97        const NR: usize = 32;
 98        const STATE_HALF_LEN: usize = 4;
 99        const KEY_LEN: usize = 4;
100
101        let f_keccak =
102            |fields: Vec<_>| hash_n_to_hash_no_pad::<F, KeccakPermutation<_>>(&fields).elements;
103        let f_poseidon =
104            |fields: Vec<_>| hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&fields).elements;
105
106        [f_keccak, f_poseidon].into_iter().for_each(|f| {
107            (0..10).for_each(|_| {
108                let state: [_; 2 * STATE_HALF_LEN] = random_fields();
109                let key_schedule: [[_; KEY_LEN]; NR] = std::array::from_fn(|_| random_fields());
110
111                let out_state = feistel_cipher(state, &key_schedule, &f);
112                let reverse_key_schedule = key_schedule.into_iter().rev().collect::<Vec<_>>();
113                let maybe_original_state = feistel_inv_cipher(out_state, &reverse_key_schedule, &f);
114
115                assert_eq!(state, maybe_original_state);
116            })
117        })
118    }
119}