sqlmodel-mysql 0.4.0

//! MySQL authentication implementations.
//!
//! This module implements the MySQL authentication plugins:
//! - `mysql_native_password`: SHA1-based (legacy, MySQL < 8.0 default)
//! - `caching_sha2_password`: SHA256-based (MySQL 8.0+ default)
//!
//! # mysql_native_password
//!
//! Password scramble algorithm:
//! ```text
//! SHA1(password) XOR SHA1(seed + SHA1(SHA1(password)))
//! ```
//!
//! # caching_sha2_password
//!
//! Fast auth (if cached on server):
//! ```text
//! XOR(SHA256(password), SHA256(SHA256(SHA256(password)) + seed))
//! ```
//!
//! Full auth requires TLS or RSA public key encryption.

use sha1::Sha1;
use sha2::{Digest, Sha256};

use rand::rand_core::UnwrapErr;
use rand::rngs::SysRng;

use rsa::RsaPublicKey;
use rsa::pkcs1::DecodeRsaPublicKey;
use rsa::pkcs8::DecodePublicKey;

/// Well-known authentication plugin names.
pub mod plugins {
    /// SHA1-based authentication (legacy default)
    pub const MYSQL_NATIVE_PASSWORD: &str = "mysql_native_password";
    /// SHA256-based authentication (MySQL 8.0+ default)
    pub const CACHING_SHA2_PASSWORD: &str = "caching_sha2_password";
    /// RSA-based SHA256 authentication
    pub const SHA256_PASSWORD: &str = "sha256_password";
    /// MySQL clear password (for debugging/testing only)
    pub const MYSQL_CLEAR_PASSWORD: &str = "mysql_clear_password";
}

/// Response codes for caching_sha2_password protocol.
pub mod caching_sha2 {
    /// Request for public key (client should send 0x02)
    pub const REQUEST_PUBLIC_KEY: u8 = 0x02;
    /// Fast auth success
    pub const FAST_AUTH_SUCCESS: u8 = 0x03;
    /// Full auth needed (switch to secure channel or RSA)
    pub const PERFORM_FULL_AUTH: u8 = 0x04;
}

/// Compute mysql_native_password authentication response.
///
/// Algorithm: `SHA1(password) XOR SHA1(seed + SHA1(SHA1(password)))`
///
/// # Arguments
/// * `password` - The user's password (UTF-8)
/// * `auth_data` - The 20-byte scramble from the server
///
/// # Returns
/// The 20-byte authentication response, or empty vec if password is empty.
pub fn mysql_native_password(password: &str, auth_data: &[u8]) -> Vec<u8> {
    if password.is_empty() {
        return vec![];
    }

    // Ensure we only use first 20 bytes of auth_data
    let seed = if auth_data.len() > 20 {
        &auth_data[..20]
    } else {
        auth_data
    };

    // Stage 1: SHA1(password)
    let mut hasher = Sha1::new();
    hasher.update(password.as_bytes());
    let stage1: [u8; 20] = hasher.finalize().into();

    // Stage 2: SHA1(SHA1(password))
    let mut hasher = Sha1::new();
    hasher.update(stage1);
    let stage2: [u8; 20] = hasher.finalize().into();

    // Stage 3: SHA1(seed + stage2)
    let mut hasher = Sha1::new();
    hasher.update(seed);
    hasher.update(stage2);
    let stage3: [u8; 20] = hasher.finalize().into();

    // Final: stage1 XOR stage3
    stage1
        .iter()
        .zip(stage3.iter())
        .map(|(a, b)| a ^ b)
        .collect()
}

