mssql-tds 0.1.0

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use crate::{query::metadata::ColumnMetadata, token::tokens::SqlCollation};
use core::fmt;
use std::{fmt::Debug, fmt::Display};
use tracing::warn;

use super::{
    lcid_encoding::lcid_to_encoding,
    sqldatatypes::{TypeInfoVariant, is_unicode_type},
};

/// Character encoding used by a [`SqlString`].
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum EncodingType {
    /// UTF-8 encoding.
    Utf8,
    /// UTF-16LE encoding.
    Utf16,
    /// Encoding derived from SQL collation LCID.
    LcidBased(SqlCollation),
    /// Placeholder set before the connection collation is known.
    // This is to be used when we want to have an empty encoding, which
    // is later written over the protocol by getting the collation from the connection.
    DelayedSet,
}

/// Encoded string value from a TDS character column.
#[derive(PartialEq, Clone)]
pub struct SqlString {
    /// Raw encoded bytes.
    pub bytes: Vec<u8>,
    encoding_type: EncodingType,
}

/// Maps a collation's LCID to an encoding, falling back to Windows-1252 with a
/// warning when the LCID is not one we map.
fn lcid_encoding_or_fallback(collation: SqlCollation) -> &'static encoding_rs::Encoding {
    // LCID lives in the lower 20 bits of collation.info.
    let lcid = collation.info & 0x000F_FFFF;
    match lcid_to_encoding(lcid) {
        Ok(encoding) => encoding,
        Err(e) => {
            warn!(
                "Unsupported LCID 0x{:04X} ({}), falling back to Windows-1252. Error: {}",
                lcid, lcid, e
            );
            encoding_rs::WINDOWS_1252
        }
    }
}

/// Encodes `text` for the wire under `collation`'s narrow encoding: UTF-8 when
/// the collation is UTF-8-aware, or its single-byte LCID codepage otherwise
/// (falling back to Windows-1252 for an LCID this crate does not map).
///
/// Mirrors the encoding step [`get_encoding_type`] performs for a materialized
/// value, for a caller that must produce collation-correct wire bytes
/// directly instead of routing through the parameter serializer -- e.g. a
/// data-at-execution write, which streams bytes to the wire before the normal
/// parameter-serialization path runs. Does *not* mirror that serializer's own
/// `VARCHAR | CHAR | TEXT` arm (`tds_value_serializer.rs`): that arm predates
/// the UTF-8-collation flag and always encodes through the single-byte LCID
/// codepage regardless of `collation.utf8()`, so a UTF-8-collation database
/// gets correct UTF-8 wire bytes from a value streamed through this function
/// but single-byte-miscoded bytes from the same value bound inline. Tracked
/// under AB#47590.
pub fn encode_narrow(text: &str, collation: SqlCollation) -> Vec<u8> {
    if collation.utf8() {
        return text.as_bytes().to_vec();
    }
    let (encoded, encoding_used, had_errors) = lcid_encoding_or_fallback(collation).encode(text);
    if had_errors {
        warn!(
            "Encountered encoding errors while converting string to LCID 0x{:04X} ({}) encoding. \
             Some characters may have been replaced.",
            collation.info & 0x000F_FFFF,
            encoding_used.name()
        );
    }
    encoded.into_owned()
}

impl EncodingType {
    /// The encoding these bytes are in, or `None` when the collation is not yet
    /// known ([`EncodingType::DelayedSet`]).
    ///
    /// Exists so a writer handed borrowed wire bytes by
    /// [`RowWriter::write_string`](crate::datatypes::row_writer::RowWriter::write_string)
    /// can transcode straight into its own buffer instead of building an owned
    /// [`SqlString`] first.
    ///
    /// Decoding through this substitutes U+FFFD on malformed input, whereas
    /// [`SqlString::to_utf8_string`] panics on invalid UTF-8 under
    /// [`EncodingType::Utf8`]. Use this when replacement is the wanted
    /// behaviour, not as a drop-in for `to_utf8_string`.
    pub fn encoding(&self) -> Option<&'static encoding_rs::Encoding> {
        match self {
            EncodingType::Utf8 => Some(encoding_rs::UTF_8),
            EncodingType::Utf16 => Some(encoding_rs::UTF_16LE),
            EncodingType::LcidBased(collation) => Some(lcid_encoding_or_fallback(*collation)),
            EncodingType::DelayedSet => None,
        }
    }
}

