tiberius 0.13.0

use super::TypeInfoTvp;
use super::{encode_all_headers_tx, encode_b_varchar, Encode};
use crate::{tds::codec::ColumnData, BytesMutWithTypeInfo, Result};
use bytes::{BufMut, BytesMut};
use enumflags2::{bitflags, BitFlags};
use std::borrow::BorrowMut;
use std::borrow::Cow;

#[bitflags]
#[repr(u8)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum RpcStatus {
    ByRefValue = 1 << 0,
    DefaultValue = 1 << 1,
    // reserved
    Encrypted = 1 << 3,
}

#[bitflags]
#[repr(u16)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum RpcOption {
    WithRecomp = 1 << 0,
    NoMeta = 1 << 1,
    ReuseMeta = 1 << 2,
}

#[derive(Debug)]
pub struct TokenRpcRequest<'a> {
    proc_id: RpcProcIdValue<'a>,
    flags: BitFlags<RpcOption>,
    params: Vec<RpcParam<'a>>,
    transaction_desc: [u8; 8],
}

impl<'a> TokenRpcRequest<'a> {
    pub fn new<I>(proc_id: I, params: Vec<RpcParam<'a>>, transaction_desc: [u8; 8]) -> Self
    where
        I: Into<RpcProcIdValue<'a>>,
    {
        Self {
            proc_id: proc_id.into(),
            flags: BitFlags::empty(),
            params,
            transaction_desc,
        }
    }
}

/// The value carried by an [`RpcParam`]. A scalar column value, or a
/// table-valued parameter (TVP).
#[derive(Debug)]
pub enum RpcValue<'a> {
    /// An ordinary scalar parameter value.
    Scalar(ColumnData<'a>),
    /// A table-valued parameter. As per the TDS grammar, `TYPE_INFO_TVP`
    /// carries both the type metadata and the data rows.
    Table(TypeInfoTvp<'a>),
}

#[derive(Debug)]
pub struct RpcParam<'a> {
    pub name: Cow<'a, str>,
    pub flags: BitFlags<RpcStatus>,
    pub value: RpcValue<'a>,
}

/// 2.2.6.6 RPC Request
#[allow(dead_code)]
#[repr(u8)]
#[derive(Clone, Copy, Debug)]
pub enum RpcProcId {
    CursorOpen = 2,
    CursorFetch = 7,
    CursorClose = 9,
    ExecuteSQL = 10,
    Prepare = 11,
    Execute = 12,
    PrepExec = 13,
    Unprepare = 15,
}

#[derive(Debug)]
#[allow(dead_code)]
pub enum RpcProcIdValue<'a> {
    Name(Cow<'a, str>),
    Id(RpcProcId),
}

impl<'a, S> From<S> for RpcProcIdValue<'a>
where
    S: Into<Cow<'a, str>>,
{
    fn from(s: S) -> Self {
        Self::Name(s.into())
    }
}

impl<'a> From<RpcProcId> for RpcProcIdValue<'a> {
    fn from(id: RpcProcId) -> Self {
        Self::Id(id)
    }
}

impl<'a> Encode<BytesMut> for TokenRpcRequest<'a> {
    fn encode(self, dst: &mut BytesMut) -> Result<()> {
        encode_all_headers_tx(dst, self.transaction_desc);

        match self.proc_id {
            RpcProcIdValue::Id(ref id) => {
                let val = (0xffff_u32) | ((*id as u16) as u32) << 16;
                dst.put_u32_le(val);
            }
            RpcProcIdValue::Name(ref name) => {
                // ProcName is a US_VARCHAR: a u16 little-endian character count
                // followed by that many UTF-16 code units. The count is a u16, so
                // a name longer than 65535 code units cannot be represented;
                // reject it rather than overflow the counter and desync the wire.
                let units = name.encode_utf16().count();
                if units > u16::MAX as usize {
                    return Err(crate::Error::Protocol(
                        format!(
                            "RPC procedure name is too long ({units} UTF-16 code units, max 65535)"
                        )
                        .into(),
                    ));
                }

                let len_pos = dst.len();
                dst.put_u16_le(0u16);
                let mut length = 0_u16;

                for chr in name.encode_utf16() {
                    dst.put_u16_le(chr);
                    length += 1;
                }

                let dst: &mut [u8] = dst.borrow_mut();
                let mut dst = &mut dst[len_pos..];
                dst.put_u16_le(length);
            }
        }

        dst.put_u16_le(self.flags.bits());

        for param in self.params.into_iter() {
            param.encode(dst)?;
        }

        Ok(())
    }
}

impl<'a> Encode<BytesMut> for RpcParam<'a> {
    fn encode(self, dst: &mut BytesMut) -> Result<()> {
        // ParamName is a B_VARCHAR (u8 code-unit count); reject an over-long
        // name rather than wrap the counter and desync the wire.
        encode_b_varchar(dst, &self.name)?;

        dst.put_u8(self.flags.bits());

        match self.value {
            RpcValue::Scalar(value) => {
                let mut dst_ti = BytesMutWithTypeInfo::new(dst);
                value.encode(&mut dst_ti)?;
            }
            RpcValue::Table(value) => value.encode(dst)?,
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::super::ALL_HEADERS_LEN_TX;
    use super::*;
    use crate::tds::codec::ColumnData;

    fn scalar(value: ColumnData<'static>) -> RpcValue<'static> {
        RpcValue::Scalar(value)
    }

    #[test]
    fn encodes_named_proc_header() {
        let req = TokenRpcRequest::new(
            "dbo.usp_MyProc",
            vec![RpcParam {
                name: Cow::Borrowed("@id"),
                flags: BitFlags::empty(),
                value: scalar(ColumnData::I32(Some(1))),
            }],
            [0u8; 8],
        );

        let mut buf = BytesMut::new();
        req.encode(&mut buf).unwrap();

        // Skip the ALL_HEADERS block, positioned right at the ProcName.
        let name_pos = ALL_HEADERS_LEN_TX;
        let len = u16::from_le_bytes([buf[name_pos], buf[name_pos + 1]]);
        assert_eq!(len as usize, "dbo.usp_MyProc".encode_utf16().count());

        // Verify the UTF-16 payload matches the proc name.
        let mut chars = Vec::new();
        let mut off = name_pos + 2;
        for _ in 0..len {
            chars.push(u16::from_le_bytes([buf[off], buf[off + 1]]));
            off += 2;
        }
        assert_eq!(String::from_utf16(&chars).unwrap(), "dbo.usp_MyProc");

        // Option flags (u16) follow the name.
        let flags = u16::from_le_bytes([buf[off], buf[off + 1]]);
        assert_eq!(flags, 0);
    }

    #[test]
    fn named_and_by_id_differ_only_in_proc_slot() {
        let by_id = {
            let req = TokenRpcRequest::new(RpcProcId::ExecuteSQL, vec![], [0u8; 8]);
            let mut buf = BytesMut::new();
            req.encode(&mut buf).unwrap();
            buf
        };

        // By-id encodes 0xFFFF followed by the proc id in the high word.
        let val = u32::from_le_bytes([
            by_id[ALL_HEADERS_LEN_TX],
            by_id[ALL_HEADERS_LEN_TX + 1],
            by_id[ALL_HEADERS_LEN_TX + 2],
            by_id[ALL_HEADERS_LEN_TX + 3],
        ]);
        assert_eq!(val & 0xffff, 0xffff);
        assert_eq!((val >> 16) as u16, RpcProcId::ExecuteSQL as u16);
    }

    #[test]
    fn encodes_param_name_and_by_ref_flag() {
        let param = RpcParam {
            name: Cow::Borrowed("@out"),
            flags: BitFlags::from_flag(RpcStatus::ByRefValue),
            value: scalar(ColumnData::I32(Some(7))),
        };

        let mut buf = BytesMut::new();
        param.encode(&mut buf).unwrap();

        // First byte is the param-name length (in UTF-16 code units).
        assert_eq!(buf[0] as usize, "@out".encode_utf16().count());

        let mut chars = Vec::new();
        let mut off = 1usize;
        for _ in 0..buf[0] {
            chars.push(u16::from_le_bytes([buf[off], buf[off + 1]]));
            off += 2;
        }
        assert_eq!(String::from_utf16(&chars).unwrap(), "@out");

        // Status flags byte carries the ByRefValue bit.
        assert_eq!(buf[off], RpcStatus::ByRefValue as u8);
    }

    // ProcName length is a u16; a name longer than 65535 UTF-16 code units must
    // error rather than overflow the counter (and desync the wire).
    #[test]
    fn rejects_over_long_proc_name() {
        let long = "a".repeat(u16::MAX as usize + 1);
        let req = TokenRpcRequest::new(long, vec![], [0u8; 8]);

        let mut buf = BytesMut::new();
        let err = req.encode(&mut buf).unwrap_err();
        assert!(matches!(err, crate::Error::Protocol(_)), "got {err:?}");
    }

    // A param name is a B_VARCHAR (u8 length); a name longer than 255 UTF-16 code
    // units must error rather than wrap the counter and desync the wire.
    #[test]
    fn rejects_over_long_param_name() {
        let param = RpcParam {
            name: Cow::Owned("@".to_string() + &"a".repeat(255)),
            flags: BitFlags::empty(),
            value: scalar(ColumnData::I32(Some(1))),
        };

        let mut buf = BytesMut::new();
        let err = param.encode(&mut buf).unwrap_err();
        assert!(matches!(err, crate::Error::Protocol(_)), "got {err:?}");
    }

    // A param name of exactly 255 units is the boundary and must still encode,
    // writing the correct u8 length prefix.
    #[test]
    fn encodes_max_length_param_name() {
        let name = "a".repeat(255);
        let param = RpcParam {
            name: Cow::Owned(name.clone()),
            flags: BitFlags::empty(),
            value: scalar(ColumnData::I32(Some(1))),
        };

        let mut buf = BytesMut::new();
        param.encode(&mut buf).expect("255-unit name must encode");
        assert_eq!(buf[0], 255);
    }
}