kcptun-rust 1.1.0

//! Quantum Permutation Pad (QPP) - compatible with github.com/xtaci/qpp
//! Based on Kuang's Quantum Permutation Pad for quantum-resistant encryption.

use aes::Aes256;
use cipher::{BlockEncryptMut, KeyInit};
use cipher::generic_array::GenericArray;
use hmac::{Hmac, Mac};
use num_bigint::BigUint;
use pbkdf2::pbkdf2_hmac;
use sha1::Sha1;
use sha2::Sha256;
use std::convert::TryInto;

type HmacSha256 = Hmac<Sha256>;

const PAD_IDENTIFIER: &str = "QPP_";
const PM_SELECTOR_IDENTIFIER: &[u8] = b"PERMUTATION_MATRIX_SELECTOR";
const SHUFFLE_SALT: &[u8] = b"___QUANTUM_PERMUTATION_PAD_SHUFFLE_SALT___";
const PRNG_SALT: &[u8] = b"___QUANTUM_PERMUTATION_PAD_PRNG_SALT___";
const CHUNK_DERIVE_SALT: &[u8] = b"___QUANTUM_PERMUTATION_PAD_SEED_DERIVE___";
const PBKDF2_LOOPS: u32 = 128;
const CHUNK_DERIVE_LOOPS: u32 = 1024;
const PAD_SWITCH: u8 = 8;
const QUBITS: u8 = 8;
const MATRIX_BYTES: usize = 1 << QUBITS; // 256

/// PRNG state - xoshiro256** compatible with Go
#[derive(Clone)]
pub struct Rand {
    pub xoshiro: [u64; 4],
    pub seed64: u64,
    pub count: u8,
}

fn rol64(x: u64, k: i32) -> u64 {
    (x << k) | (x >> (64 - k))
}

/// xoshiro256** - must match Go's prng.go
fn xoshiro256ss(s: &mut [u64; 4]) -> u64 {
    let result = rol64(s[1].wrapping_mul(5), 7).wrapping_mul(9);
    let t = s[1] << 17;
    s[2] ^= s[0];
    s[3] ^= s[1];
    s[1] ^= s[2];
    s[0] ^= s[3];
    s[2] ^= t;
    s[3] = rol64(s[3], 45);
    result
}

/// Create PRNG from seed - matches Go CreatePRNG
pub fn create_prng(seed: &[u8]) -> Rand {
    let mut mac = <HmacSha256 as Mac>::new_from_slice(seed).expect("HMAC key");
    Mac::update(&mut mac, PM_SELECTOR_IDENTIFIER);
    let sum = mac.finalize().into_bytes();

    let mut xoshiro = [0u8; 32];
    pbkdf2_hmac::<Sha1>(&sum, PRNG_SALT, PBKDF2_LOOPS, &mut xoshiro);

    let mut state = [
        u64::from_le_bytes(xoshiro[0..8].try_into().unwrap()),
        u64::from_le_bytes(xoshiro[8..16].try_into().unwrap()),
        u64::from_le_bytes(xoshiro[16..24].try_into().unwrap()),
        u64::from_le_bytes(xoshiro[24..32].try_into().unwrap()),
    ];
    let seed64 = xoshiro256ss(&mut state);

    Rand {
        xoshiro: state,
        seed64,
        count: 0,
    }
}

/// Quantum Permutation Pad - encryption/decryption pads
pub struct QuantumPermutationPad {
    pads: Vec<u8>,
    rpads: Vec<u8>,
    num_pads: u16,
}

fn qpp_minimum_seed_length(qubits: u8) -> usize {
    // 2^qubits factorial in bits
    let n: usize = 1 << qubits;
    let mut log2_fact = 0f64;
    for i in 2..=n {
        log2_fact += (i as f64).log2();
    }
    let bit_len = log2_fact.ceil() as usize;
    let byte_len = (bit_len + 7) / 8;
    byte_len.max(1)
}

fn fill(pad: &mut [u8]) {
    for (i, b) in pad.iter_mut().enumerate() {
        *b = i as u8;
    }
}

fn reverse(pad: &[u8], rpad: &mut [u8]) {
    for (i, &p) in pad.iter().enumerate() {
        rpad[p as usize] = i as u8;
    }
}