impl SqlString {
    /// Creates a `SqlString` from raw bytes and an encoding type.
    pub fn new(bytes: Vec<u8>, encoding_type: EncodingType) -> Self {
        SqlString {
            bytes,
            encoding_type,
        }
    }

    /// Splits into the raw encoded bytes and their encoding.
    pub fn into_parts(self) -> (Vec<u8>, EncodingType) {
        (self.bytes, self.encoding_type)
    }

    /// Creates a UTF-16LE–encoded `SqlString` from a Rust `String`.
    pub fn from_utf8_string(string: String) -> Self {
        let utf16_bytes = string
            .encode_utf16()
            .flat_map(|f| f.to_le_bytes())
            .collect::<Vec<u8>>();
        SqlString::new(utf16_bytes, EncodingType::Utf16)
    }

    /// Decodes the stored bytes into a Rust `String` according to the encoding type.
    pub fn to_utf8_string(&self) -> String {
        Self::decode(&self.bytes, self.encoding_type)
    }

    /// Decodes wire bytes in `encoding_type` into a Rust `String`.
    ///
    /// Lets a writer handed borrowed bytes by
    /// [`RowWriter::write_string`](crate::datatypes::row_writer::RowWriter::write_string)
    /// decode them without first copying into an owned [`SqlString`].
    pub fn decode(bytes: &[u8], encoding_type: EncodingType) -> String {
        match encoding_type {
            // TODO: Investigation needed. When creating a Utf8 strings from the vector, the string is weirdly encoded.
            // UTF16 decode works better.
            EncodingType::Utf8 => String::from_utf8(bytes.to_vec()).unwrap(),
            EncodingType::Utf16 => {
                // Use encoding_rs for efficient UTF-16LE decoding without intermediate Vec<u16> allocation
                let (decoded, _, _) = encoding_rs::UTF_16LE.decode(bytes);
                decoded.into_owned()
            }
            EncodingType::LcidBased(collation) => {
                // Extract LCID from the lower 20 bits of collation.info
                let lcid = collation.info & 0x000F_FFFF;
                let encoding = lcid_encoding_or_fallback(collation);

                // Decode bytes using the determined encoding
                let (decoded, _used_encoding, had_errors) = encoding.decode(bytes);

                if had_errors {
                    warn!(
                        "Encountered decoding errors while converting LCID 0x{:04X} ({}) encoded data. \
                         Some characters may have been replaced with U+FFFD.",
                        lcid, lcid
                    );
                }

                decoded.into_owned()
            }
            EncodingType::DelayedSet => {
                // DelayedSet encoding is not defined, so we return the bytes as a UTF-8 string.
                unimplemented!("DelayedSet encoding conversion to UTF8 not implemented");
            }
        }
    }

    /// Returns true if this SqlString is already encoded as UTF-16
    #[inline]
    pub fn is_utf16(&self) -> bool {
        matches!(self.encoding_type, EncodingType::Utf16)
    }

    /// Returns the raw UTF-16 bytes if already encoded, otherwise None
    /// This avoids re-encoding strings that are already in UTF-16 format
    #[inline]
    pub fn as_utf16_bytes(&self) -> Option<&[u8]> {
        if self.is_utf16() {
            Some(&self.bytes)
        } else {
            None
        }
    }

    /// Returns the raw bytes when they should be written directly to the wire
    /// without encoding conversion. This is the case for DelayedSet and LcidBased
    /// encodings where the bytes are already in the correct wire format.
    #[inline]
    pub fn as_raw_wire_bytes(&self) -> Option<&[u8]> {
        match &self.encoding_type {
            EncodingType::DelayedSet | EncodingType::LcidBased(_) => Some(&self.bytes),
            _ => None,
        }
    }

    /// Returns the encoding type of this SqlString
    #[inline]
    pub fn encoding_type(&self) -> &EncodingType {
        &self.encoding_type
    }
}

impl Debug for SqlString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.encoding_type {
            EncodingType::LcidBased(_) => write!(f, "{:?}", self.bytes),
            EncodingType::DelayedSet => write!(f, "DelayedSet encoded: {:?}", self.bytes.len()),
            _ => write!(f, "{:?}", self.to_utf8_string()),
        }
    }
}