/// Compute caching_sha2_password fast authentication response.
///
/// Algorithm: `XOR(SHA256(password), SHA256(SHA256(SHA256(password)) + seed))`
///
/// # Arguments
/// * `password` - The user's password (UTF-8)
/// * `auth_data` - The scramble from the server (typically 20 bytes + NUL)
///
/// # Returns
/// The 32-byte authentication response, or empty vec if password is empty.
pub fn caching_sha2_password(password: &str, auth_data: &[u8]) -> Vec<u8> {
    if password.is_empty() {
        return vec![];
    }

    // Remove trailing NUL if present (MySQL sends 20-byte scramble + NUL = 21 bytes)
    // Only strip if length is 21 and ends with NUL, to avoid modifying valid 20-byte seeds
    let seed = if auth_data.len() == 21 && auth_data.last() == Some(&0) {
        &auth_data[..20]
    } else {
        auth_data
    };

    // SHA256(password)
    let mut hasher = Sha256::new();
    hasher.update(password.as_bytes());
    let password_hash: [u8; 32] = hasher.finalize().into();

    // SHA256(SHA256(password))
    let mut hasher = Sha256::new();
    hasher.update(password_hash);
    let password_hash_hash: [u8; 32] = hasher.finalize().into();

    // SHA256(SHA256(SHA256(password)) + seed)
    let mut hasher = Sha256::new();
    hasher.update(password_hash_hash);
    hasher.update(seed);
    let scramble: [u8; 32] = hasher.finalize().into();

    // XOR(SHA256(password), scramble)
    password_hash
        .iter()
        .zip(scramble.iter())
        .map(|(a, b)| a ^ b)
        .collect()
}

/// Generate a random nonce for client-side use.
///
/// Uses `SysRng` (OS entropy; renamed from `OsRng` in rand 0.10) for
/// cryptographically secure random generation. `SysRng` is fallible-only
/// (`TryRng`) in rand 0.10; `UnwrapErr` adapts it to the infallible `Rng`
/// surface `fill_bytes` needs.
pub fn generate_nonce(length: usize) -> Vec<u8> {
    use rand::Rng;
    let mut bytes = vec![0u8; length];
    UnwrapErr(SysRng).fill_bytes(&mut bytes);
    bytes
}

/// Scramble password for sha256_password plugin using RSA encryption.
///
/// This is used when full authentication is required for caching_sha2_password
/// or sha256_password plugins without TLS.
///
/// # Arguments
/// * `password` - The user's password
/// * `seed` - The authentication seed from server
/// * `public_key` - RSA public key from server (PEM format)
///
/// # Returns
/// The encrypted password, or error if encryption fails.
///
/// This is used for full authentication for `caching_sha2_password`/`sha256_password`
/// when the connection is not secured by TLS.
pub fn sha256_password_rsa(
    password: &str,
    seed: &[u8],
    public_key_pem: &[u8],
    use_oaep: bool,
) -> Result<Vec<u8>, String> {
    // MySQL expects: RSA_encrypt(password_with_nul XOR seed_rotation)
    let mut pw = password.as_bytes().to_vec();
    pw.push(0); // NUL terminator

    if seed.is_empty() {
        return Err("Seed is empty".to_string());
    }

    for (i, b) in pw.iter_mut().enumerate() {
        *b ^= seed[i % seed.len()];
    }

    // Server usually returns a PEM public key for sha256_password/caching_sha2_password.
    let pem = std::str::from_utf8(public_key_pem)
        .map_err(|e| format!("Public key is not valid UTF-8 PEM: {e}"))?;

    // Try both common encodings.
    let pub_key = RsaPublicKey::from_public_key_pem(pem)
        .or_else(|_| RsaPublicKey::from_pkcs1_pem(pem))
        .map_err(|e| format!("Failed to parse RSA public key PEM: {e}"))?;

    let encrypted = if use_oaep {
        // MySQL 8.0.5+ uses OAEP padding for caching_sha2_password.
        let padding = rsa::Oaep::<Sha1>::new();
        pub_key
            .encrypt(&mut UnwrapErr(SysRng), padding, &pw)
            .map_err(|e| format!("RSA OAEP encryption failed: {e}"))?
    } else {
        let padding = rsa::Pkcs1v15Encrypt;
        pub_key
            .encrypt(&mut UnwrapErr(SysRng), padding, &pw)
            .map_err(|e| format!("RSA PKCS1v1.5 encryption failed: {e}"))?
    };

    Ok(encrypted)
}

