sqlx_sqlserver/protocol/
pre_login.rs

1use crate::{Encrypt, MssqlConnectOptions};
2use thiserror::Error;
3
4use super::packet::{encode_message, PacketFrameError, PacketType};
5
6/// TDS pre-login option token.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8#[repr(u8)]
9pub enum PreLoginOptionToken {
10    /// Protocol version.
11    Version = 0x00,
12    /// Encryption negotiation.
13    Encryption = 0x01,
14    /// Instance name.
15    Instance = 0x02,
16    /// Thread ID.
17    ThreadId = 0x03,
18    /// Multiple active result sets flag.
19    Mars = 0x04,
20    /// Option-list terminator.
21    Terminator = 0xff,
22}
23
24impl TryFrom<u8> for PreLoginOptionToken {
25    type Error = PreLoginError;
26
27    fn try_from(value: u8) -> Result<Self, Self::Error> {
28        match value {
29            0x00 => Ok(Self::Version),
30            0x01 => Ok(Self::Encryption),
31            0x02 => Ok(Self::Instance),
32            0x03 => Ok(Self::ThreadId),
33            0x04 => Ok(Self::Mars),
34            0xff => Ok(Self::Terminator),
35            _ => Err(PreLoginError::UnknownToken(value)),
36        }
37    }
38}
39
40impl From<PreLoginOptionToken> for u8 {
41    fn from(value: PreLoginOptionToken) -> Self {
42        value as u8
43    }
44}
45
46/// One pre-login option and its raw payload.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct PreLoginOption {
49    /// Option token.
50    pub token: PreLoginOptionToken,
51    /// Raw option payload.
52    pub data: Vec<u8>,
53}
54
55/// Assembles a TDS pre-login option table and payload.
56///
57/// Each table entry is encoded as `token`, big-endian `offset`, and big-endian
58/// `length`, followed by a `0xff` terminator and then the concatenated option
59/// payloads. Offsets are relative to the beginning of the pre-login message.
60pub fn assemble_options(options: &[PreLoginOption]) -> Result<Vec<u8>, PreLoginError> {
61    let table_len = options
62        .len()
63        .checked_mul(5)
64        .and_then(|len| len.checked_add(1))
65        .ok_or(PreLoginError::MessageTooLarge)?;
66
67    let mut offset = u16::try_from(table_len).map_err(|_| PreLoginError::MessageTooLarge)?;
68    let payload_len = options
69        .iter()
70        .map(|option| option.data.len())
71        .try_fold(0usize, |sum, len| {
72            sum.checked_add(len).ok_or(PreLoginError::MessageTooLarge)
73        })?;
74
75    let total_len = table_len
76        .checked_add(payload_len)
77        .ok_or(PreLoginError::MessageTooLarge)?;
78
79    u16::try_from(total_len).map_err(|_| PreLoginError::MessageTooLarge)?;
80
81    let mut out = Vec::with_capacity(total_len);
82
83    for option in options {
84        if option.token == PreLoginOptionToken::Terminator {
85            return Err(PreLoginError::TerminatorOption);
86        }
87
88        let len = u16::try_from(option.data.len()).map_err(|_| PreLoginError::MessageTooLarge)?;
89
90        out.push(option.token.into());
91        out.extend_from_slice(&offset.to_be_bytes());
92        out.extend_from_slice(&len.to_be_bytes());
93
94        offset = offset
95            .checked_add(len)
96            .ok_or(PreLoginError::MessageTooLarge)?;
97    }
98
99    out.push(PreLoginOptionToken::Terminator.into());
100
101    for option in options {
102        out.extend_from_slice(&option.data);
103    }
104
105    Ok(out)
106}
107
108/// Parses a TDS pre-login option table and payload.
109pub fn parse_options(input: &[u8]) -> Result<Vec<PreLoginOption>, PreLoginError> {
110    let terminator = input
111        .iter()
112        .position(|byte| *byte == u8::from(PreLoginOptionToken::Terminator))
113        .ok_or(PreLoginError::MissingTerminator)?;
114
115    if terminator % 5 != 0 {
116        return Err(PreLoginError::TruncatedOptionTable);
117    }
118
119    let mut options = Vec::with_capacity(terminator / 5);
120
121    for entry in input[..terminator].chunks_exact(5) {
122        let token = PreLoginOptionToken::try_from(entry[0])?;
123        let offset = usize::from(u16::from_be_bytes([entry[1], entry[2]]));
124        let len = usize::from(u16::from_be_bytes([entry[3], entry[4]]));
125        let end = offset
126            .checked_add(len)
127            .ok_or(PreLoginError::OptionOutOfBounds { offset, len })?;
128
129        let data = input
130            .get(offset..end)
131            .ok_or(PreLoginError::OptionOutOfBounds { offset, len })?
132            .to_vec();
133
134        options.push(PreLoginOption { token, data });
135    }
136
137    Ok(options)
138}
139
140/// Maps SQL Server connection encryption preferences to the TDS pre-login byte.
141pub fn encode_encrypt(encrypt: Encrypt) -> u8 {
142    match encrypt {
143        Encrypt::Off => 0x00,
144        Encrypt::On => 0x01,
145        Encrypt::NotSupported => 0x02,
146        Encrypt::Required => 0x03,
147    }
148}
149
150/// Maps a TDS pre-login encryption byte to a connection encryption preference.
151pub fn decode_encrypt(value: u8) -> Result<Encrypt, PreLoginError> {
152    match value {
153        0x00 => Ok(Encrypt::Off),
154        0x01 => Ok(Encrypt::On),
155        0x02 => Ok(Encrypt::NotSupported),
156        0x03 => Ok(Encrypt::Required),
157        _ => Err(PreLoginError::InvalidEncrypt(value)),
158    }
159}
160
161/// Builds an unframed TDS PRELOGIN payload from connection options.
162pub fn build_pre_login_payload(options: &MssqlConnectOptions) -> Result<Vec<u8>, PreLoginError> {
163    let mut pre_login_options = vec![
164        PreLoginOption {
165            token: PreLoginOptionToken::Version,
166            data: vec![0, 0, 0, 0, 0, 0],
167        },
168        PreLoginOption {
169            token: PreLoginOptionToken::Encryption,
170            data: vec![encode_encrypt(options.encrypt())],
171        },
172        PreLoginOption {
173            token: PreLoginOptionToken::Mars,
174            data: vec![0],
175        },
176    ];
177
178    if let Some(instance) = options.instance() {
179        let mut data = instance.as_bytes().to_vec();
180        data.push(0);
181        pre_login_options.push(PreLoginOption {
182            token: PreLoginOptionToken::Instance,
183            data,
184        });
185    }
186
187    assemble_options(&pre_login_options)
188}
189
190/// Builds framed TDS PRELOGIN packet bytes from connection options.
191pub fn build_pre_login_packet(options: &MssqlConnectOptions) -> Result<Vec<u8>, PreLoginError> {
192    let payload = build_pre_login_payload(options)?;
193
194    encode_message(
195        PacketType::PRE_LOGIN,
196        &payload,
197        usize::try_from(options.requested_packet_size())
198            .map_err(|_| PreLoginError::MessageTooLarge)?,
199    )
200    .map_err(PreLoginError::Packet)
201}
202
203/// Extracts the server encryption response from a PRELOGIN payload.
204pub fn parse_server_encrypt(input: &[u8]) -> Result<Encrypt, PreLoginError> {
205    let options = parse_options(input)?;
206    let encryption = options
207        .iter()
208        .find(|option| option.token == PreLoginOptionToken::Encryption)
209        .and_then(|option| option.data.first().copied())
210        .ok_or(PreLoginError::MissingEncryption)?;
211
212    decode_encrypt(encryption)
213}
214
215/// Error returned while decoding a pre-login helper value.
216#[derive(Debug, Error, PartialEq, Eq)]
217pub enum PreLoginError {
218    /// The option token is not defined by this helper.
219    #[error("unknown TDS pre-login option token 0x{0:02x}")]
220    UnknownToken(u8),
221    /// The encryption value is not defined by TDS.
222    #[error("invalid TDS pre-login encryption value 0x{0:02x}")]
223    InvalidEncrypt(u8),
224    /// The option table did not include a terminator byte.
225    #[error("TDS pre-login option table is missing its terminator")]
226    MissingTerminator,
227    /// The option table terminator appeared in the middle of an option entry.
228    #[error("TDS pre-login option table is truncated")]
229    TruncatedOptionTable,
230    /// A regular pre-login option used the terminator token.
231    #[error("TDS pre-login terminator cannot be encoded as an option")]
232    TerminatorOption,
233    /// An option offset and length point outside the message buffer.
234    #[error("TDS pre-login option points outside the message: offset {offset}, length {len}")]
235    OptionOutOfBounds {
236        /// Option payload offset.
237        offset: usize,
238        /// Option payload length.
239        len: usize,
240    },
241    /// The assembled pre-login message exceeds the protocol's 16-bit offsets.
242    #[error("TDS pre-login message is too large")]
243    MessageTooLarge,
244    /// Packet framing failed.
245    #[error(transparent)]
246    Packet(#[from] PacketFrameError),
247    /// The server PRELOGIN response did not include an encryption option.
248    #[error("TDS pre-login response is missing its encryption option")]
249    MissingEncryption,
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255
256    #[test]
257    fn encryption_values_round_trip() {
258        for encrypt in [
259            Encrypt::NotSupported,
260            Encrypt::Off,
261            Encrypt::On,
262            Encrypt::Required,
263        ] {
264            assert_eq!(encrypt, decode_encrypt(encode_encrypt(encrypt)).unwrap());
265        }
266    }
267
268    #[test]
269    fn encryption_values_match_tds_wire_values() {
270        assert_eq!(0x00, encode_encrypt(Encrypt::Off));
271        assert_eq!(0x01, encode_encrypt(Encrypt::On));
272        assert_eq!(0x02, encode_encrypt(Encrypt::NotSupported));
273        assert_eq!(0x03, encode_encrypt(Encrypt::Required));
274
275        assert_eq!(Encrypt::Off, decode_encrypt(0x00).unwrap());
276        assert_eq!(Encrypt::On, decode_encrypt(0x01).unwrap());
277        assert_eq!(Encrypt::NotSupported, decode_encrypt(0x02).unwrap());
278        assert_eq!(Encrypt::Required, decode_encrypt(0x03).unwrap());
279    }
280
281    #[test]
282    fn rejects_unknown_encryption_value() {
283        assert_eq!(
284            Err(PreLoginError::InvalidEncrypt(0x7f)),
285            decode_encrypt(0x7f)
286        );
287    }
288
289    #[test]
290    fn decodes_known_option_tokens() {
291        assert_eq!(
292            PreLoginOptionToken::Encryption,
293            PreLoginOptionToken::try_from(0x01).unwrap()
294        );
295        assert_eq!(
296            PreLoginOptionToken::Terminator,
297            PreLoginOptionToken::try_from(0xff).unwrap()
298        );
299    }
300
301    #[test]
302    fn assembles_option_table_with_big_endian_offsets() {
303        let bytes = assemble_options(&[
304            PreLoginOption {
305                token: PreLoginOptionToken::Version,
306                data: vec![0, 0, 0, 0, 0, 0],
307            },
308            PreLoginOption {
309                token: PreLoginOptionToken::Encryption,
310                data: vec![encode_encrypt(Encrypt::On)],
311            },
312        ])
313        .unwrap();
314
315        assert_eq!(
316            vec![
317                0x00, 0x00, 0x0b, 0x00, 0x06, // VERSION at offset 11, len 6
318                0x01, 0x00, 0x11, 0x00, 0x01, // ENCRYPTION at offset 17, len 1
319                0xff, // terminator
320                0, 0, 0, 0, 0, 0,    // version payload
321                0x01, // encryption payload
322            ],
323            bytes
324        );
325    }
326
327    #[test]
328    fn parses_option_table_payloads() {
329        let options = parse_options(&[
330            0x00, 0x00, 0x0b, 0x00, 0x06, 0x01, 0x00, 0x11, 0x00, 0x01, 0xff, 0, 0, 0, 0, 0, 0,
331            0x03,
332        ])
333        .unwrap();
334
335        assert_eq!(
336            vec![
337                PreLoginOption {
338                    token: PreLoginOptionToken::Version,
339                    data: vec![0, 0, 0, 0, 0, 0],
340                },
341                PreLoginOption {
342                    token: PreLoginOptionToken::Encryption,
343                    data: vec![0x03],
344                },
345            ],
346            options
347        );
348    }
349
350    #[test]
351    fn builds_pre_login_payload_from_connection_options() {
352        let options = MssqlConnectOptions::parse_url(
353            "mssql://localhost/master?encrypt=not_supported&instance=SQLEXPRESS",
354        )
355        .unwrap();
356        let payload = build_pre_login_payload(&options).unwrap();
357        let parsed = parse_options(&payload).unwrap();
358
359        assert!(parsed.iter().any(|option| {
360            option.token == PreLoginOptionToken::Encryption && option.data == vec![0x02]
361        }));
362        assert!(parsed.iter().any(|option| {
363            option.token == PreLoginOptionToken::Instance && option.data == b"SQLEXPRESS\0"
364        }));
365    }
366
367    #[test]
368    fn extracts_server_encrypt_option() {
369        let payload = assemble_options(&[PreLoginOption {
370            token: PreLoginOptionToken::Encryption,
371            data: vec![0x02],
372        }])
373        .unwrap();
374
375        assert_eq!(
376            Encrypt::NotSupported,
377            parse_server_encrypt(&payload).unwrap()
378        );
379    }
380
381    #[test]
382    fn rejects_pre_login_option_out_of_bounds() {
383        let err = parse_options(&[0x01, 0x00, 0x10, 0x00, 0x01, 0xff]).unwrap_err();
384
385        assert_eq!(PreLoginError::OptionOutOfBounds { offset: 16, len: 1 }, err);
386    }
387}
sqlx_sqlserver/protocol/pre_login.rs

sqlx_sqlserver/protocol/
pre_login.rs