impl Display for SqlString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let EncodingType::LcidBased(_) = self.encoding_type {
            write!(f, "{:?}", self.bytes)
        } else {
            write!(f, "{}", self.to_utf8_string())
        }
    }
}

/// Determines the character encoding for a column from its metadata.
pub fn get_encoding_type(metadata: &ColumnMetadata) -> EncodingType {
    let collation = match metadata.type_info.type_info_variant {
        TypeInfoVariant::PartialLen(_, _, collation, _, _) => collation,
        TypeInfoVariant::VarLenString(_, _, collation) => collation,
        _ => None,
    };

    if is_unicode_type(metadata.data_type) {
        EncodingType::Utf16
    } else if collation.is_some() && collation.unwrap().utf8() {
        EncodingType::Utf8
    } else {
        EncodingType::LcidBased(collation.unwrap())
    }
}

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

    #[test]
    fn test_sql_string_new() {
        let bytes = vec![72, 0, 101, 0, 108, 0, 108, 0, 111, 0];
        let sql_str = SqlString::new(bytes.clone(), EncodingType::Utf16);
        assert_eq!(sql_str.bytes, bytes);
    }

    #[test]
    fn test_from_utf8_string() {
        let input = "Hello World".to_string();
        let sql_str = SqlString::from_utf8_string(input.clone());
        assert_eq!(sql_str.to_utf8_string(), input);
    }

    #[test]
    fn test_to_utf8_string_utf16() {
        let bytes = vec![72, 0, 105, 0];
        let sql_str = SqlString::new(bytes, EncodingType::Utf16);
        assert_eq!(sql_str.to_utf8_string(), "Hi");
    }

    #[test]
    fn test_to_utf8_string_utf8() {
        let bytes = "Test".as_bytes().to_vec();
        let sql_str = SqlString::new(bytes, EncodingType::Utf8);
        assert_eq!(sql_str.to_utf8_string(), "Test");
    }

    #[test]
    fn test_sql_string_clone() {
        let sql_str = SqlString::from_utf8_string("Clone test".to_string());
        let cloned = sql_str.clone();
        assert_eq!(sql_str.bytes, cloned.bytes);
    }

    fn collation(lcid: u32) -> SqlCollation {
        SqlCollation {
            info: lcid,
            lcid_language_id: lcid as i32,
            col_flags: 0,
            sort_id: 0,
        }
    }

    #[test]
    fn encoding_maps_the_unicode_variants() {
        assert_eq!(EncodingType::Utf8.encoding(), Some(encoding_rs::UTF_8));
        assert_eq!(EncodingType::Utf16.encoding(), Some(encoding_rs::UTF_16LE));
    }

    #[test]
    fn encoding_is_none_until_the_collation_is_known() {
        assert_eq!(EncodingType::DelayedSet.encoding(), None);
    }

    #[test]
    fn encoding_resolves_a_known_lcid() {
        // 0x0409 (en-US) maps to Windows-1252, and the fallback would also
        // produce Windows-1252, so pin an LCID whose encoding is distinct from
        // the fallback to prove the lookup actually ran.
        let encoding = EncodingType::LcidBased(collation(0x0419)).encoding();
        assert_eq!(encoding, Some(lcid_to_encoding(0x0419).unwrap()));
        assert_ne!(encoding, Some(encoding_rs::WINDOWS_1252));
    }

    #[test]
    fn encoding_falls_back_for_an_unmapped_lcid() {
        let unmapped = 0x000F_FFFF;
        assert!(lcid_to_encoding(unmapped).is_err(), "LCID must be unmapped");
        assert_eq!(
            EncodingType::LcidBased(collation(unmapped)).encoding(),
            Some(encoding_rs::WINDOWS_1252)
        );
    }

    #[test]
    fn encoding_agrees_with_to_utf8_string_for_lcid_bytes() {
        // The accessor has to decode to the same text `to_utf8_string` would,
        // otherwise a writer using it would silently diverge from the owned path.
        let encoding_type = EncodingType::LcidBased(collation(0x0419));
        let bytes = vec![0xCF, 0xF0, 0xE8, 0xE2, 0xE5, 0xF2];

        let via_accessor = encoding_type
            .encoding()
            .expect("LCID encoding is known")
            .decode(&bytes)
            .0
            .into_owned();

        assert_eq!(
            via_accessor,
            SqlString::new(bytes, encoding_type).to_utf8_string()
        );
    }

    #[test]
    fn decode_matches_to_utf8_string_across_encodings() {
        // A writer decoding borrowed bytes must land on exactly the text the
        // owned path produces, or the two row-write paths silently diverge.
        let cases = [
            (b"h\0i\0".to_vec(), EncodingType::Utf16),
            (b"hi".to_vec(), EncodingType::Utf8),
            (
                vec![0xCF, 0xF0, 0xE8, 0xE2, 0xE5, 0xF2],
                EncodingType::LcidBased(collation(0x0419)),
            ),
        ];

        for (bytes, encoding_type) in cases {
            assert_eq!(
                SqlString::decode(&bytes, encoding_type),
                SqlString::new(bytes.clone(), encoding_type).to_utf8_string(),
                "mismatch for {encoding_type:?}"
            );
        }
    }

    #[test]
    fn test_sql_string_debug_utf16() {
        let sql_str = SqlString::from_utf8_string("Debug".to_string());
        let debug_str = format!("{sql_str:?}");
        assert!(debug_str.contains("Debug"));
    }

    #[test]
    fn test_sql_string_debug_delayed_set() {
        let sql_str = SqlString::new(vec![1, 2, 3, 4, 5], EncodingType::DelayedSet);
        let debug_str = format!("{sql_str:?}");
        assert!(debug_str.contains("DelayedSet"));
        assert!(debug_str.contains("5"));
    }

    #[test]
    fn test_sql_string_display_utf16() {
        let sql_str = SqlString::from_utf8_string("Display".to_string());
        let display_str = format!("{sql_str}");
        assert_eq!(display_str, "Display");
    }

    #[test]
    fn test_sql_string_equality() {
        let sql_str1 = SqlString::from_utf8_string("Equal".to_string());
        let sql_str2 = SqlString::from_utf8_string("Equal".to_string());
        let sql_str3 = SqlString::from_utf8_string("Different".to_string());
        assert_eq!(sql_str1, sql_str2);
        assert_ne!(sql_str1, sql_str3);
    }

    #[test]
    fn test_from_utf8_string_empty() {
        let sql_str = SqlString::from_utf8_string(String::new());
        assert_eq!(sql_str.to_utf8_string(), "");
        assert!(sql_str.bytes.is_empty());
    }

    #[test]
    fn test_from_utf8_string_special_chars() {
        let input = "Hello! @#$%^&*()".to_string();
        let sql_str = SqlString::from_utf8_string(input.clone());
        assert_eq!(sql_str.to_utf8_string(), input);
    }

    #[test]
    fn test_from_utf8_string_unicode() {
        let input = "Hello World".to_string();
        let sql_str = SqlString::from_utf8_string(input.clone());
        assert_eq!(sql_str.to_utf8_string(), input);
    }

    #[test]
    fn test_sql_string_new_utf8() {
        let bytes = "UTF8 String".as_bytes().to_vec();
        let sql_str = SqlString::new(bytes.clone(), EncodingType::Utf8);
        assert_eq!(sql_str.bytes, bytes);
        assert_eq!(sql_str.to_utf8_string(), "UTF8 String");
    }

    #[test]
    fn test_sql_string_new_delayed_set() {
        let bytes = vec![1, 2, 3, 4];
        let sql_str = SqlString::new(bytes.clone(), EncodingType::DelayedSet);
        assert_eq!(sql_str.bytes, bytes);
    }

    // ========================================================================
    // LCID Encoding Tests
    // ========================================================================

    #[test]
    fn test_lcid_based_encoding_us_english() {
        // Test US English (Windows-1252) encoding
        // "Hello, World!" in Windows-1252
        let text = b"Hello, World!";
        let collation = SqlCollation {
            info: 0x0409, // US English LCID
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        let sql_str = SqlString::new(text.to_vec(), EncodingType::LcidBased(collation));
        assert_eq!(sql_str.to_utf8_string(), "Hello, World!");
    }

    #[test]
    fn test_lcid_based_encoding_special_chars_windows1252() {
        // Test special characters in Windows-1252
        // "Café résumé naïve" with special chars
        let text = b"Caf\xe9 r\xe9sum\xe9 na\xefve"; // é = 0xE9, ï = 0xEF in Windows-1252
        let collation = SqlCollation {
            info: 0x0409, // US English LCID
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        let sql_str = SqlString::new(text.to_vec(), EncodingType::LcidBased(collation));
        assert_eq!(sql_str.to_utf8_string(), "Café résumé naïve");
    }

    #[test]
    fn test_lcid_based_encoding_japanese() {
        // Test Japanese Shift_JIS encoding
        // "こんにちは" (Konnichiwa) in Shift_JIS: 82B1 82F1 82C9 82BF 82CD
        let text = vec![0x82, 0xB1, 0x82, 0xF1, 0x82, 0xC9, 0x82, 0xBF, 0x82, 0xCD];
        let collation = SqlCollation {
            info: 0x0411, // Japanese LCID
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        let sql_str = SqlString::new(text, EncodingType::LcidBased(collation));
        assert_eq!(sql_str.to_utf8_string(), "こんにちは");
    }

    #[test]
    fn test_lcid_based_encoding_with_flags() {
        // Test LCID extraction with flags set in upper bits
        // US English LCID (0x0409) with flags (0x00D00409)
        let text = b"Test";
        let collation = SqlCollation {
            info: 0x00D0_0409, // LCID with comparison flags
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        let sql_str = SqlString::new(text.to_vec(), EncodingType::LcidBased(collation));
        // Should still decode as US English (lower 20 bits = 0x0409)
        assert_eq!(sql_str.to_utf8_string(), "Test");
    }

    #[test]
    fn test_lcid_based_encoding_empty_string() {
        // Test empty string
        let text = vec![];
        let collation = SqlCollation {
            info: 0x0409, // US English LCID
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        let sql_str = SqlString::new(text, EncodingType::LcidBased(collation));
        assert_eq!(sql_str.to_utf8_string(), "");
    }

    #[test]
    fn test_is_utf16() {
        let utf16_str = SqlString::from_utf8_string("test".to_string());
        assert!(utf16_str.is_utf16());

        let utf8_str = SqlString::new(b"test".to_vec(), EncodingType::Utf8);
        assert!(!utf8_str.is_utf16());
    }

    #[test]
    fn test_as_utf16_bytes() {
        let utf16_str = SqlString::from_utf8_string("Hi".to_string());
        let bytes = utf16_str.as_utf16_bytes();
        assert!(bytes.is_some());
        assert_eq!(bytes.unwrap(), &[72, 0, 105, 0]); // "Hi" in UTF-16LE

        let utf8_str = SqlString::new(b"test".to_vec(), EncodingType::Utf8);
        assert!(utf8_str.as_utf16_bytes().is_none());
    }

    #[test]
    fn encode_narrow_uses_the_lcid_codepage_for_a_non_utf8_collation() {
        let collation = SqlCollation {
            info: 0x0409, // US English LCID -> Windows-1252
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        assert_eq!(encode_narrow("Caf\u{e9}", collation), b"Caf\xe9");
    }

    #[test]
    fn encode_narrow_passes_through_utf8_for_a_utf8_collation() {
        let collation = SqlCollation {
            info: 0x0409,
            lcid_language_id: 0,
            col_flags: 0x40, // fUTF8
            sort_id: 0,
        };
        assert_eq!(
            encode_narrow("Caf\u{e9}", collation),
            "Caf\u{e9}".as_bytes()
        );
    }

    #[test]
    fn encode_narrow_falls_back_to_windows_1252_for_an_unmapped_lcid() {
        let collation = SqlCollation {
            info: 0x000F_FFFF,
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        assert_eq!(encode_narrow("Caf\u{e9}", collation), b"Caf\xe9");
    }

    /// U+65E5 has no Windows-1252 representation. `encoding_rs` substitutes
    /// an HTML numeric character reference, not `?`, and `encode_narrow`
    /// warns on this rather than substituting silently.
    #[test]
    fn encode_narrow_substitutes_ncr_for_a_character_the_codepage_cannot_represent() {
        let collation = SqlCollation {
            info: 0x0409, // US English LCID -> Windows-1252
            lcid_language_id: 0,
            col_flags: 0,
            sort_id: 0,
        };
        assert_eq!(encode_narrow("\u{65e5}", collation), b"&#26085;");
    }
}