/// XOR password with seed for cleartext transmission over TLS.
///
/// When the connection is secured with TLS, some auth methods allow sending
/// the password XOR'd with the seed (or even cleartext).
pub fn xor_password_with_seed(password: &str, seed: &[u8]) -> Vec<u8> {
    let password_bytes = password.as_bytes();
    let mut result = Vec::with_capacity(password_bytes.len() + 1);

    for (i, &byte) in password_bytes.iter().enumerate() {
        let seed_byte = seed.get(i % seed.len()).copied().unwrap_or(0);
        result.push(byte ^ seed_byte);
    }

    // NUL terminator
    result.push(0);

    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mysql_native_password_empty() {
        let result = mysql_native_password("", &[0; 20]);
        assert!(result.is_empty());
    }

    #[test]
    fn test_mysql_native_password() {
        // Known test vector from MySQL protocol documentation
        // Seed: 20 bytes of zeros
        // Password: "secret"
        let seed = [0u8; 20];
        let result = mysql_native_password("secret", &seed);

        // Should produce 20 bytes
        assert_eq!(result.len(), 20);

        // The result should be deterministic
        let result2 = mysql_native_password("secret", &seed);
        assert_eq!(result, result2);
    }

    #[test]
    fn test_mysql_native_password_real_seed() {
        // Test with a realistic scramble
        let seed = [
            0x3d, 0x4c, 0x5e, 0x2f, 0x1a, 0x0b, 0x7c, 0x8d, 0x9e, 0xaf, 0x10, 0x21, 0x32, 0x43,
            0x54, 0x65, 0x76, 0x87, 0x98, 0xa9,
        ];

        let result = mysql_native_password("mypassword", &seed);
        assert_eq!(result.len(), 20);

        // Different password should give different result
        let result2 = mysql_native_password("otherpassword", &seed);
        assert_ne!(result, result2);
    }

    #[test]
    fn test_caching_sha2_password_empty() {
        let result = caching_sha2_password("", &[0; 20]);
        assert!(result.is_empty());
    }

    #[test]
    fn test_caching_sha2_password() {
        let seed = [0u8; 20];
        let result = caching_sha2_password("secret", &seed);

        // Should produce 32 bytes (SHA-256 output)
        assert_eq!(result.len(), 32);

        // Should be deterministic
        let result2 = caching_sha2_password("secret", &seed);
        assert_eq!(result, result2);
    }

    #[test]
    fn test_caching_sha2_password_with_nul() {
        // MySQL often sends seed with trailing NUL
        let mut seed = vec![0u8; 20];
        seed.push(0); // Trailing NUL

        let result = caching_sha2_password("secret", &seed);
        assert_eq!(result.len(), 32);

        // Should be same as without NUL
        let result2 = caching_sha2_password("secret", &seed[..20]);
        assert_eq!(result, result2);
    }

    #[test]
    fn test_generate_nonce() {
        let nonce1 = generate_nonce(20);
        let nonce2 = generate_nonce(20);

        assert_eq!(nonce1.len(), 20);
        assert_eq!(nonce2.len(), 20);

        // Should be different (extremely high probability)
        assert_ne!(nonce1, nonce2);
    }

    #[test]
    fn test_xor_password_with_seed() {
        let password = "test";
        let seed = [1, 2, 3, 4, 5, 6, 7, 8];

        let result = xor_password_with_seed(password, &seed);

        // Should be password length + 1 (NUL terminator)
        assert_eq!(result.len(), 5);

        // Last byte should be NUL
        assert_eq!(result[4], 0);

        // XOR is reversible
        let recovered: Vec<u8> = result[..4]
            .iter()
            .enumerate()
            .map(|(i, &b)| b ^ seed[i % seed.len()])
            .collect();
        assert_eq!(recovered, password.as_bytes());
    }

    /// Test-only RSA public key in SPKI (`BEGIN PUBLIC KEY`) form — the shape
    /// MySQL serves for `caching_sha2_password`/`sha256_password` full auth.
    const SPKI_PUBLIC_KEY: &str = "-----BEGIN PUBLIC KEY-----\n\
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArJ59U1RtRan3BEJqsItb\n\
tD3nDU7ThwlNaJ42vRSEt/UzFO5/yemxUz3ogTUcsDMgXeLVjQjwwV0+lh+s9IWc\n\
9fU6nu4Q8yW7Pc/SDpDkFBdEtLAOIjfSMv0CzoPQB0A0njVFe7l7SuyrMWQ/19N5\n\
iEtqZQmP2y7h5a23XgPGogXHm0XnKpueZn9KFXhK2lNhZj9IUuQLmhzrH0pov8Ae\n\
FknazXZxL5aoAG+cJIoHKsf9NcGTzj0Hewb36YlgVi+yZ5NRkhgjklQ8E6IL+aaW\n\
yEOsCBtS8kCd/nHP4t6ZeyExdpggPGQ2nJq18jG+sttI+3AnzQhXtR2Adq9LKVdN\n\
RQIDAQAB\n\
-----END PUBLIC KEY-----\n";

    /// The same key in PKCS#1 (`BEGIN RSA PUBLIC KEY`) form — the fallback
    /// encoding `sha256_password_rsa` accepts.
    const PKCS1_PUBLIC_KEY: &str = "-----BEGIN RSA PUBLIC KEY-----\n\
MIIBCgKCAQEArJ59U1RtRan3BEJqsItbtD3nDU7ThwlNaJ42vRSEt/UzFO5/yemx\n\
Uz3ogTUcsDMgXeLVjQjwwV0+lh+s9IWc9fU6nu4Q8yW7Pc/SDpDkFBdEtLAOIjfS\n\
Mv0CzoPQB0A0njVFe7l7SuyrMWQ/19N5iEtqZQmP2y7h5a23XgPGogXHm0XnKpue\n\
Zn9KFXhK2lNhZj9IUuQLmhzrH0pov8AeFknazXZxL5aoAG+cJIoHKsf9NcGTzj0H\n\
ewb36YlgVi+yZ5NRkhgjklQ8E6IL+aaWyEOsCBtS8kCd/nHP4t6ZeyExdpggPGQ2\n\
nJq18jG+sttI+3AnzQhXtR2Adq9LKVdNRQIDAQAB\n\
-----END RSA PUBLIC KEY-----\n";

    /// The RSA full-auth path must keep working across `rsa` major bumps: the
    /// `pem` feature was folded into `encoding` in 0.10, so both PEM encodings
    /// MySQL can serve must still parse, and both padding modes must produce a
    /// modulus-sized, randomized ciphertext.
    #[test]
    fn test_sha256_password_rsa_accepts_both_pem_encodings() {
        let seed = [
            0x3d, 0x4c, 0x5e, 0x2f, 0x1a, 0x0b, 0x7c, 0x8d, 0x9e, 0xaf, 0x10, 0x21, 0x32, 0x43,
            0x54, 0x65, 0x76, 0x87, 0x98, 0xa9,
        ];

        for pem in [SPKI_PUBLIC_KEY, PKCS1_PUBLIC_KEY] {
            for use_oaep in [true, false] {
                let out = sha256_password_rsa("hunter2", &seed, pem.as_bytes(), use_oaep)
                    .expect("RSA encryption with a 2048-bit MySQL-style key");
                // RSA output is always exactly the modulus size (2048 bits).
                assert_eq!(out.len(), 256, "pem={pem} oaep={use_oaep}");
                // Both paddings are randomized: two encryptions must differ.
                let again = sha256_password_rsa("hunter2", &seed, pem.as_bytes(), use_oaep)
                    .expect("second encryption");
                assert_ne!(out, again, "padding must be randomized");
            }
        }
    }

    #[test]
    fn test_sha256_password_rsa_rejects_bad_input() {
        assert!(sha256_password_rsa("pw", &[], SPKI_PUBLIC_KEY.as_bytes(), true).is_err());
        assert!(sha256_password_rsa("pw", &[1, 2, 3], b"not a pem", true).is_err());
    }

    #[test]
    fn test_plugin_names() {
        assert_eq!(plugins::MYSQL_NATIVE_PASSWORD, "mysql_native_password");
        assert_eq!(plugins::CACHING_SHA2_PASSWORD, "caching_sha2_password");
        assert_eq!(plugins::SHA256_PASSWORD, "sha256_password");
    }
}