feat: Add SSRP support for MSSQL named instances

Co-authored-by: contact <contact@ophir.dev>

Author: cursoragent

sqlx-core/src/mssql/connection/mod.rs

@@ -14,6 +14,7 @@ use std::sync::Arc;
 mod establish;
 mod executor;
 mod prepare;
+mod ssrp;
 mod stream;
 mod tls_prelogin_stream_wrapper;
 

sqlx-core/src/mssql/connection/ssrp.rs

@@ -0,0 +1,160 @@
+use crate::error::Error;
+use sqlx_rt::{timeout, UdpSocket};
+use std::collections::HashMap;
+use std::time::Duration;
+
+const SSRP_PORT: u16 = 1434;
+const CLNT_UCAST_INST: u8 = 0x04;
+const SVR_RESP: u8 = 0x05;
+const SSRP_TIMEOUT: Duration = Duration::from_secs(1);
+
+pub(crate) async fn resolve_instance_port(server: &str, instance: &str) -> Result<u16, Error> {
+    let mut request = Vec::with_capacity(1 + instance.len() + 1);
+    request.push(CLNT_UCAST_INST);
+    request.extend_from_slice(instance.as_bytes());
+    request.push(0);
+
+    let socket = UdpSocket::bind("0.0.0.0:0").await.map_err(|e| {
+        err_protocol!("failed to bind UDP socket for SSRP: {}", e)
+    })?;
+
+    socket
+        .send_to(&request, (server, SSRP_PORT))
+        .await
+        .map_err(|e| {
+            err_protocol!("failed to send SSRP request to {}:{}: {}", server, SSRP_PORT, e)
+        })?;
+
+    let mut buffer = [0u8; 1024];
+    let bytes_read = timeout(SSRP_TIMEOUT, socket.recv(&mut buffer))
+        .await
+        .map_err(|_| {
+            err_protocol!(
+                "SSRP request to {} for instance {} timed out after {:?}",
+                server,
+                instance,
+                SSRP_TIMEOUT
+            )
+        })?
+        .map_err(|e| {
+            err_protocol!(
+                "failed to receive SSRP response from {} for instance {}: {}",
+                server,
+                instance,
+                e
+            )
+        })?;
+
+    if bytes_read < 3 {
+        return Err(err_protocol!(
+            "SSRP response too short: {} bytes",
+            bytes_read
+        ));
+    }
+
+    if buffer[0] != SVR_RESP {
+        return Err(err_protocol!(
+            "invalid SSRP response type: expected 0x05, got 0x{:02x}",
+            buffer[0]
+        ));
+    }
+
+    let response_size = u16::from_le_bytes([buffer[1], buffer[2]]) as usize;
+    if response_size + 3 > bytes_read {
+        return Err(err_protocol!(
+            "SSRP response size mismatch: expected {} bytes, got {}",
+            response_size + 3,
+            bytes_read
+        ));
+    }
+
+    let response_data = String::from_utf8(buffer[3..(3 + response_size)].to_vec())
+        .map_err(|e| err_protocol!("SSRP response is not valid UTF-8: {}", e))?;
+
+    parse_ssrp_response(&response_data, instance)
+}
+
+fn parse_ssrp_response(data: &str, instance_name: &str) -> Result<u16, Error> {
+    let instances: Vec<&str> = data.split(";;").collect();
+
+    for instance_data in instances {
+        if instance_data.is_empty() {
+            continue;
+        }
+
+        let tokens: Vec<&str> = instance_data.split(';').collect();
+        let mut properties: HashMap<&str, &str> = HashMap::new();
+
+        let mut i = 0;
+        while i + 1 < tokens.len() {
+            let key = tokens[i];
+            let value = tokens[i + 1];
+            properties.insert(key, value);
+            i += 2;
+        }
+
+        if let Some(name) = properties.get("InstanceName") {
+            if name.eq_ignore_ascii_case(instance_name) {
+                if let Some(tcp_port_str) = properties.get("tcp") {
+                    return tcp_port_str.parse::<u16>().map_err(|e| {
+                        err_protocol!(
+                            "invalid TCP port '{}' in SSRP response: {}",
+                            tcp_port_str,
+                            e
+                        )
+                    });
+                } else {
+                    return Err(err_protocol!(
+                        "instance '{}' found but no TCP port available",
+                        instance_name
+                    ));
+                }
+            }
+        }
+    }
+
+    Err(err_protocol!(
+        "instance '{}' not found in SSRP response",
+        instance_name
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_ssrp_response_single_instance() {
+        let data = "ServerName;MYSERVER;InstanceName;SQLEXPRESS;IsClustered;No;Version;15.0.2000.5;tcp;1433;;";
+        let port = parse_ssrp_response(data, "SQLEXPRESS").unwrap();
+        assert_eq!(port, 1433);
+    }
+
+    #[test]
+    fn test_parse_ssrp_response_multiple_instances() {
+        let data = "ServerName;SRV1;InstanceName;INST1;IsClustered;No;Version;15.0.2000.5;tcp;1433;;ServerName;SRV1;InstanceName;INST2;IsClustered;No;Version;16.0.1000.6;tcp;1434;np;\\\\SRV1\\pipe\\MSSQL$INST2\\sql\\query;;";
+        let port = parse_ssrp_response(data, "INST2").unwrap();
+        assert_eq!(port, 1434);
+    }
+
+    #[test]
+    fn test_parse_ssrp_response_case_insensitive() {
+        let data = "ServerName;MYSERVER;InstanceName;SQLExpress;IsClustered;No;Version;15.0.2000.5;tcp;1433;;";
+        let port = parse_ssrp_response(data, "sqlexpress").unwrap();
+        assert_eq!(port, 1433);
+    }
+
+    #[test]
+    fn test_parse_ssrp_response_instance_not_found() {
+        let data = "ServerName;MYSERVER;InstanceName;SQLEXPRESS;IsClustered;No;Version;15.0.2000.5;tcp;1433;;";
+        let result = parse_ssrp_response(data, "NOTFOUND");
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_parse_ssrp_response_no_tcp_port() {
+        let data = "ServerName;MYSERVER;InstanceName;SQLEXPRESS;IsClustered;No;Version;15.0.2000.5;;";
+        let result = parse_ssrp_response(data, "SQLEXPRESS");
+        assert!(result.is_err());
+    }
+}

sqlx-core/src/mssql/connection/stream.rs

@@ -51,7 +51,13 @@ pub(crate) struct MssqlStream {
 
 impl MssqlStream {
     pub(super) async fn connect(options: &MssqlConnectOptions) -> Result<Self, Error> {
-        let tcp_stream = TcpStream::connect((&*options.host, options.port)).await?;
+        let port = if let Some(ref instance) = options.instance {
+            super::ssrp::resolve_instance_port(&options.host, instance).await?
+        } else {
+            options.port
+        };
+
+        let tcp_stream = TcpStream::connect((&*options.host, port)).await?;
         let wrapped_stream = TlsPreloginWrapper::new(tcp_stream);
         let inner = BufStream::new(MaybeTlsStream::Raw(wrapped_stream));
 

sqlx-core/src/mssql/options/mod.rs

@@ -10,8 +10,14 @@ mod parse;
 ///
 /// Connection strings should be in the form:
 /// ```text
-/// mssql://[username[:password]@]host/database[?instance=instance_name&packet_size=packet_size&client_program_version=client_program_version&client_pid=client_pid&hostname=hostname&app_name=app_name&server_name=server_name&client_interface_name=client_interface_name&language=language]
+/// mssql://[username[:password]@]host[:port]/database[?param1=value1&param2=value2...]
 /// ```
+///
+/// When connecting to a named instance, use the `instance` parameter:
+/// ```text
+/// mssql://user:pass@localhost/mydb?instance=SQLEXPRESS
+/// ```
+/// The port will be automatically discovered using the SQL Server Resolution Protocol (SSRP).
 #[derive(Debug, Clone)]
 pub struct MssqlConnectOptions {
     pub(crate) host: String,

sqlx-core/src/mssql/options/parse.rs

@@ -23,7 +23,7 @@ impl FromStr for MssqlConnectOptions {
     /// - `database`: The name of the database to connect to.
     ///
     /// Supported query parameters:
-    /// - `instance`: SQL Server named instance.
+    /// - `instance`: SQL Server named instance. When specified, the port is automatically discovered using the SQL Server Resolution Protocol (SSRP).
     /// - `encrypt`: Controls connection encryption:
     ///   - `strict`: Requires encryption and validates the server certificate.
     ///   - `mandatory` or `true` or `yes`: Requires encryption but doesn't validate the server certificate.
@@ -41,9 +41,10 @@ impl FromStr for MssqlConnectOptions {
     /// - `client_interface_name`: Name of the client interface, sent to the server for logging purposes.
     /// - `language`: Sets the language for server messages. Affects date formats and system messages.
     ///
-    /// Example:
+    /// Examples:
     /// ```text
     /// mssql://user:pass@localhost:1433/mydb?encrypt=strict&app_name=MyApp&packet_size=4096
+    /// mssql://user:pass@localhost/mydb?instance=SQLEXPRESS
     /// ```
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         let url: Url = s.parse().map_err(Error::config)?;

sqlx-rt/src/rt_async_std.rs

@@ -1,8 +1,8 @@
 pub use async_std::{
     self, fs, future::timeout, io::prelude::ReadExt as AsyncReadExt,
     io::prelude::WriteExt as AsyncWriteExt, io::Read as AsyncRead, io::Write as AsyncWrite,
-    net::TcpStream, sync::Mutex as AsyncMutex, task::sleep, task::spawn, task::spawn_blocking,
-    task::yield_now,
+    net::TcpStream, net::UdpSocket, sync::Mutex as AsyncMutex, task::sleep, task::spawn,
+    task::spawn_blocking, task::yield_now,
 };
 
 #[cfg(unix)]

sqlx-rt/src/rt_tokio.rs

@@ -1,7 +1,7 @@
 pub use tokio::{
     self, fs, io::AsyncRead, io::AsyncReadExt, io::AsyncWrite, io::AsyncWriteExt, io::ReadBuf,
-    net::TcpStream, runtime::Handle, sync::Mutex as AsyncMutex, task::spawn, task::yield_now,
-    time::sleep, time::timeout,
+    net::TcpStream, net::UdpSocket, runtime::Handle, sync::Mutex as AsyncMutex, task::spawn,
+    task::yield_now, time::sleep, time::timeout,
 };
 
 #[cfg(unix)]