fn seed_to_chunks(seed: &[u8], qubits: u8) -> Vec<Vec<u8>> {
    let min_len = qpp_minimum_seed_length(qubits);
    let seed = if seed.len() < 32 {
        let mut derived = vec![0u8; 32];
        pbkdf2_hmac::<Sha1>(seed, CHUNK_DERIVE_SALT, PBKDF2_LOOPS, &mut derived);
        derived
    } else {
        seed.to_vec()
    };

    let chunk_count = ((min_len + 31) / 32).max(1);
    let mut chunks = Vec::with_capacity(chunk_count);
    let mut seed_idx = 0;
    for _ in 0..chunk_count {
        let mut chunk = vec![0u8; 32];
        for j in 0..32 {
            chunk[j] = seed[seed_idx % seed.len()];
            seed_idx += 1;
        }
        let mut derived = vec![0u8; chunk.len()];
        pbkdf2_hmac::<Sha1>(&chunk, CHUNK_DERIVE_SALT, CHUNK_DERIVE_LOOPS, &mut derived);
        chunks.push(derived);
    }
    chunks
}

fn shuffle(chunk: &[u8], pad: &mut [u8], pad_id: u16, blocks: &mut [Aes256]) {
    let mut mac = <HmacSha256 as Mac>::new_from_slice(chunk).expect("HMAC key");
    let msg = format!("{}{:b}", PAD_IDENTIFIER, pad_id);
    Mac::update(&mut mac, msg.as_bytes());
    let sum = mac.finalize().into_bytes();
    let mut sum = sum.to_vec();

    for i in (1..pad.len()).rev() {
        for block in blocks.iter_mut() {
            for off in (0..sum.len()).step_by(16) {
                if off + 16 <= sum.len() {
                    let mut block_data: GenericArray<u8, _> =
                        GenericArray::clone_from_slice(&sum[off..off + 16]);
                    block.encrypt_block_mut(&mut block_data);
                    sum[off..off + 16].copy_from_slice(block_data.as_slice());
                }
            }
        }
        let bigrand = BigUint::from_bytes_be(&sum);
        let modulus = BigUint::from(i + 1);
        let rem = &bigrand % &modulus;
        let digits = rem.to_u64_digits();
        let j = digits.first().copied().unwrap_or(0) as usize;
        pad.swap(i, j);
    }
}

impl QuantumPermutationPad {
    /// Create new QPP with seed and number of pads (must match Go peer)
    pub fn new(seed: &[u8], num_pads: u16) -> Self {
        let chunks = seed_to_chunks(seed, QUBITS);
        let mut blocks: Vec<Aes256> = Vec::with_capacity(chunks.len());
        for chunk in &chunks {
            let mut aeskey = [0u8; 32];
            pbkdf2_hmac::<Sha1>(chunk, SHUFFLE_SALT, PBKDF2_LOOPS, &mut aeskey);
            blocks.push(Aes256::new_from_slice(&aeskey).expect("AES256"));
        }

        let mut pads = vec![0u8; num_pads as usize * MATRIX_BYTES];
        let mut rpads = vec![0u8; num_pads as usize * MATRIX_BYTES];

        for i in 0..num_pads {
            let pad_start = (i as usize) * MATRIX_BYTES;
            let pad_end = pad_start + MATRIX_BYTES;
            let (pad, rpad) = (
                &mut pads[pad_start..pad_end],
                &mut rpads[pad_start..pad_end],
            );
            fill(pad);
            shuffle(&chunks[i as usize % chunks.len()], pad, i, &mut blocks);
            reverse(pad, rpad);
        }

        QuantumPermutationPad {
            pads,
            rpads,
            num_pads,
        }
    }

    /// Encrypt data in place using the given PRNG (stream-compatible)
    pub fn encrypt_with_prng(&self, data: &mut [u8], rand: &mut Rand) {
        if data.is_empty() {
            return;
        }
        let size = data.len();
        let num_pads = self.num_pads;
        let mut r = rand.seed64;
        let mut count = rand.count;
        let [mut s0, mut s1, mut s2, mut s3] = rand.xoshiro;

        let mut offset = 0;

        // handle unaligned bytes (count != 0)
        if count != 0 {
            while offset < data.len() {
                let rr = (r >> (count << 3)) as u8;
                let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
                data[offset] = self.pads[base + (data[offset] ^ rr) as usize];
                count += 1;
                if count == PAD_SWITCH {
                    r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
                    let t = s1 << 17;
                    s2 ^= s0;
                    s3 ^= s1;
                    s1 ^= s2;
                    s0 ^= s3;
                    s2 ^= t;
                    s3 = rol64(s3, 45);
                    count = 0;
                    offset += 1;
                    break;
                }
                offset += 1;
            }
        }

        // 8-byte aligned blocks, 16 bytes at a time
        let mut idx = offset;
        let repeat = (data.len() - idx) >> 4;
        for _ in 0..repeat {
            {
                let d = &mut data[idx..idx + 16];
                let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
                d[0] = self.pads[base + (d[0] ^ r as u8) as usize];
                d[1] = self.pads[base + (d[1] ^ (r >> 8) as u8) as usize];
                d[2] = self.pads[base + (d[2] ^ (r >> 16) as u8) as usize];
                d[3] = self.pads[base + (d[3] ^ (r >> 24) as u8) as usize];
                d[4] = self.pads[base + (d[4] ^ (r >> 32) as u8) as usize];
                d[5] = self.pads[base + (d[5] ^ (r >> 40) as u8) as usize];
                d[6] = self.pads[base + (d[6] ^ (r >> 48) as u8) as usize];
                d[7] = self.pads[base + (d[7] ^ (r >> 56) as u8) as usize];
            }
            r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
            let t = s1 << 17;
            s2 ^= s0;
            s3 ^= s1;
            s1 ^= s2;
            s0 ^= s3;
            s2 ^= t;
            s3 = rol64(s3, 45);
            let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
            {
                let d = &mut data[idx + 8..idx + 16];
                d[0] = self.pads[base + (d[0] ^ r as u8) as usize];
                d[1] = self.pads[base + (d[1] ^ (r >> 8) as u8) as usize];
                d[2] = self.pads[base + (d[2] ^ (r >> 16) as u8) as usize];
                d[3] = self.pads[base + (d[3] ^ (r >> 24) as u8) as usize];
                d[4] = self.pads[base + (d[4] ^ (r >> 32) as u8) as usize];
                d[5] = self.pads[base + (d[5] ^ (r >> 40) as u8) as usize];
                d[6] = self.pads[base + (d[6] ^ (r >> 48) as u8) as usize];
                d[7] = self.pads[base + (d[7] ^ (r >> 56) as u8) as usize];
            }
            r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
            let t = s1 << 17;
            s2 ^= s0;
            s3 ^= s1;
            s1 ^= s2;
            s0 ^= s3;
            s2 ^= t;
            s3 = rol64(s3, 45);
            idx += 16;
        }

        // remaining 8-byte block
        if idx + 8 <= data.len() {
            let d = &mut data[idx..idx + 8];
            let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
            d[0] = self.pads[base + (d[0] ^ r as u8) as usize];
            d[1] = self.pads[base + (d[1] ^ (r >> 8) as u8) as usize];
            d[2] = self.pads[base + (d[2] ^ (r >> 16) as u8) as usize];
            d[3] = self.pads[base + (d[3] ^ (r >> 24) as u8) as usize];
            d[4] = self.pads[base + (d[4] ^ (r >> 32) as u8) as usize];
            d[5] = self.pads[base + (d[5] ^ (r >> 40) as u8) as usize];
            d[6] = self.pads[base + (d[6] ^ (r >> 48) as u8) as usize];
            d[7] = self.pads[base + (d[7] ^ (r >> 56) as u8) as usize];
            r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
            let t = s1 << 17;
            s2 ^= s0;
            s3 ^= s1;
            s1 ^= s2;
            s0 ^= s3;
            s2 ^= t;
            s3 = rol64(s3, 45);
            idx += 8;
        }

        // tail bytes
        for i in idx..data.len() {
            let rr = (r >> (count << 3)) as u8;
            let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
            data[i] = self.pads[base + (data[i] ^ rr) as usize];
            count += 1;
        }

        rand.xoshiro = [s0, s1, s2, s3];
        rand.seed64 = r;
        rand.count = ((rand.count as usize + size) & (PAD_SWITCH as usize - 1)) as u8;
    }

    /// Decrypt data in place using the given PRNG
    pub fn decrypt_with_prng(&self, data: &mut [u8], rand: &mut Rand) {
        if data.is_empty() {
            return;
        }
        let size = data.len();
        let num_pads = self.num_pads;
        let mut r = rand.seed64;
        let mut count = rand.count;
        let [mut s0, mut s1, mut s2, mut s3] = rand.xoshiro;

        let mut offset = 0;

        if count != 0 {
            while offset < data.len() {
                let rr = (r >> (count << 3)) as u8;
                let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
                data[offset] = self.rpads[base + data[offset] as usize] ^ rr;
                count += 1;
                if count == PAD_SWITCH {
                    r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
                    let t = s1 << 17;
                    s2 ^= s0;
                    s3 ^= s1;
                    s1 ^= s2;
                    s0 ^= s3;
                    s2 ^= t;
                    s3 = rol64(s3, 45);
                    count = 0;
                    offset += 1;
                    break;
                }
                offset += 1;
            }
        }

        let mut idx = offset;
        let repeat = (data.len() - idx) >> 4;
        for _ in 0..repeat {
            {
                let d = &mut data[idx..idx + 8];
                let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
                let (rr0, rr1, rr2, rr3) = (r as u8, (r >> 8) as u8, (r >> 16) as u8, (r >> 24) as u8);
                let (rr4, rr5, rr6, rr7) =
                    ((r >> 32) as u8, (r >> 40) as u8, (r >> 48) as u8, (r >> 56) as u8);
                d[0] = self.rpads[base + d[0] as usize] ^ rr0;
                d[1] = self.rpads[base + d[1] as usize] ^ rr1;
                d[2] = self.rpads[base + d[2] as usize] ^ rr2;
                d[3] = self.rpads[base + d[3] as usize] ^ rr3;
                d[4] = self.rpads[base + d[4] as usize] ^ rr4;
                d[5] = self.rpads[base + d[5] as usize] ^ rr5;
                d[6] = self.rpads[base + d[6] as usize] ^ rr6;
                d[7] = self.rpads[base + d[7] as usize] ^ rr7;
            }
            r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
            let t = s1 << 17;
            s2 ^= s0;
            s3 ^= s1;
            s1 ^= s2;
            s0 ^= s3;
            s2 ^= t;
            s3 = rol64(s3, 45);
            let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
            {
                let d = &mut data[idx + 8..idx + 16];
                let (rr0, rr1, rr2, rr3) = (r as u8, (r >> 8) as u8, (r >> 16) as u8, (r >> 24) as u8);
                let (rr4, rr5, rr6, rr7) =
                    ((r >> 32) as u8, (r >> 40) as u8, (r >> 48) as u8, (r >> 56) as u8);
                d[0] = self.rpads[base + d[0] as usize] ^ rr0;
                d[1] = self.rpads[base + d[1] as usize] ^ rr1;
                d[2] = self.rpads[base + d[2] as usize] ^ rr2;
                d[3] = self.rpads[base + d[3] as usize] ^ rr3;
                d[4] = self.rpads[base + d[4] as usize] ^ rr4;
                d[5] = self.rpads[base + d[5] as usize] ^ rr5;
                d[6] = self.rpads[base + d[6] as usize] ^ rr6;
                d[7] = self.rpads[base + d[7] as usize] ^ rr7;
            }
            r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
            let t = s1 << 17;
            s2 ^= s0;
            s3 ^= s1;
            s1 ^= s2;
            s0 ^= s3;
            s2 ^= t;
            s3 = rol64(s3, 45);
            idx += 16;
        }

        if idx + 8 <= data.len() {
            let d = &mut data[idx..idx + 8];
            let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
            let (rr0, rr1, rr2, rr3) = (r as u8, (r >> 8) as u8, (r >> 16) as u8, (r >> 24) as u8);
            let (rr4, rr5, rr6, rr7) =
                ((r >> 32) as u8, (r >> 40) as u8, (r >> 48) as u8, (r >> 56) as u8);
            d[0] = self.rpads[base + d[0] as usize] ^ rr0;
            d[1] = self.rpads[base + d[1] as usize] ^ rr1;
            d[2] = self.rpads[base + d[2] as usize] ^ rr2;
            d[3] = self.rpads[base + d[3] as usize] ^ rr3;
            d[4] = self.rpads[base + d[4] as usize] ^ rr4;
            d[5] = self.rpads[base + d[5] as usize] ^ rr5;
            d[6] = self.rpads[base + d[6] as usize] ^ rr6;
            d[7] = self.rpads[base + d[7] as usize] ^ rr7;
            r = rol64(s1.wrapping_mul(5), 7).wrapping_mul(9);
            let t = s1 << 17;
            s2 ^= s0;
            s3 ^= s1;
            s1 ^= s2;
            s0 ^= s3;
            s2 ^= t;
            s3 = rol64(s3, 45);
            idx += 8;
        }

        for i in idx..data.len() {
            let rr = (r >> (count << 3)) as u8;
            let base = (r as u16 % num_pads) as usize * MATRIX_BYTES;
            data[i] = self.rpads[base + data[i] as usize] ^ rr;
            count += 1;
        }

        rand.xoshiro = [s0, s1, s2, s3];
        rand.seed64 = r;
        rand.count = ((rand.count as usize + size) & (PAD_SWITCH as usize - 1)) as u8;
    }
}